"""Takumi EB1-A runtime powered by GPT-5-High router."""

from __future__ import annotations

from dataclasses import dataclass
from functools import lru_cache
from typing import TYPE_CHECKING, Any, Mapping, MutableMapping, Optional, Sequence, cast

from ember.models.schemas import ChatResponse, UsageStats
from .responses import ensure_responses_output

if TYPE_CHECKING:  # pragma: no cover - import only for type checking
    from ember.operators.gpt5.router import GPT5HighRouter, RouterResult


@dataclass(frozen=True, slots=True)
class Eb1ARunResult:
    response: ChatResponse
    decision_trace: Sequence[Mapping[str, Any]]
    usage_trace: Sequence[Mapping[str, Any]]
    directive_log: Sequence[Mapping[str, Any]]


class Eb1ARuntime:
    """Host reusable GPT-5-high router for EB1-A runs."""

    def __init__(
        self,
        controller_binding: Optional[Any] = None,
        cheap_binding: Optional[Any] = None,
    ) -> None:
        from ember.operators.gpt5.router import GPT5HighRouter

        self._router = GPT5HighRouter(
            controller_binding=controller_binding,
            cheap_binding=cheap_binding,
        )

    def run(
        self,
        prompt: str,
        *,
        system: Optional[str] = None,
        metadata: Optional[Mapping[str, Any]] = None,
        shared_options: Optional[Mapping[str, Any]] = None,
        request_options: Optional[Mapping[str, Any]] = None,
        responses_request: Optional[Any] = None,
        context_messages: Optional[Sequence[Mapping[str, str]]] = None,
    ) -> Eb1ARunResult:
        result: "RouterResult" = self._router(
            prompt,
            system=system,
            metadata=metadata,
            shared_options=shared_options,
            request_options=request_options,
            responses_request=responses_request,
            context_messages=context_messages,
        )

        return Eb1ARunResult(
            response=result.response,
            decision_trace=[call.__dict__ for call in result.decision_trace],
            usage_trace=list(result.usage_trace),
            directive_log=[directive.model_dump() for directive in result.directive_log],
        )


@lru_cache(maxsize=1)
def get_eb1_a_runtime() -> Eb1ARuntime:
    """Return a process-wide cached EB1-A runtime instance."""

    return Eb1ARuntime()


def run_eb1_a(
    prompt: str,
    *,
    system: Optional[str] = None,
    metadata: Optional[Mapping[str, Any]] = None,
    shared_options: Optional[Mapping[str, Any]] = None,
    request_options: Optional[Mapping[str, Any]] = None,
    responses_request: Optional[Any] = None,
    context_messages: Optional[Sequence[Mapping[str, str]]] = None,
) -> ChatResponse:
    return _run_eb1_a_internal(
        prompt,
        system=system,
        metadata=metadata,
        shared_options=shared_options,
        request_options=request_options,
        responses_request=responses_request,
        context_messages=context_messages,
        runtime=None,
    )


def _run_eb1_a_internal(
    prompt: str,
    *,
    system: Optional[str],
    metadata: Optional[Mapping[str, Any]],
    shared_options: Optional[Mapping[str, Any]],
    request_options: Optional[Mapping[str, Any]],
    responses_request: Optional[Any],
    context_messages: Optional[Sequence[Mapping[str, str]]],
    runtime: Optional[Eb1ARuntime],
) -> ChatResponse:
    runtime = runtime or get_eb1_a_runtime()
    runtime_metadata = _prepare_runtime_metadata(metadata)
    result = runtime.run(
        prompt,
        system=system,
        metadata=runtime_metadata,
        shared_options=shared_options,
        request_options=request_options,
        responses_request=responses_request,
        context_messages=context_messages,
    )
    response = result.response
    if not isinstance(response, ChatResponse):
        raise TypeError("Eb1ARuntime returned non-ChatResponse payload")

    response.model_id = response.model_id or "takumi/eb1-a"
    ensure_responses_output(
        response,
        default_model="takumi/eb1-a",
        metadata={"runtime": "takumi/eb1-a"},
    )

    _attach_router_metadata(response, result, runtime_metadata)
    return response


def _prepare_runtime_metadata(
    metadata: Optional[Mapping[str, Any]],
) -> Mapping[str, Any]:
    base: dict[str, Any] = {"runtime": "takumi/eb1-a"}
    if metadata:
        base.update(dict(metadata))
    return base


def _attach_router_metadata(
    response: ChatResponse,
    result: Eb1ARunResult,
    runtime_metadata: Mapping[str, Any],
) -> None:
    raw = _ensure_raw_mapping(response.raw_output)
    raw.setdefault("local_backend", True)
    raw.setdefault("runtime", "takumi/eb1-a")
    raw.setdefault("request_metadata", dict(runtime_metadata))

    eb1_a_payload = raw.setdefault("eb1_a", {})
    if not isinstance(eb1_a_payload, MutableMapping):
        eb1_a_payload = {}
        raw["eb1_a"] = eb1_a_payload

    decision_trace = tuple(result.decision_trace)
    usage_trace = tuple(result.usage_trace)
    directive_log = tuple(result.directive_log)
    total_latency_ms = _aggregate_latency(decision_trace)
    aggregated_usage = response.usage

    eb1_a_payload.update(
        {
            "decision_trace": [dict(entry) for entry in decision_trace],
            "usage_trace": [dict(entry) for entry in usage_trace],
            "directive_log": [dict(entry) for entry in directive_log],
            "total_latency_ms": total_latency_ms,
        }
    )
    if aggregated_usage:
        eb1_a_payload["aggregated_usage"] = {
            "prompt_tokens": aggregated_usage.prompt_tokens,
            "completion_tokens": aggregated_usage.completion_tokens,
            "total_tokens": aggregated_usage.total_tokens,
            "cost_usd": aggregated_usage.cost_usd,
            "actual_cost_usd": aggregated_usage.actual_cost_usd,
        }

    response.raw_output = raw


def _ensure_raw_mapping(value: Any) -> MutableMapping[str, Any]:
    if isinstance(value, MutableMapping):
        return cast(MutableMapping[str, Any], value)
    if isinstance(value, Mapping):
        return dict(value)
    return {}


def _aggregate_latency(decision_trace: Sequence[Mapping[str, Any]]) -> float:
    total = 0.0
    for entry in decision_trace:
        latency = entry.get("latency_ms")
        if latency is None:
            continue
        try:
            total += float(latency)
        except (TypeError, ValueError):
            continue
    return total


__all__ = ["Eb1ARuntime", "Eb1ARunResult", "get_eb1_a_runtime", "run_eb1_a"]
