"""BFCL mock-function environment + scoring backends (native & official).

The env records the model's calls (official mock APIs are deterministic);
scoring backends compare them against ground truth. Everything here is
registered plugins: @register_env / @register_backend.
"""

import json
from typing import Any, Dict, List, Optional

from ...data.sample import ChatMessage, Sample
from ...eval.registry import EvalRegistry
from ..loop import Environment, register_env

BACKEND_REGISTRY = EvalRegistry('scoring backend')


def register_backend(name: str):
    def decorator(fn):
        BACKEND_REGISTRY.register(name, fn)
        return fn

    return decorator


def get_backend(name: str):
    return BACKEND_REGISTRY.get(name)


def _parse_ground_truth(raw) -> Dict[str, Any]:
    if isinstance(raw, str):
        try:
            return json.loads(raw)
        except (ValueError, TypeError):
            return {}
    return raw or {}


@register_backend('bfcl_official')
def bfcl_official(calls: List[Dict[str, Any]], sample: Sample) -> Optional[bool]:
    """Strict backend: bfcl_eval official ast_checker. None => package
    missing (caller falls back to native). Ground truth and function
    descriptions pass through VERBATIM; only the model side is adapted to
    the official decoded shape (incl. the dot->underscore name convention
    the checker itself applies via convert_func_name)."""
    try:
        from bfcl_eval.eval_checker.ast_eval.ast_checker import ast_checker
    except ImportError:
        return None
    gt = _parse_ground_truth(sample.target)
    if isinstance(gt, dict):
        possible = gt.get('ground_truth')
        gt_calls = gt.get('tool_calls')
        if possible is None and gt_calls:
            possible = [{(c.get('function', c) or {}).get('name', ''):
                         (c.get('function', c) or {}).get('arguments', {})} for c in gt_calls]
    else:
        possible = gt  # official list format, verbatim
    if not possible:
        return len(calls) == 0  # irrelevance: correct = call nothing
    if isinstance(possible, str):
        try:
            possible = json.loads(possible)
        except (ValueError, TypeError):
            return None

    lang = {'java': 'Java', 'javascript': 'JavaScript'}.get(
        str((sample.metadata or {}).get('language', '')).lower(), 'Python')
    category = str((sample.metadata or {}).get('test_category', 'simple'))
    convention = str((sample.metadata or {}).get('bfcl_model_convention',
                                                 'gpt-4o-2024-11-20-FC'))
    # same underscore convention convert_func_name applies to descriptions
    key = (lambda n: n.replace('.', '_')) if 'FC' in convention else (lambda n: n)
    model_output = [{key(c['name']): c['arguments']} for c in calls]

    raw_funcs = (sample.metadata or {}).get('functions')
    funcs = None
    if raw_funcs:
        funcs = json.loads(raw_funcs) if isinstance(raw_funcs, str) else raw_funcs
    if not funcs:
        funcs = [{'name': t.name, 'description': t.description or '',
                  'parameters': t.parameters} for t in (sample.tools or [])]
    try:
        result = ast_checker(funcs, model_output, possible, lang, category, convention)
        return bool(result.get('valid'))
    except Exception:
        # official checker chokes on some v3 schemas; evalscope's adapter
        # behaves identically: exception -> not valid
        return False


@register_backend('bfcl_native')
def bfcl_native(calls: List[Dict[str, Any]], sample: Sample) -> Optional[bool]:
    """Lenient native comparison: exact call-sequence match."""
    gt = _parse_ground_truth(sample.target)
    if isinstance(gt, list):
        want = [list(e.keys())[0] for e in gt] if gt else []
        got = [c['name'] for c in calls]
        return want == got
    return None  # dict-form targets: handled by env_reward's native path


@register_env('bfcl_mock')
class BFCLEnvironment(Environment):
    """Records the model's calls; official mock APIs are deterministic.
    final_state() exposes calls + ground truth for the scorer."""

    name = 'bfcl_mock'

    def __init__(self):
        self.calls: List[Dict[str, Any]] = []
        self.ground_truth: Dict[str, Any] = {}

    def reset(self, sample: Sample) -> List[ChatMessage]:
        self.calls = []
        self.ground_truth = _parse_ground_truth(sample.target)
        return []

    async def step(self, tool_calls, text: str, sample: Sample) -> List[ChatMessage]:
        obs = []
        for call in tool_calls:
            self.calls.append({'name': call.name, 'arguments': call.arguments_dict})
            obs.append(ChatMessage(
                role='tool',
                content=json.dumps({'role': 'function', 'name': call.name,
                                    'content': json.dumps({'status': 'ok'})})))
        return obs

    def final_state(self) -> Dict[str, Any]:
        return {'calls': self.calls, 'ground_truth': self.ground_truth}