"""Shared interfaces between Keiro modules.

This module defines the contracts that decouple the bounded modules:

    Gateway  --EB1Engine-->  Ensemble  --ModelInvoker-->  Upstream

All types here are either frozen dataclasses (value objects) or Protocols
(structural interfaces). No implementation logic. No side effects.
Modules depend on these types, never on each other's concrete classes.

The dependency rule: keiro.protocol imports from ember.models.schemas
(the open-source layer) for UsageStats. It never imports from any keiro
submodule.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Iterator, Mapping, Protocol, Sequence, runtime_checkable

from ember.models.schemas import UsageStats

# ---------------------------------------------------------------------------
# SSE error event types
# ---------------------------------------------------------------------------
# The Responses API uses ``"response.error"``; the GNN server historically
# emitted bare ``"error"``.  All layers (client, gateway, model_api) must
# accept BOTH to avoid silent empty responses.
SSE_ERROR_TYPES: frozenset[str] = frozenset({"response.error", "error"})


# ---------------------------------------------------------------------------
# ToolCall: normalized tool/function call type
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class ToolCall:
    """A single tool/function call in normalized form.

    Attributes:
        id: Provider-assigned call identifier (``call_id`` or ``id``).
        name: Function name to invoke.
        arguments: Function arguments as a JSON string.
    """

    id: str
    name: str
    arguments: str

    def to_chat_dict(self) -> dict[str, object]:
        """Serialize to OpenAI Chat Completions tool_call format."""
        return {
            "id": self.id,
            "type": "function",
            "function": {"name": self.name, "arguments": self.arguments},
        }


# ---------------------------------------------------------------------------
# Shared parsers for Responses API tool call extraction
# ---------------------------------------------------------------------------


def responses_output_items(raw_output: object) -> Sequence[Mapping[str, object]]:
    """Drill into a ChatResponse raw_output to find Responses API output items.

    Handles the common shape::

        raw_output["responses_payload"]["output"] -> list[dict]

    Args:
        raw_output: The ``ChatResponse.raw_output`` attribute (dict or None).

    Returns:
        The output items list, or an empty sequence if the path is absent.
    """
    if not isinstance(raw_output, dict):
        return ()
    payload = raw_output.get("responses_payload")
    if not isinstance(payload, dict):
        return ()
    output = payload.get("output")
    if not isinstance(output, list):
        return ()
    return output


def parse_responses_tool_calls(
    items: Sequence[Mapping[str, object]],
) -> tuple[ToolCall, ...]:
    """Parse Responses API ``function_call`` items into ToolCall objects.

    Filters *items* to those with ``type == "function_call"`` and normalizes
    the ``call_id``/``id`` field ambiguity.

    Args:
        items: Output items from a Responses API payload.

    Returns:
        Tuple of parsed ToolCall objects.
    """
    calls: list[ToolCall] = []
    for item in items:
        if not isinstance(item, Mapping):
            continue
        if item.get("type") != "function_call":
            continue
        call_id = item.get("call_id") or item.get("id") or ""
        calls.append(
            ToolCall(
                id=str(call_id),
                name=str(item.get("name", "")),
                arguments=str(item.get("arguments", "{}")),
            )
        )
    return tuple(calls)


def merge_tool_calls_into_responses_output(
    output: Sequence[Mapping[str, object]] | None,
    tool_calls: Sequence[Mapping[str, Any]] | None,
) -> list[dict[str, object]]:
    """Return ``output`` with any missing chat-style tool calls added.

    Existing ``function_call`` items win. Missing calls are appended using
    the Responses API item shape.
    """
    merged = [dict(item) for item in output or () if isinstance(item, Mapping)]
    if not tool_calls:
        return merged

    existing_ids = {
        call.id for call in parse_responses_tool_calls(merged) if call.id
    }
    for tool_call in tool_calls:
        if not isinstance(tool_call, Mapping):
            continue
        call_id = str(tool_call.get("id") or tool_call.get("call_id") or "")
        if call_id and call_id in existing_ids:
            continue
        function = tool_call.get("function")
        merged.append(
            {
                "type": "function_call",
                "call_id": call_id,
                "name": (
                    str(function.get("name", ""))
                    if isinstance(function, Mapping) else ""
                ),
                "arguments": (
                    str(function.get("arguments", "{}"))
                    if isinstance(function, Mapping) else "{}"
                ),
                "status": "completed",
            }
        )
        if call_id:
            existing_ids.add(call_id)
    return merged


# ---------------------------------------------------------------------------
# Value objects: Gateway <-> Ensemble
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class EB1Request:
    """What the gateway sends to the ensemble engine.

    This is the only type that crosses the gateway/ensemble boundary.
    It contains everything the ensemble needs to execute, with no HTTP
    or transport concerns.

    Attributes:
        prompt: The user's input text.
        system: Optional system prompt. The ensemble prepends the EB1
            identity directive automatically if not already present.
        context_messages: Prior conversation turns as dicts with 'role'
            and 'content' keys. Empty tuple for single-turn requests.
        model: Plan selector string. Maps to an ensemble plan internally.
            Examples: 'eb1-preview', 'eb1-pro', 'eb1-fast', 'eb1-codex'.
        tools: Tool definitions in OpenAI function-calling format.
            Candidates receive tools to produce tool_calls; the judge
            evaluates candidates without tools (verifier mode).
        tool_choice: Tool selection constraint ('auto', 'required',
            'none', or a specific tool dict). Forwarded to candidates.
        response_format: Structured output schema (JSON mode).
        reasoning: Reasoning effort configuration for supporting models.
            Example: {'effort': 'high'}.
        max_output_tokens: Hard cap on judge output length.
        stream: Whether the judge should stream its response.
    """

    prompt: str
    system: str | None = None
    context_messages: tuple[dict, ...] = ()
    model: str = "eb1-preview"
    tools: tuple[dict, ...] | None = None
    tool_choice: str | dict | None = None
    response_format: dict | None = None
    reasoning: dict | None = None
    max_output_tokens: int | None = None
    stream: bool = False


@dataclass(frozen=True)
class QualitySignal:
    """Structured quality signal for ensemble responses.

    Always present on EB1 responses. Tells the caller exactly what
    happened during ensemble execution: was the judge invoked? How
    many candidates ran? Were there any degradations?

    The signal is surfaced in the HTTP response as ``eb1_quality``
    so benchmark harnesses can detect degraded results.

    Attributes:
        judge_invoked: Whether the judge was actually called. False
            when budget expired, judge crashed, or context overflowed.
        synthesis_method: How the final response was produced.
            ``"judge_synthesis"`` (judge rewrote), ``"judge_verification"``
            (judge selected a candidate verbatim), or
            ``"candidate_fallback"`` (judge skipped, raw candidate).
        ensemble_breadth: Number of candidates the judge actually saw.
            May be less than ``candidates_attempted`` on timeouts.
        candidates_attempted: Configured candidate count for the plan.
        candidates_succeeded: Candidates that completed before timeout.
        degradations: Typed degradation events that occurred during
            execution. Each entry has ``kind`` (enum string) and
            ``detail`` (human-readable explanation). Empty tuple when
            the ensemble executed at full quality.
    """

    judge_invoked: bool = True
    synthesis_method: str = "judge_synthesis"
    ensemble_breadth: int = 0
    candidates_attempted: int = 0
    candidates_succeeded: int = 0
    degradations: tuple[dict[str, str], ...] = ()

    def to_dict(self) -> dict[str, object]:
        """Serialize to a plain dict for inclusion in raw_output."""
        return {
            "judge_invoked": self.judge_invoked,
            "synthesis_method": self.synthesis_method,
            "ensemble_breadth": self.ensemble_breadth,
            "candidates_attempted": self.candidates_attempted,
            "candidates_succeeded": self.candidates_succeeded,
            "degradations": list(self.degradations),
        }

    @classmethod
    def from_dict(cls, data: Mapping[str, object]) -> "QualitySignal":
        """Reconstruct from a raw_output dict with type coercion."""
        return cls(
            judge_invoked=bool(data.get("judge_invoked", True)),
            synthesis_method=str(
                data.get("synthesis_method", "judge_synthesis"),
            ),
            ensemble_breadth=int(data.get("ensemble_breadth", 0)),
            candidates_attempted=int(
                data.get("candidates_attempted", 0),
            ),
            candidates_succeeded=int(
                data.get("candidates_succeeded", 0),
            ),
            degradations=tuple(data.get("degradations", ())),
        )


@dataclass(frozen=True)
class EB1Result:
    """What the ensemble engine returns to the gateway.

    Contains the synthesized response and metadata about the ensemble
    execution. The gateway converts this into the appropriate HTTP
    response format (Chat Completions, Responses API, or Messages).

    Attributes:
        text: The judge's synthesized response text.
        model: Canonical model identifier (e.g., 'eb1-preview').
        usage: Aggregated token usage across all candidates and judge.
        tool_calls: Tool calls extracted from judge output, if any.
        finish_reason: Why generation stopped ('stop', 'tool_calls',
            'length', 'content_filter').
        metadata: Ensemble execution details for debugging. Includes
            plan shape, candidate models, latencies, and cost breakdown.
            The gateway may strip sensitive fields for non-admin users.
        quality: Structured quality signal describing ensemble execution
            integrity. Always populated; check ``degradations`` for
            any quality-reducing events.
    """

    text: str
    model: str
    usage: UsageStats
    tool_calls: tuple[ToolCall, ...] = ()
    finish_reason: str = "stop"
    metadata: dict = field(default_factory=dict)
    quality: QualitySignal = field(default_factory=QualitySignal)


@dataclass(frozen=True)
class EB1StreamChunk:
    """One chunk in a streaming EB1 response.

    During streaming, candidates execute to completion (buffered), then
    the judge streams its synthesis. Each chunk carries incremental text
    and optional usage deltas for per-chunk billing.

    Attributes:
        text: Incremental text content. Empty string for non-text events.
        usage_delta: Incremental token counts for this chunk, if available.
        finish_reason: Set on the final chunk to indicate completion.
        metadata: Per-chunk metadata (sparse; most chunks have none).
    """

    text: str = ""
    usage_delta: UsageStats | None = None
    finish_reason: str | None = None
    metadata: dict = field(default_factory=dict)


# ---------------------------------------------------------------------------
# Value objects: Ensemble <-> Upstream
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class ModelCall:
    """What the ensemble sends to invoke one upstream model.

    Represents a single provider API call. The ensemble creates one per
    candidate plus one for the judge. The upstream module routes it to
    the correct provider, applies rate limiting, and selects credentials.

    Attributes:
        model_id: Canonical model identifier (e.g., 'gpt-5.2',
            'claude-opus-4-6', 'gemini-3-pro-preview').
        prompt: The prompt text for this specific model call.
        system: System prompt for this call. Candidates get a persona
            prompt; the judge gets the synthesis instructions.
        context_messages: Conversation history as dicts.
        tools: Tool definitions for function calling.
        tool_choice: Tool selection constraint ('auto', 'required', 'none',
            or a specific tool dict).
        response_format: Structured output schema.
        reasoning: Reasoning effort configuration.
        max_output_tokens: Output length cap for this call.
    """

    model_id: str
    prompt: str
    system: str | None = None
    context_messages: tuple[dict, ...] = ()
    tools: tuple[dict, ...] | None = None
    tool_choice: str | dict | None = None
    response_format: dict | None = None
    reasoning: dict | None = None
    max_output_tokens: int | None = None


@dataclass(frozen=True)
class ModelResult:
    """What comes back from one upstream model call.

    The ensemble collects these from candidates, renders them into
    the judge prompt, then invokes the judge to produce the final
    EB1Result.

    Attributes:
        text: The model's response text.
        model_id: The model that actually served the request (may differ
            from ModelCall.model_id if aliased).
        usage: Token usage for this individual call.
        tool_calls: Tool calls returned by the model, if any.
        finish_reason: Why this model stopped generating.
        latency_ms: Wall-clock time for this call in milliseconds.
        output_items: Responses API output items from the provider
            response, if available.
        thinking_trace: Extended thinking trace, if the model supports it.
        metadata: Provider metadata (model version, request ID, etc.).
    """

    text: str
    model_id: str
    usage: UsageStats
    tool_calls: tuple[ToolCall, ...] = ()
    finish_reason: str = "stop"
    latency_ms: float = 0.0
    output_items: tuple[Mapping[str, object], ...] | None = None
    thinking_trace: str | None = None
    metadata: Mapping[str, object] = field(default_factory=dict)


# ---------------------------------------------------------------------------
# Protocol: EB1Engine (Gateway -> Ensemble)
# ---------------------------------------------------------------------------


@runtime_checkable
class EB1Engine(Protocol):
    """Interface the gateway uses to invoke the ensemble.

    Implementations: ``keiro.ensemble.Eb1Ensemble``.

    The gateway never imports ensemble internals. It constructs an
    ``EB1Request``, calls ``run()`` or ``stream()``, and receives
    ``EB1Result`` or ``EB1StreamChunk`` values. All plan selection,
    candidate orchestration, and judge synthesis happen behind this
    interface.
    """

    def run(self, request: EB1Request) -> EB1Result:
        """Execute the ensemble synchronously.

        Args:
            request: Fully specified ensemble request.

        Returns:
            Synthesized result with aggregated usage.

        Raises:
            ProviderAPIError: If all candidates fail.
            ValueError: If the request is malformed.
        """
        ...

    def stream(self, request: EB1Request) -> Iterator[EB1StreamChunk]:
        """Execute the ensemble with streaming judge output.

        Candidates run to completion (buffered), then the judge streams
        its synthesis. The final chunk has ``finish_reason`` set.

        Args:
            request: Fully specified ensemble request with ``stream=True``.

        Yields:
            Incremental chunks of the judge's synthesized response.

        Raises:
            ProviderAPIError: If all candidates fail.
        """
        ...


# ---------------------------------------------------------------------------
# Protocol: ModelInvoker (Ensemble -> Upstream)
# ---------------------------------------------------------------------------


@runtime_checkable
class ModelInvoker(Protocol):
    """Interface the ensemble uses to call upstream providers.

    Implementations: ``keiro.upstream.UpstreamPool``.

    The ensemble never imports provider SDKs, credential stores, or
    rate limiters. It calls ``invoke()`` and gets back a ``ModelResult``.
    The upstream module handles credential selection, AIMD rate limiting,
    retries, and provider SDK translation behind this interface.
    """

    def invoke(self, call: ModelCall) -> ModelResult:
        """Call a single upstream model synchronously.

        Args:
            call: Model call specification.

        Returns:
            Model response with usage and latency.

        Raises:
            ProviderAPIError: If the provider returns an error after
                retries are exhausted.
        """
        ...

    def invoke_stream(self, call: ModelCall) -> Iterator[ModelResult]:
        """Call a single upstream model with streaming.

        Each yielded ``ModelResult`` contains an incremental text chunk.
        The final result has ``finish_reason`` set and complete usage.

        Args:
            call: Model call specification.

        Yields:
            Incremental model response chunks.
        """
        ...

    def available(self, model_id: str) -> bool:
        """Check if a model can be called right now.

        Returns False if credentials are missing or the provider's
        AIMD limiter has no capacity. Used by the ensemble during
        plan compilation to skip unavailable candidates.

        Args:
            model_id: Canonical model identifier.

        Returns:
            True if the model can likely be invoked successfully.
        """
        ...
