"""Benchmark router driven by gpt-5-high."""

from __future__ import annotations

import json
import time
from collections.abc import Mapping, MutableMapping, Sequence
from dataclasses import dataclass, field
from enum import Enum

from pydantic import BaseModel, ConfigDict, Field, ValidationError

from ember.api.models import ModelBinding
from ember.models.providers.utils import ResponsesRequest
from ember.models.schemas import ChatResponse, UsageStats
from ember.non import build_graph
from ember.operators.base import Operator
from keiro.ensemble.eb1 import Eb1Graph, PromptContext


class RouterAction(str, Enum):
    """Actions that the controller may request."""

    NANO = "nano"
    SELF = "self"
    GRAPH = "graph"


class RouterDirective(BaseModel):
    """Schema the controller must satisfy when issuing directives."""

    model_config = ConfigDict(validate_by_name=True)

    action: RouterAction
    reason: str = Field(min_length=1)
    continue_loop: bool = Field(default=False, alias="continue")
    latency_budget_ms: float | None = Field(default=None, gt=0)
    params: dict[str, object] = Field(default_factory=dict)
    graph_spec: str | None = None
    graph_flows: list[str] | None = None
    final_answer: str | None = None


@dataclass
class RoutedCall:
    """Single invocation captured in the decision trace."""

    step: int
    action: str
    model: str | None
    latency_ms: float
    extra: MutableMapping[str, object] = field(default_factory=dict)


@dataclass
class RouterResult:
    """Final outcome produced by the benchmark router."""

    response: ChatResponse
    decision_trace: Sequence[RoutedCall]
    directive_log: Sequence[RouterDirective]
    usage_trace: Sequence[Mapping[str, object]]


_MIN_BUDGET_MS: float = 120_000.0
"""Floor for controller-issued latency budgets (120 s).

The controller is an LLM that may hallucinate arbitrarily tight budgets
(e.g. 300 ms) that are physically impossible for even a single downstream
LLM call.  Any value below this floor is clamped up.
"""

_MAX_ITERATIONS: int = 8
"""Hard cap on routing loop iterations to prevent runaway loops."""


def _enforce_budget(started_at: float, budget_ms: float) -> None:
    """Raise if elapsed wall time exceeds the latency budget."""
    elapsed_ms = (time.perf_counter() - started_at) * 1000.0
    if elapsed_ms > budget_ms:
        raise RuntimeError(
            f"Router exceeded latency budget: "
            f"{elapsed_ms:.1f}ms > {budget_ms:.1f}ms"
        )


class GPT5HighRouter(Operator):
    """Operator that orchestrates adaptive routing with gpt-5-high."""

    controller_model: ModelBinding
    cheap_model: ModelBinding

    def __init__(
        self,
        *,
        controller_model: str = "gpt-5-high",
        cheap_model: str = "gpt-5-nano",
        controller_params: Mapping[str, object] | None = None,
        cheap_params: Mapping[str, object] | None = None,
        controller_binding: ModelBinding | None = None,
        cheap_binding: ModelBinding | None = None,
    ) -> None:
        from ember.api.models import models as _models_facade

        if controller_binding is not None:
            self.controller_model = controller_binding
        else:
            self.controller_model = _models_facade.instance(
                controller_model,
                **(controller_params or {}),
            )

        if cheap_binding is not None:
            self.cheap_model = cheap_binding
        else:
            self.cheap_model = _models_facade.instance(
                cheap_model, **(cheap_params or {})
            )

    def forward(
        self,
        prompt: str,
        *,
        system: str | None = None,
        metadata: Mapping[str, object] | None = None,
        shared_options: Mapping[str, object] | None = None,
        request_options: Mapping[str, object] | None = None,
        responses_request: ResponsesRequest | None = None,
        context_messages: Sequence[Mapping[str, str]] | None = None,
    ) -> RouterResult:
        decision_trace: list[RoutedCall] = []
        directives: list[RouterDirective] = []
        usage_trace: list[Mapping[str, object]] = []
        routing_started_at = time.perf_counter()
        active_budget_ms: float | None = None

        base_options: dict[str, object] = {
            **(shared_options or {}),
            **(request_options or {}),
        }

        controller_prompt, controller_system = self._build_controller_prompt(
            prompt,
            system,
            metadata,
        )
        latest_response: ChatResponse | None = None
        loop_counter = 0
        total_usage: UsageStats | None = None

        def record_usage(label: str, response: ChatResponse | None) -> None:
            nonlocal total_usage
            if response is None or response.usage is None:
                usage_trace.append({"stage": label, "usage": None})
                return
            if total_usage is None:
                total_usage = UsageStats()
            total_usage.add(response.usage)
            usage_trace.append(
                {
                    "stage": label,
                    "prompt_tokens": response.usage.prompt_tokens,
                    "completion_tokens": response.usage.completion_tokens,
                    "total_tokens": response.usage.total_tokens,
                    "cost_usd": response.usage.cost_usd,
                    "actual_cost_usd": response.usage.actual_cost_usd,
                }
            )

        while True:
            if active_budget_ms is not None:
                _enforce_budget(routing_started_at, active_budget_ms)

            loop_counter += 1
            call, directive, controller_response = self._invoke_controller(
                len(decision_trace) + 1,
                controller_prompt,
                controller_system,
            )
            decision_trace.append(call)
            directives.append(directive)
            record_usage("controller", controller_response)
            if directive.latency_budget_ms is not None:
                clamped = max(directive.latency_budget_ms, _MIN_BUDGET_MS)
                if active_budget_ms is None:
                    active_budget_ms = clamped
                else:
                    active_budget_ms = min(active_budget_ms, clamped)

            if directive.action is RouterAction.SELF:
                latest_response = controller_response
            elif directive.action is RouterAction.NANO:
                merged_params = self._merge_action_params(base_options, directive.params)
                latest_response = self._run_cheap(
                    prompt,
                    system,
                    merged_params,
                    decision_trace,
                )
            elif directive.action is RouterAction.GRAPH:
                latest_response = self._run_graph(
                    prompt,
                    system,
                    self._merge_action_params(base_options, directive.params),
                    directive.graph_spec,
                    directive.graph_flows,
                    shared_options,
                    request_options,
                    responses_request,
                    context_messages,
                    decision_trace,
                )
            else:  # pragma: no cover
                raise RuntimeError(f"Unsupported action: {directive.action}")

            if active_budget_ms is not None:
                _enforce_budget(routing_started_at, active_budget_ms)

            if not directive.continue_loop:
                if directive.action is RouterAction.SELF and not directive.final_answer:
                    raise RuntimeError(
                        "Controller self action with continue=false must include final_answer."
                    )
                final = directive.final_answer or (latest_response.text if latest_response else "")
                if latest_response:
                    latest_response = ChatResponse(
                        data=final,
                        output=latest_response.output,
                        response_id=latest_response.response_id,
                        usage=latest_response.usage,
                        model_id=latest_response.model_id,
                        raw_output=latest_response.raw_output,
                        created_at=latest_response.created_at,
                        latency_ms=latest_response.latency_ms,
                        started_at=latest_response.started_at,
                        completed_at=latest_response.completed_at,
                    )
                else:
                    latest_response = ChatResponse(data=final)
                break

            controller_prompt, controller_system = self._augment_prompt(
                controller_prompt,
                controller_system,
                directive,
                latest_response,
            )
            if loop_counter >= _MAX_ITERATIONS:
                raise RuntimeError(
                    f"Router exceeded maximum allowed iterations ({_MAX_ITERATIONS})"
                )

        if latest_response is None:
            raise RuntimeError("Router did not produce a response")

        record_usage("final", latest_response)
        if total_usage is not None:
            latest_response.usage = total_usage

        return RouterResult(
            response=latest_response,
            decision_trace=tuple(decision_trace),
            directive_log=tuple(directives),
            usage_trace=tuple(usage_trace),
        )

    def _build_controller_prompt(
        self,
        prompt: str,
        system: str | None,
        metadata: Mapping[str, object] | None,
    ) -> tuple[str, str]:
        controller_system = self._controller_system_prompt()
        user_payload = {
            "prompt": prompt,
            "system": system,
            "metadata": metadata or {},
        }
        user_text = json.dumps(user_payload)
        controller_prompt = f"<<request>>\n{user_text}\n<<end-request>>\n"
        return controller_prompt, controller_system

    def _augment_prompt(
        self,
        prompt_text: str,
        system_text: str,
        directive: RouterDirective,
        latest_response: ChatResponse | None,
    ) -> tuple[str, str]:
        history_blocks = [prompt_text.rstrip()]
        directive_payload = json.dumps(directive.model_dump(by_alias=True))
        history_blocks.append(f"<<directive>>\n{directive_payload}\n<<end-directive>>")
        if latest_response is not None:
            result_payload = {
                "type": "result",
                "model": latest_response.model_id,
                "output": latest_response.data,
            }
            history_blocks.append(f"<<result>>\n{json.dumps(result_payload)}\n<<end-result>>")
        augmented_prompt = "\n".join(history_blocks) + "\n"
        return augmented_prompt, system_text

    def _controller_system_prompt(self) -> str:
        return (
            "You are the routing controller for the Ember benchmark. "
            "Decide among three actions:\n"
            "1. 'nano' -> delegate to gpt-5-nano for simple tasks.\n"
            "2. 'self' -> answer directly with your current completion.\n"
            "3. 'graph' -> construct a compact NON graph specification to run "
            "multi-model ensembles with a judge.\n"
            "Always return a JSON object with fields: action, reason, continue, "
            "latency_budget_ms (optional), params (dict), graph_spec (string when "
            "action is graph), graph_flows (list of strings, optional), and final_answer "
            "(string when continue=false)."
        )

    def _invoke_controller(
        self,
        step: int,
        prompt_text: str,
        system_text: str | None,
    ) -> tuple[RoutedCall, RouterDirective, ChatResponse]:
        start = time.perf_counter()
        response_wrapper = self.controller_model.response(
            prompt_text,
            system=system_text,
        )
        latency_ms = (time.perf_counter() - start) * 1000.0
        chat_response = response_wrapper.chat_response

        directive = self._parse_directive(chat_response.data)
        trace_entry = RoutedCall(
            step=step,
            action="controller",
            model=chat_response.model_id,
            latency_ms=latency_ms,
        )
        return trace_entry, directive, chat_response

    def _parse_directive(self, payload: str) -> RouterDirective:
        try:
            data = json.loads(payload)
        except json.JSONDecodeError as exc:
            raise RuntimeError(f"Controller returned invalid JSON: {payload[:200]}") from exc
        try:
            return RouterDirective.model_validate(data)
        except ValidationError as exc:
            raise RuntimeError(f"Controller directive failed validation: {exc}") from exc

    def _run_cheap(
        self,
        prompt: str,
        system: str | None,
        params: Mapping[str, object],
        trace: list[RoutedCall],
    ) -> ChatResponse:
        start = time.perf_counter()
        clean_params = {
            k: v
            for k, v in params.items()
            if k not in ("prompt", "system", "model")
        }
        response_wrapper = self.cheap_model.response(
            prompt, system=system, **clean_params
        )
        latency_ms = (time.perf_counter() - start) * 1000.0
        trace.append(
            RoutedCall(
                step=len(trace) + 1,
                action="nano",
                model=response_wrapper.model_id,
                latency_ms=latency_ms,
            )
        )
        return response_wrapper.chat_response

    def _run_graph(
        self,
        prompt: str,
        system: str | None,
        params: Mapping[str, object],
        graph_spec: str | None,
        graph_flows: Sequence[str] | None,
        shared_options: Mapping[str, object] | None,
        request_options: Mapping[str, object] | None,
        responses_request: ResponsesRequest | None,
        context_messages: Sequence[Mapping[str, str]] | None,
        trace: list[RoutedCall],
    ) -> ChatResponse:
        if not graph_spec:
            raise RuntimeError("Graph action requires graph_spec")

        spec_lines = self._split_graph_spec(graph_spec)
        graph = build_graph(spec_lines, flows=graph_flows)
        plan = graph.to_eb1_plan()
        runner = Eb1Graph(plan, shared_options=shared_options)
        context = PromptContext(
            prompt=prompt,
            system=system,
            shared_options=params,
            request_options=request_options,
            responses_request=responses_request,
            context_messages=tuple(context_messages or ()),
        )

        start = time.perf_counter()
        result = runner(context)
        latency_ms = (time.perf_counter() - start) * 1000.0
        trace.append(
            RoutedCall(
                step=len(trace) + 1,
                action="graph",
                model=result.model_id,
                latency_ms=latency_ms,
            )
        )
        if params:
            result.raw_output = self._merge_raw_output(result.raw_output, params)
        return result

    @staticmethod
    def _merge_action_params(
        base: Mapping[str, object],
        directive_params: Mapping[str, object],
    ) -> dict[str, object]:
        merged: dict[str, object] = dict(base)
        merged.update(directive_params)
        return merged

    def _split_graph_spec(self, spec: str) -> list[str]:
        lines = [line.strip() for line in spec.splitlines() if line.strip()]
        if not lines:
            raise RuntimeError("graph_spec must include at least one line")
        if len(lines) > 8:
            raise RuntimeError("graph_spec exceeds maximum supported nodes (8 lines)")
        return lines

    @staticmethod
    def _merge_raw_output(
        raw_output: object | None,
        params: Mapping[str, object],
    ) -> Mapping[str, object]:
        if isinstance(raw_output, Mapping):
            merged: dict[str, object] = dict(raw_output)
        else:
            merged = {"raw_output": raw_output}

        existing = merged.get("directive_params")
        if existing is None:
            merged["directive_params"] = dict(params)
        elif isinstance(existing, dict):
            existing.update(params)
        else:
            raise TypeError("directive_params must be a dict when present")
        return merged


__all__ = ["GPT5HighRouter", "RouterDirective", "RouterResult", "RoutedCall"]
