"""Upstream model pool implementing the :class:`ModelInvoker` protocol.

Bridges the Keiro protocol layer to the existing ``ModelRegistry`` in
``ember.models``.  During migration both call paths coexist: the ensemble
currently uses ``run_plan()`` (which calls ``models`` directly), while
this adapter enables new callers to go through the clean protocol boundary.

Once the ensemble is refactored to accept ``ModelInvoker``, this becomes
the sole entry point for all upstream calls, giving us a single place to
add credential pooling, AIMD rate limiting, and circuit breaking.

Example::

    from keiro.upstream import UpstreamPool
    from keiro.protocol import ModelCall

    pool = UpstreamPool()
    result = pool.invoke(ModelCall(model_id="gpt-5.2", prompt="Hello"))
    print(result.text)
"""

from __future__ import annotations

import logging
import time
from typing import TYPE_CHECKING, Iterator

from ember.models.schemas import UsageStats

from keiro.protocol import (
    ModelCall,
    ModelResult,
    parse_responses_tool_calls,
    responses_output_items,
)

if TYPE_CHECKING:
    from ember.models.runtime.registry import ModelRegistry

logger = logging.getLogger(__name__)


def _get_registry() -> ModelRegistry:
    """Lazily acquire the global ModelRegistry.

    Deferred to avoid import-time side effects from provider discovery.
    The registry is accessed through the public ``ember.api.models`` facade
    which handles initialization.

    Returns:
        The active ModelRegistry instance.
    """
    from ember.api.models import Models

    facade = Models()
    return facade._registry


def _usage_from_response(response: object) -> UsageStats:
    """Extract UsageStats from a ChatResponse.

    Args:
        response: ChatResponse from the ModelRegistry.

    Returns:
        UsageStats extracted from the response, or empty UsageStats.
    """
    usage = getattr(response, "usage", None)
    if isinstance(usage, UsageStats):
        return usage
    return UsageStats()


class UpstreamPool:
    """Adapter wrapping ``ModelRegistry`` behind the ``ModelInvoker`` protocol.

    Translates :class:`ModelCall` into ``registry.invoke_model()`` calls and
    converts the results into :class:`ModelResult`. Satisfies the
    ``ModelInvoker`` protocol via structural typing.

    During migration, this delegates to the existing ``ModelRegistry``.
    Future enhancements will add:
    - Multi-key credential pool per provider.
    - AIMD rate limiting per upstream endpoint.
    - Circuit breakers for degraded providers.
    - Straggler detection and speculative retries.

    Args:
        registry: Explicit ModelRegistry instance. If None, the global
            registry is resolved lazily on first call.
    """

    def __init__(self, registry: ModelRegistry | None = None) -> None:
        self._registry = registry

    def _get_registry(self) -> ModelRegistry:
        """Return the registry, resolving lazily if needed."""
        if self._registry is None:
            self._registry = _get_registry()
        return self._registry

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

        Translates ``ModelCall`` fields into ``registry.invoke_model()``
        kwargs and wraps the response in ``ModelResult``.

        Args:
            call: Model call specification.

        Returns:
            Model response with usage and latency.

        Raises:
            ProviderAPIError: If the provider returns an error.
        """
        registry = self._get_registry()

        kwargs: dict = {}
        if call.system is not None:
            kwargs["system"] = call.system
        if call.context_messages:
            kwargs["context"] = list(call.context_messages)
        if call.tools is not None:
            kwargs["tools"] = list(call.tools)
        if call.response_format is not None:
            kwargs["response_format"] = call.response_format
        if call.reasoning is not None:
            kwargs["reasoning"] = call.reasoning
        if call.max_output_tokens is not None:
            kwargs["max_output_tokens"] = call.max_output_tokens

        start = time.perf_counter()
        response = registry.invoke_model(call.model_id, call.prompt, **kwargs)
        elapsed_ms = (time.perf_counter() - start) * 1000

        text: str = getattr(response, "data", "") or ""
        usage = _usage_from_response(response)
        raw_output = getattr(response, "raw_output", None)

        items = responses_output_items(raw_output)
        tool_calls = parse_responses_tool_calls(items)
        finish_reason = "tool_calls" if tool_calls else "stop"

        return ModelResult(
            text=text,
            model_id=call.model_id,
            usage=usage,
            tool_calls=tool_calls,
            finish_reason=finish_reason,
            latency_ms=elapsed_ms,
        )

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

        Currently falls back to synchronous invocation and yields the
        full result as a single chunk. True per-token streaming will be
        wired when the registry's ``stream_model()`` method is integrated.

        Args:
            call: Model call specification.

        Yields:
            Single ModelResult containing the complete response.
        """
        yield self.invoke(call)

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

        Attempts to resolve the model in the registry. Returns False if
        the model is unknown or provider instantiation fails.

        Args:
            model_id: Canonical model identifier.

        Returns:
            True if the model can likely be invoked successfully.
        """
        try:
            registry = self._get_registry()
            registry.get_model(model_id)
            return True
        except Exception:
            logger.debug("Model %s not available", model_id, exc_info=True)
            return False


