"""Local EB1 ensemble provider — calls upstream models directly.

Eliminates the gateway HTTP indirection that ``TakumiProvider`` uses.
EB1 ensemble logic already lives in ``ember.takumi.eb1``; this provider
simply exposes it through the standard :class:`BaseProvider` interface so
the model registry can route ``eb1-*`` models locally.

Two execution paths:

* **Ensemble** (eb1, eb1-preview, eb1-pro, eb1-codex, ...): delegates to
  :func:`~ember.takumi.eb1.run_plan` with the plan from
  :func:`~ember.takumi.plans.select_plan`.
* **SWE routing** (eb1-swe-preview, eb1-swe-fast-preview, ...): classifies
  query complexity via :func:`~ember.gateway.eb1_resolver.resolve_eb1_swe_model`
  then invokes the resolved concrete model through the registry.
"""

from __future__ import annotations

import logging
from collections.abc import Mapping, Sequence
from typing import Any, Generator

from ember._internal.exceptions import ProviderAPIError
from ember.models.providers.base import BaseProvider
from ember.models.providers.utils import (
    ResponsesMessage,
    build_prompt_components,
    parse_responses_request,
)
from ember.models.schemas import ChatResponse

_LOG = logging.getLogger(__name__)


class Eb1EnsembleProvider(BaseProvider):
    """Execute EB1 ensemble plans locally without a gateway HTTP hop.

    Follows the same contract as :class:`TakumiEb1ARuntimeProvider` but
    handles the full range of ``eb1-*`` and ``eb2-*`` model names.
    """

    requires_api_key = False
    supports_responses_api = True
    supports_streaming = True

    # ------------------------------------------------------------------
    # BaseProvider interface
    # ------------------------------------------------------------------

    def complete(self, prompt: str, model: str, **kwargs: Any) -> ChatResponse:
        params = self._extract_invocation_kwargs(kwargs)
        return self._run_ensemble_from_prompt(model, prompt, params)

    def stream_complete(
        self,
        prompt: str,
        model: str,
        **kwargs: Any,
    ) -> Generator[str, None, ChatResponse]:
        response = self.complete(prompt, model, **kwargs)
        text = response.text or ""
        if text:
            yield text
        return response

    def complete_responses_payload(
        self,
        payload: Mapping[str, Any],
        **kwargs: Any,
    ) -> ChatResponse:
        if not isinstance(payload, Mapping):
            raise ProviderAPIError(
                "Responses payload must be a mapping",
                context={"provider": "eb1_ensemble"},
            )

        request = parse_responses_request(payload)
        model = request.model or "eb1"

        if self._is_swe_model(model):
            return self._run_swe_routing(model, request, payload, **kwargs)

        # Tool-use bypass: the ensemble strips tools from candidates and
        # produces text synthesis, which is incompatible with agentic tool
        # calling.  When tools are present, route to the judge model
        # directly so the provider handles tool_use natively.
        if request.tools:
            return self._run_tool_use_passthrough(model, payload)

        components = build_prompt_components(request)
        prompt_text = (
            self._prompt_from_responses_input(request.input)
            or components.prompt
            or ""
        )

        context_messages: tuple[Mapping[str, str], ...] = tuple(
            {"role": entry["role"], "content": entry["content"]}
            for entry in components.context
        )
        if not context_messages:
            context_messages = self._context_from_responses_messages(
                request.messages,
            )

        merged_kwargs: dict[str, Any] = dict(kwargs)
        merged_kwargs.setdefault("system", components.system)
        merged_kwargs["responses_request"] = request

        if request.metadata:
            metadata_payload = dict(request.metadata)
            existing = merged_kwargs.get("metadata")
            if isinstance(existing, Mapping):
                combined = dict(metadata_payload)
                combined.update(dict(existing))
                merged_kwargs["metadata"] = combined
            else:
                merged_kwargs.setdefault("metadata", metadata_payload)

        shared_options = dict(merged_kwargs.get("shared_options") or {})
        if request.reasoning and "reasoning" not in shared_options:
            shared_options["reasoning"] = dict(request.reasoning)
        if request.temperature is not None and "temperature" not in shared_options:
            shared_options["temperature"] = request.temperature
        if request.top_p is not None and "top_p" not in shared_options:
            shared_options["top_p"] = request.top_p
        if shared_options:
            merged_kwargs["shared_options"] = shared_options

        request_options = dict(merged_kwargs.get("request_options") or {})
        if request.response_format and "response_format" not in request_options:
            request_options["response_format"] = dict(request.response_format)
        if request.max_output_tokens is not None and "max_output_tokens" not in request_options:
            request_options["max_output_tokens"] = request.max_output_tokens
        if request_options:
            merged_kwargs["request_options"] = request_options

        if context_messages:
            merged_kwargs.setdefault(
                "context_messages",
                tuple(dict(entry) for entry in context_messages),
            )

        params = self._extract_invocation_kwargs(merged_kwargs)
        return self._run_ensemble_from_prompt(model, prompt_text, params)

    # ------------------------------------------------------------------
    # Ensemble execution
    # ------------------------------------------------------------------

    def _run_ensemble_from_prompt(
        self,
        model: str,
        prompt: str,
        params: dict[str, Any],
    ) -> ChatResponse:
        """Run the full EB1 ensemble: candidates + judge synthesis."""
        from keiro.ensemble.eb1 import run_plan
        from keiro.ensemble.plans import select_plan

        plan = select_plan(model)
        response = run_plan(
            plan,
            prompt,
            system=params["system"],
            shared_options=params["shared_options"],
            request_options=params["request_options"],
            responses_request=params["responses_request"],
            context_messages=params["context_messages"],
        )
        response.model_id = response.model_id or f"eb1/{model}"
        return response

    # ------------------------------------------------------------------
    # SWE routing
    # ------------------------------------------------------------------

    @staticmethod
    def _is_swe_model(model: str) -> bool:
        suffix = model.lower()
        if "/" in suffix:
            suffix = suffix.split("/", 1)[1]
        return suffix.startswith("eb1-swe-")

    def _run_swe_routing(
        self,
        model: str,
        request: Any,
        payload: Mapping[str, Any],
        **kwargs: Any,
    ) -> ChatResponse:
        """Route via SWE complexity classification.

        For the pro variant with complex/architectural queries, runs a full
        in-process ensemble (candidates + judge synthesis) instead of
        collapsing to a single model.  All other cases delegate to the
        gateway resolver for single-model invocation.
        """
        from keiro.ensemble.eb1_swe import classify_complexity
        from keiro.gateway.infra.eb1_resolver import _extract_swe_variant

        query = self._extract_query_text(request)
        variant = _extract_swe_variant(model)
        complexity, signals = classify_complexity(query)

        # Pro variant with complex/architectural queries: run in-process
        # ensemble instead of the gateway resolver (which cannot orchestrate
        # multi-model ensembles and would collapse to a single judge model).
        if variant == "pro" and complexity in ("complex", "architectural"):
            return self._run_swe_pro_ensemble(
                model, request, payload, complexity, signals, **kwargs
            )

        # All other cases: single-model routing via gateway resolver.
        from keiro.gateway.infra.eb1_resolver import resolve_eb1_swe_model

        routing = resolve_eb1_swe_model(model, query)

        _LOG.info(
            "SWE routing: %s -> %s (tier=%s, latency=%.1fms)",
            model,
            routing.resolved_model,
            routing.tier,
            routing.routing_latency_ms,
        )

        # Rewrite the payload model to the resolved concrete model and
        # invoke through the registry so the correct provider handles it.
        from ember._internal.context.runtime import EmberContext

        registry = EmberContext.current().model_registry
        resolved_payload = dict(payload)
        resolved_payload["model"] = routing.resolved_model
        response = registry.invoke_responses(routing.resolved_model, resolved_payload)

        # Persist routing decision in raw_output for offline analysis.
        # Trajectory files (.traj) serialize raw_output, so this metadata
        # survives for post-hoc routing analysis without separate log scraping.
        _inject_routing_metadata(response, routing.to_dict())

        return response

    def _run_swe_pro_ensemble(
        self,
        model: str,
        request: Any,
        payload: Mapping[str, Any],
        complexity: str,
        signals: tuple[str, ...],
        **kwargs: Any,
    ) -> ChatResponse:
        """Run in-process ensemble for pro variant complex/architectural queries.

        Builds an Eb1Plan from PRO_ENSEMBLE_CANDIDATES/PRO_ENSEMBLE_JUDGE,
        extracts prompt/system/context from the request, and delegates to
        run_plan().  The plan executor automatically strips tools from
        candidates and gives the full context (with tools) to the judge.
        """
        from keiro.ensemble.eb1 import CandidatePlan, Eb1Plan, JudgePlan, run_plan
        from keiro.ensemble.eb1_swe import (
            MODEL_PARAMS,
            PRO_COMPLEXITY_TEMPERATURE,
            PRO_ENSEMBLE_CANDIDATES,
            PRO_ENSEMBLE_JUDGE,
        )

        # Build the ensemble plan with SWE-specific per-model parameters.
        candidates = tuple(
            CandidatePlan(
                model=candidate_model,
                params=dict(MODEL_PARAMS.get(candidate_model, {})),
            )
            for candidate_model in PRO_ENSEMBLE_CANDIDATES
        )
        judge = JudgePlan(
            model=PRO_ENSEMBLE_JUDGE,
            params=dict(MODEL_PARAMS.get(PRO_ENSEMBLE_JUDGE, {})),
        )
        plan = Eb1Plan(candidates=candidates, judge=judge)

        # Extract prompt, system, and context using the same logic as the
        # non-SWE ensemble path (complete_responses_payload lines 91-106).
        components = build_prompt_components(request)
        prompt_text = (
            self._prompt_from_responses_input(request.input)
            or components.prompt
            or ""
        )

        context_messages: tuple[Mapping[str, str], ...] = tuple(
            {"role": entry["role"], "content": entry["content"]}
            for entry in components.context
        )
        if not context_messages:
            context_messages = self._context_from_responses_messages(
                request.messages,
            )

        temperature = PRO_COMPLEXITY_TEMPERATURE.get(complexity, 0.1)
        shared_options: dict[str, Any] = {"temperature": temperature}

        _LOG.info(
            "SWE pro ensemble: %s -> %d candidates + judge "
            "(complexity=%s, signals=%s)",
            model,
            len(PRO_ENSEMBLE_CANDIDATES),
            complexity,
            signals,
        )

        response = run_plan(
            plan,
            prompt_text,
            system=components.system,
            shared_options=shared_options,
            responses_request=request,
            context_messages=context_messages or None,
        )
        response.model_id = response.model_id or f"eb1/{model}"

        _inject_routing_metadata(response, {
            "original_model": model,
            "resolved_model": "ensemble",
            "tier": complexity,
            "confidence": 1.0,
            "policy": "swe_keyword_ensemble",
            "routing_latency_ms": 0.0,
            "metadata": {
                "routing_mode": "swe_keyword",
                "variant": "pro",
                "complexity": complexity,
                "signals": list(signals),
                "candidates": list(PRO_ENSEMBLE_CANDIDATES),
                "judge": PRO_ENSEMBLE_JUDGE,
            },
        })

        return response

    @staticmethod
    def _extract_query_text(request: Any) -> str:
        """Pull the user query from a ResponsesRequest for SWE classification."""
        if isinstance(request.input, str):
            return request.input

        if isinstance(request.input, Sequence):
            for message in reversed(request.input):
                text = _responses_message_text(message)
                if text:
                    return text

        if request.messages:
            for message in reversed(request.messages):
                text = _responses_message_text(message)
                if text:
                    return text

        return ""

    # ------------------------------------------------------------------
    # Tool-use passthrough
    # ------------------------------------------------------------------

    def _run_tool_use_passthrough(
        self,
        model: str,
        payload: Mapping[str, Any],
    ) -> ChatResponse:
        """Bypass the ensemble for tool-use payloads.

        The EB1 ensemble strips tools from candidates and synthesizes a
        text response, which is incompatible with agentic tool calling
        (SWE-agent, function-calling agents, etc.).  When tools are
        present, route directly to the judge model so the underlying
        provider handles ``tool_use`` natively.
        """
        from keiro.ensemble.model_pool import CLAUDE_OPUS
        from ember._internal.context.runtime import EmberContext

        _LOG.info(
            "Tool-use passthrough: %s -> %s (ensemble incompatible with tools)",
            model,
            CLAUDE_OPUS,
        )

        registry = EmberContext.current().model_registry
        resolved_payload = dict(payload)
        resolved_payload["model"] = CLAUDE_OPUS
        response = registry.invoke_responses(CLAUDE_OPUS, resolved_payload)

        _inject_routing_metadata(response, {
            "original_model": model,
            "resolved_model": CLAUDE_OPUS,
            "policy": "tool_use_passthrough",
            "tier": "passthrough",
            "confidence": 1.0,
            "routing_latency_ms": 0.0,
            "metadata": {"reason": "ensemble_incompatible_with_tools"},
        })

        return response

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

    @staticmethod
    def _extract_invocation_kwargs(params: Mapping[str, Any]) -> dict[str, Any]:
        """Normalize kwargs into the shape expected by ``run_plan()``."""
        working = dict(params)

        system = working.pop("system", None)
        shared_options = _to_optional_mapping(working.pop("shared_options", None))
        request_options = _to_optional_mapping(working.pop("request_options", None))
        responses_request = working.pop("responses_request", None)
        context_messages = working.pop("context_messages", None)
        working.pop("context", None)
        working.pop("stream", None)
        working.pop("metadata", None)

        shared: dict[str, Any] = dict(shared_options or {})
        request_opts: dict[str, Any] = dict(request_options or {})

        for key in list(working):
            value = working.pop(key)
            if key == "max_tokens":
                request_opts.setdefault("max_output_tokens", value)
            elif key in {"max_output_tokens", "response_format"}:
                request_opts.setdefault(key, value)
            elif key == "reasoning" and isinstance(value, Mapping):
                shared.setdefault("reasoning", dict(value))
            else:
                shared.setdefault(key, value)

        normalized_system = (
            system if system is None or isinstance(system, str) else str(system)
        )
        normalized_context = _normalize_context_messages(context_messages)

        return {
            "system": normalized_system,
            "shared_options": shared or None,
            "request_options": request_opts or None,
            "responses_request": responses_request,
            "context_messages": normalized_context if normalized_context else None,
        }

    @staticmethod
    def _prompt_from_responses_input(
        input_value: Sequence[ResponsesMessage] | str | None,
    ) -> str | None:
        if input_value is None:
            return None
        if isinstance(input_value, str):
            candidate = input_value.strip()
            return candidate or None
        for message in reversed(input_value):
            text = _responses_message_text(message)
            if text:
                return text
        return None

    @staticmethod
    def _context_from_responses_messages(
        messages: Sequence[ResponsesMessage] | None,
    ) -> tuple[dict[str, str], ...]:
        if not messages:
            return ()
        entries: list[dict[str, str]] = []
        for message in messages:
            text = _responses_message_text(message)
            if text:
                entries.append({"role": message.role, "content": text})
        return tuple(entries)

    def _get_api_key_from_env(self) -> str:  # pragma: no cover
        return ""


# ------------------------------------------------------------------
# Module-level helpers (shared with static methods above)
# ------------------------------------------------------------------


def _responses_message_text(message: ResponsesMessage) -> str | None:
    """Extract concatenated text from a ResponsesMessage."""
    fragments: list[str] = []
    for part in message.content:
        text = part.fields.get("text")
        if isinstance(text, str):
            fragments.append(text)
    combined = "".join(f.strip() for f in fragments if f.strip())
    return combined or None


def _to_optional_mapping(value: Any) -> dict[str, Any] | None:
    if value is None:
        return None
    if isinstance(value, Mapping):
        return dict(value)
    return None


def _normalize_context_messages(
    value: Any,
) -> tuple[dict[str, str], ...]:
    """Coerce context_messages into a uniform tuple of dicts."""
    if value is None:
        return ()
    if isinstance(value, str):
        text = value.strip()
        if not text:
            return ()
        return ({"role": "system", "content": text},)
    if isinstance(value, Mapping):
        role = str(value.get("role", "user")).strip() or "user"
        content = str(value.get("content", "")).strip()
        if not content:
            return ()
        return ({"role": role, "content": content},)
    if isinstance(value, Sequence):
        normalized: list[dict[str, str]] = []
        for entry in value:
            if not isinstance(entry, Mapping):
                raise ValueError(
                    "context entries must be mappings with role/content"
                )
            role = str(entry.get("role", "user")).strip() or "user"
            content = str(entry.get("content", "")).strip()
            if content:
                normalized.append({"role": role, "content": content})
        return tuple(normalized)
    raise TypeError("context must be string, mapping, or sequence of mappings")


def _inject_routing_metadata(
    response: ChatResponse,
    routing_dict: dict[str, object],
) -> None:
    """Merge routing metadata into ``response.raw_output`` for offline analysis.

    If ``raw_output`` is already a dict-like mapping (as with most providers),
    the routing info is merged under the ``"eb1_routing"`` key.  Otherwise a
    new dict is constructed preserving the original value under ``"_original"``.

    This is intentionally mutation-based: the caller already holds the
    ``ChatResponse`` reference and we avoid allocating a replacement.
    """
    routing_entry = {"eb1_routing": routing_dict}

    existing = response.raw_output
    if existing is None:
        response.raw_output = routing_entry
    elif isinstance(existing, dict):
        existing.update(routing_entry)
    elif isinstance(existing, Mapping):
        merged: dict[str, object] = dict(existing)
        merged.update(routing_entry)
        response.raw_output = merged
    else:
        response.raw_output = {"_original": existing, **routing_entry}
