"""BFCL mock-function environment: model calls declared functions; the env
executes them against official ground-truth state and compares.

Official v3 semantics (eval_checker): for AST-scorable categories the
predicted call sequence (name + args) is compared against ground_truth
tool_calls; stateful multi-turn categories additionally compare final
env state. This env implements the stateful half: it tracks a Python-dict
world, applies ground-truth effects for known calls, and exposes the model's
call sequence + final state for the scorer.
"""

import json
from typing import Any, Dict, List

from ...data.sample import ChatMessage, Sample
from ..loop import Environment


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


class BFCLEnvironment(Environment):
    """Records the model's calls; applies no real side effects (official
    mock APIs are deterministic). final_state() exposes calls + ground truth."""

    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 = []
        target = sample.target
        self.ground_truth = _parse_ground_truth(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}