"""Agent evaluation driver: a MESSAGE PUMP, not a thinking framework.

We are EVALUATING agents, not being one: the model under test does whatever
thinking it wants; this loop only (1) hands tool declarations over, (2)
executes the model's tool_calls against an Environment plugin, (3) feeds
observations back, (4) stops at max-turns / no-more-calls / env-done.
Everything is recorded as a trajectory for env_reward scorers.
"""

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

from ..data.sample import ChatMessage, Sample
from ..model.adapter import ModelAdapter
from ..model.output import ModelOutput, Usage


class Environment:
    """Minimal env contract. One instance per sample.

    Two execution models, auto-selected by the runner:
    - MESSAGE PUMP (default): implement reset/step/final_state; the runner
      drives the model turn-by-turn (bfcl_mock style).
    - SELF-RUNNING (official engines): implement run_task(adapter, sample);
      the whole simulation happens inside (official tau2/swe bundles), and
      a prediction dict is returned directly. Falls back to the pump when
      run_task is not overridden.
    """

    name = 'base'
    needs_adapter = False   # self-running engines set True; the runner passes
                            # the adapter to the constructor (no name hardcoding)

    def __init__(self, adapter=None):
        self.adapter = adapter

    def reset(self, sample: Sample) -> List[ChatMessage]:
        return []

    async def step(self, tool_calls: List[Any], text: str, sample: Sample
                   ) -> List[ChatMessage]:
        raise NotImplementedError

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

    async def run_task(self, adapter, sample: Sample, **kw) -> Optional[Dict[str, Any]]:
        """Self-running hook. Return a prediction dict (raw/trajectory/
        env_state/usage/group_key) or None to fall back to the pump."""
        return None


# ---------------- environment registry (万物皆可插件: envs too) ----------------

from ..eval.registry import EvalRegistry  # noqa: E402

ENV_REGISTRY = EvalRegistry('environment')


def register_env(name: str):
    """Class decorator: @register_env('bfcl_mock'). One env per benchmark
    family; drop a module in agent/envs/ and it auto-registers."""

    def decorator(cls):
        ENV_REGISTRY.register(name, cls)
        return cls

    return decorator


def get_env(name: str, adapter=None):
    cls = ENV_REGISTRY.get(name)
    return cls(adapter=adapter)


class Trajectory:
    """Recorded turns: role-tagged messages + per-turn usage."""

    def __init__(self):
        self.messages: List[Dict[str, Any]] = []
        self.turns = 0
        self.usage = Usage()
        self.env_state: Dict[str, Any] = {}

    def add(self, role: str, content: str, **extra) -> None:
        entry: Dict[str, Any] = {'role': role, 'content': content, 'turn': self.turns}
        entry.update(extra)
        self.messages.append(entry)


async def drive(
    adapter: ModelAdapter,
    sample: Sample,
    env: Optional[Environment] = None,
    max_turns: int = 8,
    system: str = '',
) -> Trajectory:
    """Run one sample through the model (+env if present).

    No env  -> single-turn fast path (one generate, done).
    With env-> pump until the model stops calling tools / env says done /
               max_turns reached.
    """
    traj = Trajectory()
    base: List[ChatMessage] = []
    if system:
        base.append(ChatMessage(role='system', content=system))
    base += ([ChatMessage(role='user', content=sample.input)] if isinstance(sample.input, str)
             else list(sample.input))
    if env is not None:
        base += env.reset(sample)
    if getattr(adapter, 'name', '') == 'mock' \
            and adapter.extra.get('mode') in ('boxed', 'oracle', 'fc') \
            and sample.target not in ('', None):
        base = base + [ChatMessage(role='user', content=f'MOCKTARGET::{sample.target}')]

    tools = [{'name': t.name, 'description': t.description or '',
              'parameters': t.parameters} for t in (sample.tools or [])] or None

    messages = list(base)
    t0 = time.time()
    for turn in range(max_turns + 1):
        out: ModelOutput = await adapter.generate(messages, tools=tools)
        traj.turns = turn + 1
        traj.usage = traj.usage + out.usage
        traj.add('assistant', out.text, tool_calls=[c.model_dump() for c in out.tool_calls])
        messages.append(ChatMessage(role='assistant',
                                    content=out.text or _calls_text(out.tool_calls)))
        if not out.tool_calls or env is None:
            break  # plain answer or single-turn: done
        observations = await env.step(out.tool_calls, out.text, sample)
        for obs in observations:
            traj.add(obs.role, obs.content)
            messages.append(obs)
    if env is not None:
        traj.env_state = env.final_state()
    traj.messages.insert(0, {'role': 'meta', 'content': f'turns={traj.turns} '
                           f'latency={time.time() - t0:.1f}s'})
    return traj


def _calls_text(calls) -> str:
    import json

    return json.dumps([c.to_openai()['function'] for c in calls], ensure_ascii=False)


def trajectory_to_prediction(traj: Trajectory) -> Dict[str, Any]:
    """Fold a trajectory into the runner's prediction-dict shape."""
    last_assistant = next((m for m in reversed(traj.messages) if m['role'] == 'assistant'), {})
    return {
        'raw': last_assistant.get('content', ''),
        'trajectory': traj.messages,
        'env_state': traj.env_state or None,
        'usage': traj.usage.model_dump(),
        'group_key': str(last_assistant.get('turn', '')),
    }