# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
"""One agent turn: (role, prompt, schema, model) -> AgentTurn, through the provider interface.

Failures (provider error, validation failure after retries, spend cap) are recorded in ``AgentTurn.error``
with ``estimate=None`` / ``critiques=None``; nothing is ever fabricated. ``SpendCapExceeded`` is re-raised
so that the experiment runner stops gracefully.
"""
from __future__ import annotations

import builtins
from typing import Any

from pydantic import BaseModel

from ..providers import CostLedger, LLMProvider, get_provider
from ..providers.ledger import SpendCapExceeded
from ..schemas import AgentEstimate, AgentTurn, CritiqueSet
from .roles import Role


class ProviderPool:
    """One provider adapter per model id (each with its own concurrency semaphore), sharing one ledger."""

    def __init__(self, ledger: CostLedger | None = None, force_mock: bool | None = None):
        self.ledger = ledger or CostLedger()
        self.force_mock = force_mock
        self._providers: dict[str, LLMProvider] = {}

    def get(self, model: str) -> LLMProvider:
        if model not in self._providers:
            self._providers[model] = get_provider(model, self.ledger, force_mock=self.force_mock)
        return self._providers[model]

    @property
    def mode(self) -> str:
        return "mock" if self.force_mock else "live"


async def run_turn(
    provider: LLMProvider,
    role: Role,
    view_prompt: str,
    schema: type[BaseModel],
    model: str,
    effort: str = "medium",
    round: int = 0,  # noqa: A002 - mirrors AgentTurn.round
    kind: str = "estimate",
    *,
    max_tokens: int = 4096,
    tag: str = "",
) -> AgentTurn:
    """Run one structured call and wrap it as an AgentTurn. Never raises except for SpendCapExceeded."""
    turn = AgentTurn(role=role.id, round=round, kind=kind, model=model, provider=getattr(provider, "name", "unknown"))  # type: ignore[arg-type]
    try:
        resp = await provider.complete(
            system=role.system_prompt(), user=view_prompt, schema=schema, model=model, effort=effort, max_tokens=max_tokens, tag=tag or f"{role.id}:{kind}:{round}"
        )
    except SpendCapExceeded:
        raise
    except Exception as e:  # defensive: provider.complete already catches most errors
        turn.error = f"{type(e).__name__}: {e}"[:500]
        return turn
    turn.input_tokens = int(resp.input_tokens)
    turn.output_tokens = int(resp.output_tokens)
    turn.cached_tokens = int(resp.cached_tokens)
    turn.cost_usd = float(resp.cost_usd)
    turn.latency_s = builtins.round(float(resp.latency_s), 4)
    turn.prompt_hash = resp.prompt_hash
    turn.request_id = resp.request_id
    turn.validation_retries = int(resp.validation_retries)
    turn.error = resp.error
    parsed: Any = resp.parsed
    if parsed is None:
        turn.error = turn.error or "no parsed output"
    elif isinstance(parsed, AgentEstimate):
        turn.estimate = parsed
    elif isinstance(parsed, CritiqueSet):
        turn.critiques = parsed
    else:
        turn.error = f"unexpected schema {type(parsed).__name__}"
    return turn