"""Provider adapters used by gateway runtimes."""

from __future__ import annotations

from dataclasses import replace
from typing import Any, Generator, Mapping, Sequence

from ember._internal.exceptions import ProviderAPIError
from ember.core.ir import (
    NormalizationError,
    is_ensemble_variant,
    normalize_tools,
    tools_to_dict_list,
)
from ..errors import CapabilityConflict
from ..telemetry import Eb1Telemetry, get_telemetry
from keiro.protocol import parse_responses_tool_calls
from ember.models.providers.anthropic import AnthropicProvider
from ember.models.providers.azure_openai import AzureOpenAIProvider
from ember.models.providers.base import BaseProvider
from ember.models.providers.google import GoogleProvider
from ember.models.providers.openai import OpenAIProvider
from ember.models.providers.utils import (
    ResponsesRequest,
    sanitize_responses_payload,
)
from ember.models.schemas import ChatResponse, UsageStats
from keiro.ensemble.eb1 import (
    _latency_ms,
    _usage_payload,
)

_adapter_latency_ms = _latency_ms


def _validate_streaming_tools_ir(
    model: str,
    stream: bool,
    tools: Sequence[Mapping[str, Any]] | None,
) -> None:
    """Raise CapabilityConflict if streaming + tools is unsupported for model."""
    if stream and tools and is_ensemble_variant(model):
        raise CapabilityConflict(
            f"Model {model} does not support streaming with tools. "
            "Ensemble variants don't surface worker tool calls. "
            "Use eb1-codex, eb1-a, or set stream=false.",
            label="tools_not_streamable_in_ensemble",
        )


def _normalize_request_tools_ir(
    request: ResponsesRequest,
    model: str,
) -> ResponsesRequest:
    """Return a new ResponsesRequest with normalized tools."""
    if not request.tools:
        return request

    try:
        normalized = normalize_tools(request.tools, model=model)
        tools_list = tools_to_dict_list(normalized)
    except NormalizationError as exc:
        raise ProviderAPIError(str(exc), context={"model": model, **exc.context}) from exc

    return replace(request, tools=tools_list)


_adapter_usage_payload = _usage_payload


def _adapter_begin_run(
    plan_outline: Mapping[str, Any],
    canonical_model: str,
    *,
    mode: str,
    org_limits: Mapping[str, Any] | None,
) -> tuple[Eb1Telemetry, str]:
    telemetry = get_telemetry()
    context = {
        "mode": mode,
        "requested_model": canonical_model,
        "plan": dict(plan_outline),
    }
    if org_limits:
        context["org_limits_keys"] = sorted(org_limits.keys())
    run_id = telemetry.begin_run(context=context)
    return telemetry, run_id


def _adapter_record_error(
    telemetry: Eb1Telemetry,
    run_id: str,
    path: str,
    canonical_model: str,
    error_payload: Mapping[str, Any],
    start_time: float,
) -> None:
    telemetry.record(
        run_id,
        "adapter.error",
        {
            "path": path,
            "requested_model": canonical_model,
            "latency_ms": _adapter_latency_ms(start_time),
            "error": dict(error_payload),
        },
    )


def _adapter_end_run(
    telemetry: Eb1Telemetry,
    run_id: str,
    path: str,
    outcome: Mapping[str, Any],
) -> None:
    payload = dict(outcome)
    payload.setdefault("path", path)
    telemetry.end_run(run_id, outcome=payload)


def _record_telemetry_outcome(
    telemetry: Eb1Telemetry,
    run_id: str,
    path: str,
    start_time: float,
    response: ChatResponse,
) -> None:
    """Build latency + usage outcome dict and finalize the telemetry run."""
    outcome: dict[str, Any] = {"latency_ms": _adapter_latency_ms(start_time)}
    if response.usage is not None:
        outcome["usage"] = _adapter_usage_payload(response.usage)
    _adapter_end_run(telemetry, run_id, path, outcome)


# ---------------------------------------------------------------------------
# Shared adapter base classes
# ---------------------------------------------------------------------------


class ProviderAdapter:
    """Minimal interface implemented by all provider adapters."""

    name: str = "abstract"
    supports_streaming: bool = False

    def complete(
        self,
        model: str,
        prompt_text: str,
        temperature: float = 0.0,
        **kwargs: Any,
    ) -> tuple[str, int, int, object | None]:
        raise NotImplementedError

    def stream_complete(
        self,
        model: str,
        prompt_text: str,
        temperature: float = 0.0,
        **kwargs: Any,
    ) -> Generator[str, None, ChatResponse]:
        raise NotImplementedError


class ProviderClientAdapter(ProviderAdapter):
    """Adapter that wraps a ``BaseProvider`` implementation."""

    def __init__(self, name: str, provider_cls: type[BaseProvider]) -> None:
        self.name = name
        self._provider_cls = provider_cls
        self.supports_streaming = getattr(provider_cls, "supports_streaming", False)
        self._provider_cache: dict[str, BaseProvider] = {}

    # Cache providers by api_key override.
    def _create_provider(self, api_key_override: str | None = None) -> BaseProvider:
        cache_key = api_key_override or "__default__"
        provider = self._provider_cache.get(cache_key)
        if provider is None:
            if api_key_override is None:
                provider = self._provider_cls()
            else:
                provider = self._provider_cls(api_key_override)
            self._provider_cache[cache_key] = provider
        return provider

    def complete(
        self,
        model: str,
        prompt_text: str,
        temperature: float = 0.0,
        **kwargs: Any,
    ) -> tuple[str, int, int, object | None]:
        params = dict(kwargs)
        responses_payload = params.pop("responses_payload", None)
        api_key_override = params.pop("api_key_override", None)
        provider = self._create_provider(api_key_override)

        if responses_payload is not None and getattr(provider, "supports_responses_api", False):
            payload = dict(responses_payload)
            payload.setdefault("model", model)
            payload = sanitize_responses_payload(str(payload.get("model", model)), payload)
            response = provider.complete_responses_payload(payload)
            return _response_to_tuple(response)

        response = provider.complete(
            prompt_text,
            model,
            temperature=temperature,
            **params,
        )
        return _response_to_tuple(response)

    def stream_complete(
        self,
        model: str,
        prompt_text: str,
        temperature: float = 0.0,
        **kwargs: Any,
    ) -> Generator[str, None, ChatResponse]:
        params = dict(kwargs)
        responses_payload = params.pop("responses_payload", None)
        api_key_override = params.pop("api_key_override", None)
        provider = self._create_provider(api_key_override)

        if responses_payload is not None and getattr(provider, "supports_responses_api", False):
            payload = dict(responses_payload)
            payload.setdefault("model", model)
            payload = sanitize_responses_payload(str(payload.get("model", model)), payload)
            return provider.stream_responses_payload(payload)

        return provider.stream_complete(
            prompt_text,
            model,
            temperature=temperature,
            **params,
        )


# ---------------------------------------------------------------------------
# Concrete adapters
# ---------------------------------------------------------------------------


class OpenAIAdapter(ProviderClientAdapter):
    """Adapter that routes calls to OpenAI."""

    def __init__(self, provider_cls: type[BaseProvider] | None = None) -> None:
        super().__init__("openai", provider_cls or OpenAIProvider)
        self.supports_streaming = True


class AnthropicAdapter(ProviderClientAdapter):
    """Adapter wrapping the Anthropic provider."""

    def __init__(self) -> None:
        super().__init__("anthropic", AnthropicProvider)
        self.supports_streaming = False  # Anthropic streaming handled elsewhere.


class GoogleAdapter(ProviderClientAdapter):
    """Adapter wrapping Google Gemini provider."""

    def __init__(self) -> None:
        super().__init__("google", GoogleProvider)
        self.supports_streaming = True


class AzureOpenAIAdapter(ProviderClientAdapter):
    """Adapter wrapping Azure OpenAI."""

    def __init__(self) -> None:
        super().__init__("azure_openai", AzureOpenAIProvider)
        self.supports_streaming = True


# ---------------------------------------------------------------------------
# Takumi / EB1 adapter — canonical source: takumi.py
# ---------------------------------------------------------------------------


def _response_to_tuple(response: ChatResponse) -> tuple[str, int, int, dict | None]:
    usage = response.usage or UsageStats()
    raw_metadata = response.raw_output

    if isinstance(raw_metadata, Mapping):
        responses_payload = raw_metadata.get("responses_payload")
        if isinstance(responses_payload, Mapping):
            metadata: dict | None = dict(responses_payload)
        else:
            metadata = dict(raw_metadata)
    elif raw_metadata is None:
        metadata = None
    elif hasattr(raw_metadata, "model_dump"):
        # Pydantic model - convert to dict
        metadata = raw_metadata.model_dump()
    elif hasattr(raw_metadata, "to_dict"):
        # OpenAI SDK objects have to_dict() method
        metadata = raw_metadata.to_dict()
    else:
        # Fallback: try to extract useful info or skip
        metadata = {"_type": type(raw_metadata).__name__}

    # Include tool_calls in metadata if present (for function calling support)
    tool_calls = response.tool_calls
    if not tool_calls and response.output:
        parsed = parse_responses_tool_calls(response.output)
        tool_calls = [tc.to_chat_dict() for tc in parsed]

    if tool_calls:
        if metadata is None:
            metadata = {}
        metadata["_tool_calls"] = tool_calls

    return response.data, usage.prompt_tokens, usage.completion_tokens, metadata


# Re-exports: symbols that takumi.py's ``_pa()`` resolves at runtime from
# this module, plus the TakumiAdapter class itself.
from .takumi import (  # noqa: E402,F401
    TakumiAdapter,
    _EB1_COMPLETION_SOFT_CAP,
    _apply_completion_limit,
    _codex_orchestrator,
    _filter_plan_for_local,
    _is_codex_plan,
    _provider_has_credentials,
    _run_local_eb1,
    _run_local_eb1_stream,
    _select_default_plan,
    _takumi_base_url,
)
from keiro.ensemble.eb1 import run_plan  # noqa: E402,F401
from keiro.ensemble.eb1_a import run_eb1_a  # noqa: E402,F401
from ..providers.planner import compile_plan  # noqa: E402,F401


# ---------------------------------------------------------------------------
# Public exports
# ---------------------------------------------------------------------------

__all__ = [
    "ProviderAdapter",
    "ProviderClientAdapter",
    "OpenAIAdapter",
    "AnthropicAdapter",
    "GoogleAdapter",
    "AzureOpenAIAdapter",
    "TakumiAdapter",
]
