"""EB1 ensemble engine implementing the :class:`EB1Engine` protocol.

This module bridges the Keiro protocol layer to the existing
``ember.takumi.eb1.run_plan()`` function.  During the migration period
both call paths coexist: callers that already use the Keiro protocol go
through ``Eb1Ensemble``, while legacy callers continue to use
``run_plan()`` directly.

``run_plan()`` accepts an optional ``ModelInvoker`` parameter. When
``Eb1Ensemble`` provides its invoker, model calls route through the
protocol boundary instead of ``ember.api.models``.

Example::

    from keiro.ensemble import Eb1Ensemble
    from keiro.protocol import EB1Request

    engine = Eb1Ensemble()
    result = engine.run(EB1Request(prompt="Explain quantum computing"))
    print(result.text)
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Iterator

from ember.models.schemas import UsageStats
from .eb1 import (
    EB1_IDENTITY_DIRECTIVE,
    Eb1Plan,
    run_plan,
)
from .plans import select_plan

from keiro.protocol import (
    EB1Request,
    EB1Result,
    EB1StreamChunk,
    parse_responses_tool_calls,
    responses_output_items,
)

if TYPE_CHECKING:
    from keiro.protocol import ModelInvoker

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Identity directive injection
# ---------------------------------------------------------------------------


def _ensure_identity(system: str | None) -> str:
    """Prepend the EB1 identity directive if absent.

    The identity directive ensures the judge always identifies as EB1/Keiro.
    This mirrors ``_system_prompt_for_eb1()`` in the legacy adapter but
    lives in the ensemble module where it belongs.

    Args:
        system: Caller-provided system prompt, or None.

    Returns:
        System prompt with identity directive prepended.
    """
    if system and EB1_IDENTITY_DIRECTIVE in system:
        return system
    prefix = EB1_IDENTITY_DIRECTIVE
    if system:
        return f"{prefix}\n\n{system}"
    return prefix


# ---------------------------------------------------------------------------
# Eb1Ensemble
# ---------------------------------------------------------------------------


class Eb1Ensemble:
    """EB1 ensemble engine satisfying the :class:`EB1Engine` protocol.

    Translates :class:`EB1Request` into ``run_plan()`` arguments, executes
    the ensemble, and converts the result into :class:`EB1Result`.

    During migration, this delegates to the existing ``run_plan()`` function.
    When ``invoker`` is provided, all model calls route through the
    ``ModelInvoker`` protocol instead of the ``ember.api.models`` facade.

    Args:
        invoker: Model invoker for upstream calls. When provided, all
            candidate and judge model calls route through this invoker.
            When None, ``run_plan()`` falls back to the ember models facade.
    """

    def __init__(
        self,
        *,
        invoker: ModelInvoker | None = None,
    ) -> None:
        self._invoker = invoker

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

        Selects a plan based on ``request.model``, injects the identity
        directive, and delegates to ``run_plan()``.

        Args:
            request: Fully specified ensemble request.

        Returns:
            Synthesized result with aggregated usage and metadata.

        Raises:
            ProviderAPIError: If all candidates fail.
            ValueError: If the request is malformed.
        """
        plan = select_plan(request.model)
        response = self._execute(plan, request)
        return self._to_result(response, request.model)

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

        Currently falls back to synchronous execution and yields the
        full result as a single chunk.  True judge streaming will be
        wired when the ensemble is refactored to accept ``ModelInvoker``.

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

        Yields:
            Chunks of the judge's synthesized response.
        """
        # Phase 1: synchronous fallback.  True streaming requires
        # refactoring _execute_plan() to accept a streaming ModelInvoker,
        # which is Step 5 of the migration.
        result = self.run(request)
        yield EB1StreamChunk(
            text=result.text,
            usage_delta=result.usage,
            finish_reason=result.finish_reason,
            metadata=result.metadata,
        )

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _execute(self, plan: Eb1Plan, request: EB1Request) -> object:
        """Bridge EB1Request to run_plan() and return the ChatResponse.

        Args:
            plan: Selected execution plan.
            request: Protocol request to translate.

        Returns:
            ChatResponse from run_plan().
        """
        system = _ensure_identity(request.system)

        shared_options: dict = {}
        if request.reasoning:
            shared_options["reasoning"] = request.reasoning

        request_options: dict = {}
        if request.response_format:
            request_options["response_format"] = request.response_format
        if request.max_output_tokens is not None:
            request_options["max_output_tokens"] = request.max_output_tokens

        context_messages: list[dict] = [dict(m) for m in request.context_messages]

        return run_plan(
            plan,
            request.prompt,
            system=system,
            shared_options=shared_options or None,
            request_options=request_options or None,
            context_messages=context_messages or None,
            invoker=self._invoker,
        )

    @staticmethod
    def _to_result(response: object, model: str) -> EB1Result:
        """Convert a ChatResponse to an EB1Result.

        Args:
            response: ChatResponse from run_plan().
            model: User-facing model identifier.

        Returns:
            Protocol-conformant EB1Result.
        """
        # ChatResponse is not typed here to avoid importing it in the
        # protocol module.  Access attributes dynamically during the
        # migration period.
        text: str = getattr(response, "data", "") or ""
        usage: UsageStats = getattr(response, "usage", None) or UsageStats()
        raw_output = getattr(response, "raw_output", None)

        # Build metadata from raw_output for debugging.
        metadata: dict = {}
        if isinstance(raw_output, dict):
            for key in (
                "plan",
                "candidates",
                "judge_model",
                "internal_usage",
                "mode",
            ):
                if key in raw_output:
                    metadata[key] = raw_output[key]

        items = responses_output_items(raw_output)
        tool_calls = parse_responses_tool_calls(items)

        finish_reason = "tool_calls" if tool_calls else "stop"

        return EB1Result(
            text=text,
            model=model,
            usage=usage,
            tool_calls=tool_calls,
            finish_reason=finish_reason,
            metadata=metadata,
        )
