"""Takumi gateway provider.

Forwards requests to the Takumi Gateway API which provides:
- Unified access to all major LLM providers (OpenAI, Anthropic, Google, Azure)
- EB1 and EB1-Pro ensemble models
- Rate limiting, billing, and usage tracking
- Session affinity and intelligent routing

Examples:
    >>> from keiro.providers.takumi import TakumiProvider
    >>> provider = TakumiProvider(api_key="tkm_sk_...")
    >>> response = provider.complete("Hello", "gpt-4")
    >>> response.text
    'Hello! How can I help you today?'
"""

from __future__ import annotations

import json
import logging
import uuid
from contextlib import closing
from typing import Any, Dict, Generator, Mapping, MutableMapping, Optional, Sequence, Tuple
from urllib.parse import urlparse

import requests
from tenacity import (
    retry,
    retry_if_exception,
    stop_after_attempt,
    wait_exponential,
)

from ember._internal.exceptions import ProviderAPIError
from ember.models.providers.base import BaseProvider
from ember.models.providers.openai import usage_from_openai
from ember.models.providers.utils import (
    ResponsesMessage,
    build_prompt_components,
    extract_first_choice_text,
    parse_responses_request,
    parse_usage_payload,
)
from ember.models.schemas import ChatResponse, UsageStats
from keiro.ensemble.config import get_takumi_settings
from keiro.ensemble.eb1_a import run_eb1_a

logger = logging.getLogger(__name__)

_EB1_VARIANT_TIMEOUT_STAGES = {
    "eb1": 3,
    "eb1-a": 3,
    "eb1-pro": 4,
}

_EB1_VARIANT_PREFIXES: tuple[tuple[str, str], ...] = (
    ("eb1-a", "eb1-a"),
    ("eb1-pro", "eb1-pro"),
    ("eb1", "eb1"),
)

_EB1_PER_STAGE_TIMEOUT_SECONDS = {
    "remote": 225.0,  # Takumi Cloud
    "local": 300.0,  # local backend proxied by gateway
}

# Models that require EB1-GNN server routing (not handled by local gateway)
_EB1_GNN_ROUTED_PREFIXES = frozenset(
    (
    "eb1-preview",
    "eb1-fast-preview",
    "eb1-frontier-preview",
    "eb1-delta-preview",
    )
)


def _requires_eb1_gnn_routing(model: str) -> bool:
    """Check if model requires routing through the EB1-GNN server."""
    model_lower = model.lower()
    if model_lower.startswith("takumi/"):
        model_lower = model_lower[7:]
    return any(model_lower.startswith(prefix) for prefix in _EB1_GNN_ROUTED_PREFIXES)


_TRANSIENT_ERROR_TYPES = frozenset({"server", "rate_limit", "network"})


def _is_transient_error(exc: BaseException) -> bool:
    """Return True for ProviderAPIErrors eligible for retry."""
    return (
        isinstance(exc, ProviderAPIError)
        and exc.context.get("error_type") in _TRANSIENT_ERROR_TYPES
    )


def _log_transient_retry(retry_state) -> None:
    """Log retry attempts with structured context."""
    exc = retry_state.outcome.exception()
    ctx = exc.context if isinstance(exc, ProviderAPIError) else {}
    logger.warning(
        "Transient %s error for model=%s, retrying in %.1fs (attempt %d): %s",
        ctx.get("error_type", "unknown"),
        ctx.get("model", "unknown"),
        retry_state.next_action.sleep,
        retry_state.attempt_number,
        exc,
    )


class TakumiProvider(BaseProvider):
    """Takumi Gateway provider for unified LLM access."""

    requires_api_key = True
    supports_responses_api = True

    def __init__(
        self,
        api_key: str | None = None,
        base_url: str | None = None,
    ):
        """Initialize the Takumi provider.

        Args:
            api_key: Gateway API key (tkm_sk_*) override.
            base_url: Gateway base URL. Defaults to the value in
                ``~/.ember/config.yaml`` (override with ``EMBER_CONFIG_PATH``).
        """
        super().__init__(api_key)

        self.base_url = self._resolve_base_url(base_url)

        # Load EB1-GNN server config for models requiring GNN routing
        self._eb1_gnn_url, self._eb1_gnn_api_key = self._load_eb1_gnn_config()

        # Create a session for connection pooling and keep-alive
        self._session = requests.Session()
        # Configure connection pooling parameters
        adapter = requests.adapters.HTTPAdapter(
            pool_connections=10,  # Number of connection pools to cache
            pool_maxsize=10,  # Maximum number of connections to save in the pool
            max_retries=0,  # Retries handled by tenacity on _execute_request
            pool_block=False,  # Don't block when pool is full
        )
        self._session.mount("http://", adapter)
        self._session.mount("https://", adapter)

    def _prepare_request(
        self,
        prompt: str,
        model: str,
        params: Dict[str, Any],
    ) -> tuple[Dict[str, Any], bool, Optional[float], str]:
        """Build request payload and extract shared arguments."""

        working = dict(params)
        messages = self._build_messages(prompt, working)
        payload: Dict[str, Any] = {"model": model, "messages": messages}

        stream_requested = bool(working.pop("stream", False))

        session_id = working.pop("session_id", None)
        if session_id is not None:
            payload["session_id"] = session_id

        policy = working.pop("policy", None)
        if policy is not None:
            if not isinstance(policy, dict):
                raise ValueError("policy must be a mapping when provided")
            payload["policy"] = policy

        passthrough_keys = (
            "temperature",
            "max_tokens",
            "top_p",
            "response_format",
            "tool_choice",
            "tools",
            "parallel_tool_calls",
            "metadata",
            "modalities",
            "seed",
            "user",
            "reasoning",
            "text",
            "max_output_tokens",
        )
        for key in passthrough_keys:
            if key in working:
                value = working.pop(key)
                if value is not None:
                    payload[key] = value

        if _is_gpt5_model(model) and "max_tokens" in payload and "max_output_tokens" not in payload:
            value = payload.pop("max_tokens")
            try:
                payload["max_output_tokens"] = int(value)
            except (TypeError, ValueError) as exc:
                raise ValueError("max_tokens must be an integer for GPT-5 models") from exc

        timeout_override = working.pop("timeout", None)
        if timeout_override is not None:
            try:
                timeout_override = float(timeout_override)
            except (TypeError, ValueError) as exc:
                raise ValueError("timeout must be numeric") from exc

        idempotency_key = working.pop("idempotency_key", None)
        if idempotency_key is None:
            idempotency_key = str(uuid.uuid4())

        if working:
            unsupported = ", ".join(sorted(working))
            raise ValueError(f"Unsupported TakumiProvider parameters: {unsupported}")

        return payload, stream_requested, timeout_override, idempotency_key

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=2, max=30),
        retry=retry_if_exception(_is_transient_error),
        reraise=True,
        before_sleep=_log_transient_retry,
    )
    def _execute_request(
        self,
        *,
        url: str,
        payload: Mapping[str, Any],
        headers: Mapping[str, str],
        timeout_plan: list[tuple[float, float]],
        model: str,
        stream: bool = False,
    ) -> requests.Response:
        """Execute HTTP POST with timeout-plan fallback and transient retries."""
        if not timeout_plan:
            raise ProviderAPIError(
                "Gateway timeout plan is empty",
                context={"model": model, "error_type": "invalid_request", "url": url},
            )

        resp: requests.Response | None = None
        attempts = len(timeout_plan)
        for attempt, (connect_timeout, read_timeout) in enumerate(timeout_plan, start=1):
            try:
                resp = self._session.post(
                    url,
                    json=payload,
                    headers=headers,
                    timeout=(connect_timeout, read_timeout),
                    stream=stream,
                )
                break
            except requests.Timeout as exc:
                if attempt < attempts:
                    continue
                raise ProviderAPIError(
                    f"Gateway read timed out after {read_timeout:.0f}s.",
                    context={
                        "model": model,
                        "error_type": "timeout",
                        "url": url,
                        "attempts": attempts,
                        "read_timeout_seconds": read_timeout,
                    },
                ) from exc
            except requests.RequestException as exc:
                raise ProviderAPIError(
                    f"Gateway request failed: {exc}",
                    context={
                        "model": model,
                        "error_type": "network",
                        "url": url,
                    },
                ) from exc

        if resp is None:
            raise ProviderAPIError(
                "Gateway request could not be completed (no response)",
                context={"model": model, "error_type": "unknown", "url": url},
            )

        if resp.status_code < 400:
            return resp

        context: Dict[str, Any] = {
            "model": model,
            "status": resp.status_code,
            "url": url,
        }
        if resp.status_code in (401, 403):
            context["error_type"] = "authentication"
            raise ProviderAPIError("Authentication failed", context=context)
        if resp.status_code == 429:
            context["error_type"] = "rate_limit"
            raise ProviderAPIError("Rate limit exceeded", context=context)
        if resp.status_code >= 500:
            context["error_type"] = "server"
            raise ProviderAPIError("Gateway error", context=context)

        context["error_type"] = "invalid_request"
        body = resp.text[:500] if not stream else "<streaming>"
        raise ProviderAPIError(f"Request failed: {body}", context=context)

    def complete(self, prompt: str, model: str, **kwargs: Any) -> ChatResponse:
        """Generate completion via Takumi Gateway.

        Args:
            prompt: The user prompt
            model: Model identifier (e.g., 'gpt-4', 'claude-opus-4-1', 'eb1')
            **kwargs: Optional parameters:
                system: System message
                temperature: Sampling temperature (0-2)
                max_tokens: Maximum completion tokens
                stream: Enable streaming (bool, default False)
                session_id: Session affinity key
                policy: Request policy dict (allowed_providers, max_cost_usd, etc.)

        Returns:
            ChatResponse with text, usage stats, and metadata
        """
        payload, stream_requested, timeout_override, idempotency_key = self._prepare_request(
            prompt,
            model,
            kwargs,
        )

        if stream_requested:
            raise ProviderAPIError(
                "TakumiProvider.complete does not support stream=True; use stream_complete()",
                context={"model": model, "error_type": "invalid_parameter"},
            )

        timeout_plan = self._build_timeout_plan(
            prompt,
            payload,
            model=model,
            timeout_override=timeout_override,
        )

        # Get effective URL and API key based on model (EB1 preview models use GNN server)
        effective_url, effective_key = self._get_effective_url_and_key(model)
        logger.info("Model %s routing to %s", model, effective_url)

        headers = {
            "Authorization": f"Bearer {effective_key}",
            "Content-Type": "application/json",
            "Idempotency-Key": idempotency_key,
        }

        url = self._build_endpoint_url("/v1/chat/completions", base_url=effective_url)

        resp = self._execute_request(
            url=url,
            payload=payload,
            headers=headers,
            timeout_plan=timeout_plan,
            model=model,
        )

        try:
            data = resp.json()
        except ValueError as exc:
            raise ProviderAPIError(
                "Non-JSON response from gateway",
                context={"model": model, "error_type": "invalid_response"},
            ) from exc

        text_output = extract_first_choice_text(data, model=model)
        usage_stats = parse_usage_payload(data.get("usage"), model=model)

        return ChatResponse(
            data=text_output,
            usage=usage_stats,
            model_id=data.get("model") or model,
            raw_output=data,
        )

    def complete_responses_payload(
        self,
        payload: Mapping[str, Any],
        **kwargs: Any,
    ) -> ChatResponse:
        """Invoke Takumi Gateway's /v1/responses endpoint."""

        if not isinstance(payload, Mapping):
            raise ValueError("Responses payload must be a mapping")

        request_payload: Dict[str, Any] = dict(payload)
        request_payload.pop("stream", None)
        request_payload.pop("timeout", None)

        timeout_override = kwargs.pop("timeout", None)
        idempotency_key = kwargs.pop("idempotency_key", None) or request_payload.pop(
            "idempotency_key", None
        )
        if kwargs:
            unsupported = ", ".join(sorted(kwargs))
            raise ValueError(f"Unsupported Takumi Responses overrides: {unsupported}")

        model = str(request_payload.get("model") or "").strip()
        if not model:
            raise ValueError("Responses payload must include 'model'")

        timeout_payload = dict(request_payload)
        if "max_output_tokens" in request_payload and "max_tokens" not in timeout_payload:
            timeout_payload["max_tokens"] = request_payload["max_output_tokens"]

        prompt_summary = self._summarize_responses_prompt(request_payload)
        timeout_plan = self._build_timeout_plan(
            prompt_summary,
            timeout_payload,
            model=model,
            timeout_override=timeout_override,
        )

        if not idempotency_key:
            idempotency_key = str(uuid.uuid4())

        effective_url, effective_key = self._get_effective_url_and_key(model)

        headers = {
            "Authorization": f"Bearer {effective_key}",
            "Content-Type": "application/json",
            "Idempotency-Key": idempotency_key,
        }

        url = self._build_endpoint_url("/v1/responses", base_url=effective_url)

        resp = self._execute_request(
            url=url,
            payload=request_payload,
            headers=headers,
            timeout_plan=timeout_plan,
            model=model,
        )

        try:
            data = resp.json()
        except ValueError as exc:
            raise ProviderAPIError(
                "Non-JSON response from gateway",
                context={"model": model, "error_type": "invalid_response"},
            ) from exc

        usage_stats = usage_from_openai(data.get("usage"))
        output_items = self._normalize_responses_output(data.get("output"))

        text_output = ""
        if isinstance(data.get("output_text"), str):
            text_output = str(data["output_text"])
        elif output_items:
            text_output = self._aggregate_responses_text(output_items)

        response_id = data.get("id")
        resolved_model = data.get("model") or model

        return ChatResponse(
            data=text_output,
            output=output_items,
            response_id=str(response_id) if isinstance(response_id, str) else None,
            usage=usage_stats,
            model_id=str(resolved_model),
            raw_output=data,
        )

    def stream_complete(
        self,
        prompt: str,
        model: str,
        **kwargs: Any,
    ) -> Generator[str, None, ChatResponse]:
        """Stream completion results via Takumi Gateway SSE."""

        payload, stream_requested, timeout_override, idempotency_key = self._prepare_request(
            prompt,
            model,
            kwargs,
        )

        if stream_requested:
            raise ProviderAPIError(
                "Do not pass stream=True to stream_complete; streaming is implied",
                context={"model": model, "error_type": "invalid_parameter"},
            )

        timeout_plan = self._build_timeout_plan(
            prompt,
            payload,
            model=model,
            timeout_override=timeout_override,
        )

        payload["stream"] = True
        effective_url, effective_key = self._get_effective_url_and_key(model)

        headers = {
            "Authorization": f"Bearer {effective_key}",
            "Content-Type": "application/json",
            "Accept": "text/event-stream",
            "Idempotency-Key": idempotency_key,
        }
        url = self._build_endpoint_url("/v1/chat/completions", base_url=effective_url)

        response = self._execute_request(
            url=url,
            payload=payload,
            headers=headers,
            timeout_plan=timeout_plan,
            model=model,
            stream=True,
        )

        def _generator() -> Generator[str, None, ChatResponse]:
            header_payload: Dict[str, Any] | None = None
            text_chunks: list[str] = []
            tool_call_deltas: list[Dict[str, Any]] = []
            usage_totals = {"prompt": 0, "completion": 0, "total": 0}
            raw_events: list[Dict[str, Any]] = []

            with closing(response):
                for event_type, data_str, event_id in self._iter_sse_events(response):
                    if event_type == "heartbeat":
                        raw_events.append({"event": event_type, "id": event_id, "data": {}})
                        continue
                    if event_type == "done":
                        raw_events.append({"event": event_type, "id": event_id, "data": None})
                        break

                    parsed_data: Any = None
                    if data_str:
                        parsed_data = self._decode_sse_json(
                            data_str,
                            event_type=event_type,
                            warning_prefix="Failed to parse SSE event data as JSON",
                            fallback=data_str,
                        )

                    raw_events.append({"event": event_type, "id": event_id, "data": parsed_data})

                    if event_type == "header":
                        if isinstance(parsed_data, dict):
                            header_payload = parsed_data
                        continue

                    if not data_str:
                        continue

                    if isinstance(parsed_data, dict):
                        delta = parsed_data.get("delta")
                        if delta is None:
                            text_delta = parsed_data.get("text")
                            if isinstance(text_delta, str):
                                delta = text_delta
                            elif "arguments_delta" in parsed_data:
                                tool_delta: Dict[str, Any] = {
                                    "type": parsed_data.get("type") or "tool_call.delta",
                                    "arguments_delta": parsed_data["arguments_delta"],
                                }
                                if parsed_data.get("id"):
                                    tool_delta["id"] = parsed_data["id"]
                                if parsed_data.get("name"):
                                    tool_delta["name"] = parsed_data["name"]
                                delta = tool_delta
                        if isinstance(delta, str):
                            text_chunks.append(delta)
                            yield delta
                        elif isinstance(delta, dict):
                            tool_call_deltas.append(delta)
                        usage_delta = parsed_data.get("usage_delta")
                        if isinstance(usage_delta, dict):
                            self._append_chat_usage_delta(usage_totals, usage_delta)
                    elif isinstance(parsed_data, str):
                        text_chunks.append(parsed_data)
                        yield parsed_data
                    else:
                        text_chunks.append(data_str)
                        yield data_str

            final_text = "".join(text_chunks)
            usage_stats = None
            if header_payload is not None:
                usage_stats = parse_usage_payload(header_payload.get("usage"), model=model)
            if usage_stats is None and any(usage_totals.values()):
                usage_stats = UsageStats(
                    prompt_tokens=usage_totals["prompt"],
                    completion_tokens=usage_totals["completion"],
                    total_tokens=usage_totals["total"],
                    actual_cost_usd=None,
                )

            header_model = None
            if header_payload and isinstance(header_payload, dict):
                candidate = header_payload.get("model")
                if isinstance(candidate, str) and candidate:
                    header_model = candidate

            raw_output: Dict[str, Any] = {"events": raw_events}
            if header_payload is not None:
                raw_output["header"] = header_payload
            if tool_call_deltas:
                raw_output["tool_calls"] = tool_call_deltas

            return ChatResponse(
                data=final_text,
                usage=usage_stats,
                model_id=header_model or model,
                raw_output=raw_output,
                tool_calls=tool_call_deltas if tool_call_deltas else None,
            )

        return _generator()

    def stream_responses_payload(
        self,
        payload: Mapping[str, Any],
        **kwargs: Any,
    ) -> Generator[str, None, ChatResponse]:
        """Stream Responses API events from Takumi Gateway."""

        if not isinstance(payload, Mapping):
            raise ValueError("Responses payload must be a mapping")

        request_payload: Dict[str, Any] = dict(payload)
        request_payload.pop("timeout", None)
        request_payload["stream"] = True

        timeout_override = kwargs.pop("timeout", None)
        idempotency_key = kwargs.pop("idempotency_key", None) or request_payload.pop(
            "idempotency_key", None
        )

        if kwargs:
            unsupported = ", ".join(sorted(kwargs))
            raise ValueError(f"Unsupported Takumi Responses overrides: {unsupported}")

        model = str(request_payload.get("model") or "").strip()
        if not model:
            raise ValueError("Responses payload must include 'model'")

        timeout_payload = dict(request_payload)
        if "max_output_tokens" in request_payload and "max_tokens" not in timeout_payload:
            timeout_payload["max_tokens"] = request_payload["max_output_tokens"]

        prompt_summary = self._summarize_responses_prompt(request_payload)
        timeout_plan = self._build_timeout_plan(
            prompt_summary,
            timeout_payload,
            model=model,
            timeout_override=timeout_override,
        )

        if not idempotency_key:
            idempotency_key = str(uuid.uuid4())

        effective_url, effective_key = self._get_effective_url_and_key(model)

        headers = {
            "Authorization": f"Bearer {effective_key}",
            "Content-Type": "application/json",
            "Accept": "text/event-stream",
            "Idempotency-Key": idempotency_key,
        }

        url = self._build_endpoint_url("/v1/responses", base_url=effective_url)

        response = self._execute_request(
            url=url,
            payload=request_payload,
            headers=headers,
            timeout_plan=timeout_plan,
            model=model,
            stream=True,
        )

        def _generator() -> Generator[str, None, ChatResponse]:
            parts: list[str] = []
            usage_totals = UsageStats()
            raw_events: list[Dict[str, Any]] = []
            final_payload: Dict[str, Any] | None = None

            try:
                for event_type, data_str, event_id in self._iter_sse_events(response):
                    if data_str is None:
                        continue

                    raw_event: Dict[str, Any] = {"data": data_str}
                    if event_type is not None:
                        raw_event["type"] = event_type
                    if event_id is not None:
                        raw_event["id"] = event_id
                    raw_events.append(raw_event)

                    text_fragment: Optional[str] = None
                    parsed_payload: Dict[str, Any] | None = None
                    parsed_candidate = self._decode_sse_json(
                        data_str,
                        event_type=event_type,
                        warning_prefix=(
                            "Failed to parse streaming event as JSON, "
                            "structured handling skipped"
                        ),
                        fallback=None,
                    )
                    if isinstance(parsed_candidate, dict):
                        parsed_payload = parsed_candidate

                    effective_type = event_type
                    if parsed_payload is not None:
                        effective_type = str(parsed_payload.get("type") or event_type or "")
                        if effective_type == "response.error":
                            error_obj = parsed_payload.get("error")
                            if isinstance(error_obj, dict):
                                message = str(error_obj.get("message") or "stream error")
                            else:
                                message = str(error_obj or "stream error")
                            raise ProviderAPIError(
                                f"Takumi streaming error: {message}",
                                context={"model": model, "error_type": "stream_error"},
                            )
                        if effective_type == "response.completed":
                            response_block = parsed_payload.get("response")
                            if isinstance(response_block, dict):
                                final_payload = response_block
                            else:
                                final_payload = parsed_payload

                        delta_value = parsed_payload.get("delta")
                        if isinstance(delta_value, str):
                            text_fragment = delta_value
                        usage_delta = parsed_payload.get("usage_delta")
                        if not isinstance(usage_delta, dict):
                            data_section = parsed_payload.get("data")
                            if isinstance(data_section, dict):
                                usage_delta = data_section.get("usage_delta")
                        if isinstance(usage_delta, dict):
                            self._append_responses_usage_delta(usage_totals, usage_delta)

                    if text_fragment:
                        parts.append(text_fragment)

                    yield data_str
            finally:
                response.close()

            final_usage = None
            final_output = None
            response_id: Optional[str] = None
            resolved_model = model

            if final_payload is not None:
                final_usage = usage_from_openai(final_payload.get("usage"))
                final_output = self._normalize_responses_output(final_payload.get("output"))
                candidate_id = final_payload.get("id")
                if isinstance(candidate_id, str):
                    response_id = candidate_id
                candidate_model = final_payload.get("model")
                if isinstance(candidate_model, str) and candidate_model:
                    resolved_model = candidate_model

            if final_usage is None and (
                usage_totals.prompt_tokens
                or usage_totals.completion_tokens
                or usage_totals.total_tokens
            ):
                final_usage = usage_totals

            final_text = "".join(parts)
            if not final_text and final_output:
                final_text = self._aggregate_responses_text(final_output)

            raw_output: Dict[str, Any] = {"events": raw_events}
            if final_payload is not None:
                raw_output["response"] = final_payload

            return ChatResponse(
                data=final_text,
                output=final_output,
                response_id=response_id,
                usage=final_usage,
                model_id=str(resolved_model),
                raw_output=raw_output,
            )

        return _generator()

    def _summarize_responses_prompt(self, payload: Mapping[str, Any]) -> str:
        parts: list[str] = []
        instructions = payload.get("instructions")
        if isinstance(instructions, str):
            parts.append(instructions)
        parts.extend(_collect_input_texts(payload.get("input")))
        return " ".join(parts)[:2000]

    @staticmethod
    def _normalize_responses_output(value: Any) -> Optional[list[Dict[str, Any]]]:
        if isinstance(value, list):
            return value
        return None

    @staticmethod
    def _aggregate_responses_text(output: Optional[list[Dict[str, Any]]]) -> str:
        if not output:
            return ""
        parts: list[str] = []
        for item in output:
            if not isinstance(item, dict):
                continue
            content = item.get("content")
            if isinstance(content, list):
                for part in content:
                    if isinstance(part, dict):
                        text_val = part.get("text")
                        if isinstance(text_val, str):
                            parts.append(text_val)
                    elif isinstance(part, str):
                        parts.append(part)
            elif isinstance(content, str):
                parts.append(content)
            text_val = item.get("text")
            if isinstance(text_val, str):
                parts.append(text_val)
        return "".join(parts)

    @staticmethod
    def _safe_int(value: Any) -> int:
        if value is None:
            return 0
        try:
            return int(value)
        except (TypeError, ValueError):
            logger.warning(
                "Non-numeric value in usage/token field, coercing to 0: %r",
                value,
            )
            return 0

    @staticmethod
    def _decode_sse_json(
        data_str: str,
        *,
        event_type: str | None,
        warning_prefix: str,
        fallback: Any,
    ) -> Any:
        try:
            return json.loads(data_str)
        except ValueError:
            preview = data_str[:100] + "..." if len(data_str) > 100 else data_str
            logger.warning(
                "%s (event_type=%s): %s",
                warning_prefix,
                event_type,
                preview,
            )
            return fallback

    @classmethod
    def _append_chat_usage_delta(
        cls,
        usage_totals: dict[str, int],
        usage_delta: Mapping[str, Any],
    ) -> None:
        usage_totals["prompt"] += cls._safe_int(usage_delta.get("prompt_tokens"))
        usage_totals["completion"] += cls._safe_int(usage_delta.get("completion_tokens"))
        usage_totals["total"] += cls._safe_int(usage_delta.get("total_tokens"))

    @classmethod
    def _append_responses_usage_delta(
        cls,
        usage_totals: UsageStats,
        usage_delta: Mapping[str, Any],
    ) -> None:
        usage_totals.prompt_tokens += cls._safe_int(
            usage_delta.get("prompt_tokens") or usage_delta.get("input_tokens")
        )
        usage_totals.completion_tokens += cls._safe_int(
            usage_delta.get("completion_tokens") or usage_delta.get("output_tokens")
        )
        total_hint = usage_delta.get("total_tokens")
        if total_hint is not None:
            usage_totals.total_tokens += cls._safe_int(total_hint)
        else:
            usage_totals.total_tokens = (
                usage_totals.prompt_tokens + usage_totals.completion_tokens
            )

    @staticmethod
    def _iter_sse_events(
        response: requests.Response,
    ) -> Generator[tuple[Optional[str], str, Optional[Any]], None, None]:
        """Yield (event_type, data, event_id) tuples from an SSE response."""

        event_type: Optional[str] = None
        data_parts: list[str] = []
        event_id: Optional[Any] = None

        for raw_line in response.iter_lines(decode_unicode=True):
            if raw_line is None:
                continue
            line = raw_line.rstrip("\r\n")
            if not line:
                if event_type is not None or data_parts:
                    yield event_type, "\n".join(data_parts), event_id
                event_type = None
                data_parts = []
                event_id = None
                continue
            if line.startswith(":"):
                continue
            if line.startswith("event:"):
                event_type = line[len("event:") :].strip() or None
                continue
            if line.startswith("data:"):
                data_parts.append(line[len("data:") :].lstrip())
                continue
            if line.startswith("id:"):
                candidate = line[len("id:") :].strip()
                try:
                    event_id = int(candidate)
                except ValueError:
                    event_id = candidate or None
                continue
        if event_type is not None or data_parts:
            yield event_type, "\n".join(data_parts), event_id

    def _get_api_key_from_env(self) -> str:
        """Resolve API key from the active Ember configuration."""

        from ember._internal.context.runtime import EmberContext

        return EmberContext.current().get_credential("takumi")

    def _resolve_base_url(self, explicit_base_url: Optional[str]) -> str:
        """Resolve base URL: explicit > EMBER_TAKUMI_BASE_URL env > config."""
        import os

        candidate = (
            explicit_base_url
            or os.environ.get("EMBER_TAKUMI_BASE_URL")
            or get_takumi_settings().base_url
        )
        if not isinstance(candidate, str) or not candidate.strip():
            raise ValueError(
                "Takumi base_url is not configured. Provide base_url=... or "
                "set EMBER_TAKUMI_BASE_URL or "
                "set base_url in ~/.ember/config.yaml (override with EMBER_CONFIG_PATH)."
            )
        return candidate.rstrip("/")

    def _load_eb1_gnn_config(self) -> tuple[str | None, str | None]:
        """Load EB1-GNN server config via EmberContext.

        Reads ``eb1_gnn.base_url`` and ``eb1_gnn.api_key`` from the centralized
        config system (``~/.ember/config.yaml`` through ``EmberContext``).

        Returns:
            Tuple of (base_url, api_key) for the EB1-GNN server. Either may be
            None if not configured, in which case EB1 preview models will fall
            back to the default gateway (which may not support GNN routing).
        """
        try:
            from ember._internal.context.runtime import EmberContext

            ctx = EmberContext.current()
        except Exception:
            logger.warning(
                "Failed to load EmberContext for EB1-GNN config — "
                "GNN routing will be disabled for EB1-preview models",
                exc_info=True,
            )
            return None, None

        base_url = ctx.get_config("eb1_gnn.base_url")
        if isinstance(base_url, str) and base_url.strip():
            base_url = base_url.strip().rstrip("/")
        else:
            base_url = None

        api_key = ctx.get_config("eb1_gnn.api_key")
        if isinstance(api_key, str) and api_key.strip():
            api_key = api_key.strip()
        else:
            api_key = None

        if base_url:
            logger.debug("Loaded eb1_gnn config: base_url=%s", base_url)

        return base_url, api_key

    def _get_effective_url_and_key(self, model: str) -> tuple[str, str]:
        """Get the effective base URL and API key for the given model.

        For EB1 preview models that require GNN routing, uses eb1_gnn.base_url
        if configured. Falls back to default base_url otherwise.

        Args:
            model: Model identifier.

        Returns:
            Tuple of (base_url, api_key) to use for the request.
        """
        if _requires_eb1_gnn_routing(model) and self._eb1_gnn_url:
            effective_url = self._eb1_gnn_url
            effective_key = self._eb1_gnn_api_key or self.api_key
            logger.debug("Using EB1-GNN server for model %s: %s", model, effective_url)
        else:
            effective_url = self.base_url
            effective_key = self.api_key

        return effective_url, effective_key

    def _build_endpoint_url(self, path: str, base_url: str | None = None) -> str:
        """Build full endpoint URL, avoiding double /v1 if already in base_url.

        Args:
            path: API path like "/v1/responses" or "/v1/chat/completions"
            base_url: Optional override for the base URL. Uses self.base_url if None.

        Returns:
            Full URL combining base_url and path correctly
        """
        base = (base_url or self.base_url).rstrip("/")
        # If path starts with /v1 and base already ends with /v1, strip /v1 from path
        if path.startswith("/v1/") and base.endswith("/v1"):
            path = path[3:]  # Remove leading "/v1" from path
        return f"{base}{path}"

    def _build_messages(
        self,
        prompt: str,
        kwargs: MutableMapping[str, Any],
    ) -> list[dict[str, str]]:
        system_prompt = kwargs.pop("system", None)
        context_prompt = kwargs.pop("context", None)

        messages: list[dict[str, str]] = []

        if system_prompt is not None:
            if not isinstance(system_prompt, str):
                raise ValueError("system prompt must be a string when provided")
            messages.append({"role": "system", "content": system_prompt})

        if context_prompt is not None:
            messages.extend(_context_to_messages(context_prompt))

        messages.append({"role": "user", "content": prompt})
        return messages

    def _build_timeout_plan(
        self,
        prompt: str,
        payload: Dict[str, Any],
        *,
        model: str,
        timeout_override: Optional[float],
    ) -> list[tuple[float, float]]:
        """Return (connect, read) timeout pairs ordered by attempt."""

        connect_timeout = 10.0
        if timeout_override is not None:
            return [(connect_timeout, max(1.0, float(timeout_override)))]

        is_local = self._is_local_base_url()

        if eb1_variant := self._canonical_eb1_variant(model):
            stages = _EB1_VARIANT_TIMEOUT_STAGES[eb1_variant]
            per_stage = _EB1_PER_STAGE_TIMEOUT_SECONDS["local" if is_local else "remote"]
            return [(connect_timeout, stages * per_stage)]

        read_timeout = self._estimate_standard_timeout(prompt, payload, is_local=is_local)
        return [(connect_timeout, read_timeout)]

    def _estimate_standard_timeout(
        self,
        prompt: str,
        payload: Dict[str, Any],
        *,
        is_local: bool,
    ) -> float:
        """Estimate a read timeout for non-EB1 models using lightweight heuristics."""

        est_prompt_tokens = max(1, len(prompt) // 4)
        max_tokens_hint = payload.get("max_tokens")
        try:
            completion_hint = int(max_tokens_hint) if max_tokens_hint else 1024
        except (TypeError, ValueError):
            completion_hint = 1024

        total_tokens = est_prompt_tokens + max(1, completion_hint)

        tokens_per_second = 45.0 if is_local else 55.0
        base_buffer = 60.0 if is_local else 45.0

        estimated = (total_tokens / tokens_per_second) + base_buffer
        lower_bound = 120.0 if is_local else 90.0
        upper_bound = 420.0 if is_local else 360.0

        return float(min(max(estimated, lower_bound), upper_bound))

    def _is_local_base_url(self) -> bool:
        """Return ``True`` when the provider points at a loopback Takumi gateway."""

        parsed = urlparse(self.base_url)
        host = (parsed.hostname or "").lower()
        return host in {"127.0.0.1", "localhost"}

    @staticmethod
    def _canonical_eb1_variant(model: str) -> Optional[str]:
        model_lower = model.lower()
        if model_lower.startswith("takumi/"):
            model_lower = model_lower.split("/", 1)[1]
        for prefix, variant in _EB1_VARIANT_PREFIXES:
            if model_lower.startswith(prefix):
                return variant
        return None


class TakumiEb1ARuntimeProvider(BaseProvider):
    """Local Takumi provider backed by the EB1-A router runtime."""

    requires_api_key = False
    supports_responses_api = True
    supports_streaming = True

    def complete(self, prompt: str, model: str, **kwargs: Any) -> ChatResponse:
        params = self._extract_invocation_kwargs(kwargs)
        response = run_eb1_a(
            prompt,
            system=params["system"],
            metadata=params["metadata"],
            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 "takumi/eb1-a"
        return response

    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": "takumi_eb1_a"},
            )

        request = parse_responses_request(payload)
        components = build_prompt_components(request)

        prompt_text = _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 = _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_metadata = merged_kwargs.get("metadata")
            if isinstance(existing_metadata, Mapping):
                combined = dict(metadata_payload)
                combined.update(dict(existing_metadata))
                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),
            )

        return self.complete(prompt_text, request.model, **merged_kwargs)

    def _extract_invocation_kwargs(self, params: Mapping[str, Any]) -> dict[str, Any]:
        working = dict(params)

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

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

        for key in list(working.keys()):
            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_context = self._normalize_context_messages(
            context_messages if context_messages is not None else context
        )

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

        return {
            "system": normalized_system,
            "metadata": metadata,
            "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 _to_optional_mapping(value: Any) -> Optional[dict[str, Any]]:
        if value is None:
            return None
        if isinstance(value, Mapping):
            return dict(value)
        return None

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

    @staticmethod
    def _normalize_context_messages(value: Any) -> Tuple[dict[str, str], ...]:
        if value is None:
            return tuple()
        if isinstance(value, str):
            text = value.strip()
            if not text:
                return tuple()
            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 tuple()
            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 _context_to_messages(context: Any) -> list[dict[str, str]]:
    """Normalise context payloads into chat messages."""
    if isinstance(context, str):
        return [{"role": "system", "content": context}]
    if isinstance(context, Sequence):
        normalised: list[dict[str, str]] = []
        for entry in context:
            if not isinstance(entry, Mapping):
                raise ValueError("context entries must be mappings with role/content")
            role = _normalise_role(entry.get("role"))
            content = entry.get("content")
            if not isinstance(content, str):
                raise ValueError("context message content must be text")
            normalised.append({"role": role, "content": content})
        return normalised
    raise ValueError("context must be a string or list of {role, content} objects")


def _collect_input_texts(value: Any) -> list[str]:
    texts: list[str] = []

    def _append_text(candidate: Any) -> None:
        if isinstance(candidate, str):
            texts.append(candidate)
        elif isinstance(candidate, bytes):
            texts.append(candidate.decode("utf-8", errors="ignore"))

    def _visit(candidate: Any) -> None:
        if isinstance(candidate, Mapping):
            content = candidate.get("content")
            if isinstance(content, Sequence) and not isinstance(content, (str, bytes, bytearray)):
                for part in content:
                    _visit(part)
            else:
                _append_text(content)
            _append_text(candidate.get("text"))
            return
        if isinstance(candidate, Sequence) and not isinstance(candidate, (str, bytes, bytearray)):
            for item in candidate:
                _visit(item)
            return
        _append_text(candidate)

    _visit(value)
    return texts


def _prompt_from_responses_input(
    input_value: Sequence[ResponsesMessage] | str | None,
) -> Optional[str]:
    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


def _context_from_responses_messages(
    messages: Sequence[ResponsesMessage] | None,
) -> Tuple[dict[str, str], ...]:
    if not messages:
        return tuple()
    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 _responses_message_text(message: ResponsesMessage) -> Optional[str]:
    fragments: list[str] = []
    for part in message.content:
        text = part.fields.get("text")
        if isinstance(text, str):
            fragments.append(text)
    combined = "".join(fragment.strip() for fragment in fragments if fragment.strip())
    return combined or None


def _is_gpt5_model(model: str) -> bool:
    candidate = (model or "").strip().lower()
    if candidate.startswith("takumi/"):
        candidate = candidate.split("/", 1)[1]
    return candidate == "gpt-5" or candidate.startswith("gpt-5-")


def _normalise_role(raw: Any) -> str:
    if raw is None:
        return "user"
    if isinstance(raw, str):
        role = raw.strip()
        return role or "user"
    raise ValueError("context role must be text")
