from __future__ import annotations

import logging
import time
from dataclasses import dataclass
from typing import Any, Iterable, Mapping, Optional, Protocol

from ember.models.providers.utils import ResponsesRequest

from ..domain.model import Message, Usage
from .types import OrgLimits

LOG = logging.getLogger(__name__)


@dataclass
class ExecutionResult:
    """Container describing the outcome of a provider invocation."""

    text: str | None
    stream: Iterable[str] | None
    usage: Usage
    provider_metadata: Mapping[str, object]
    responses_payload: Mapping[str, object] | None
    stream_state: Optional[object] = None
    provider_latency_ms: float = 0.0


class ProviderInvoker(Protocol):
    """Protocol consumed by :class:`Executor` for provider interactions."""

    def complete(
        self,
        model: str,
        prompt: str,
        *,
        temperature: float,
        api_key_override: str | None,
        system: str | None,
        context: list[dict[str, str]],
        responses_payload: Mapping[str, Any],
        org_limits: OrgLimits,
    ) -> tuple[str, int, int, Mapping[str, object]]: ...

    def stream_complete(
        self,
        model: str,
        prompt: str,
        *,
        temperature: float,
        api_key_override: str | None,
        system: str | None,
        context: list[dict[str, str]],
        responses_payload: Mapping[str, Any],
        org_limits: OrgLimits,
        enable_judge_stream: bool = False,
    ) -> Iterable[str]: ...


class Executor:
    """Bridges provider clients with usage accounting."""

    def __init__(
        self, *, providers: Mapping[str, ProviderInvoker], enable_judge_stream: bool = False
    ) -> None:
        self._providers = providers
        self._enable_judge_stream = enable_judge_stream

    def execute(
        self,
        *,
        org_id: str,
        request_id: str,
        provider: str,
        model: str,
        responses_request: ResponsesRequest,
        prompt_messages: list[Message],
        stream: bool,
        temperature: float | None,
        api_key_override: str | None,
        limits: OrgLimits,
        plan: Mapping[str, object],
    ) -> ExecutionResult:
        """Invoke the configured provider and capture usage metadata.

        Args:
            org_id: Organization issuing the request.
            request_id: Stable request identifier.
            provider: Selected provider slug.
            model: Canonical model identifier.
            responses_request: Normalized Responses payload.
            prompt_messages: Messages used to construct the prompt.
            stream: Whether the caller requested streaming output.
            temperature: Optional sampling temperature override.
            api_key_override: Provider-specific API key when required.
            limits: Prompt/completion budgets plus the fan-out multiplier enforced upstream.
            plan: Resolved billing plan metadata.

        Returns:
            Execution result capturing text, usage, and metadata.
        """
        client = self._providers[provider]
        LOG.info(
            "executor_start",
            extra={
                "request_id": request_id,
                "provider": provider,
                "model": model,
                "stream": stream,
            },
        )
        prompt_text, system, context = self._shape_prompt(prompt_messages)
        payload = dict(responses_request.original_payload)
        payload.setdefault("model", model)
        if stream:
            payload["stream"] = True
        else:
            payload.pop("stream", None)

        temp_value = float(temperature) if temperature is not None else 0.0
        org_limits = limits

        if stream:
            stream_wrapper = _StreamingWrapper(
                client.stream_complete(
                    model,
                    prompt_text,
                    temperature=temp_value,
                    api_key_override=api_key_override,
                    system=system,
                    context=context,
                    responses_payload=payload,
                    org_limits=org_limits,
                    enable_judge_stream=self._enable_judge_stream,
                )
            )
            LOG.info(
                "executor_streaming",
                extra={"request_id": request_id, "provider": provider},
            )
            return ExecutionResult(
                text=None,
                stream=stream_wrapper,
                usage=Usage(),
                provider_metadata={},
                responses_payload=None,
                stream_state=stream_wrapper,
            )

        t0 = time.perf_counter()
        text, pt, ct, meta = client.complete(
            model,
            prompt_text,
            temperature=temp_value,
            api_key_override=api_key_override,
            system=system,
            context=context,
            responses_payload=payload,
            org_limits=org_limits,
        )
        provider_latency_ms = (time.perf_counter() - t0) * 1000.0
        usage = Usage(
            prompt_tokens=int(pt),
            completion_tokens=int(ct),
            total_tokens=int(pt) + int(ct),
        )
        if isinstance(meta, Mapping):
            meta_mapping: Mapping[str, object] = dict(meta)
            responses_payload: Mapping[str, object] | None = meta_mapping
        else:
            meta_mapping = {"provider_metadata": meta}
            responses_payload = None
        LOG.info(
            "executor_complete",
            extra={
                "request_id": request_id,
                "provider": provider,
                "prompt_tokens": usage.prompt_tokens,
                "completion_tokens": usage.completion_tokens,
            },
        )
        return ExecutionResult(
            text=text,
            stream=None,
            usage=usage,
            provider_metadata=meta_mapping,
            responses_payload=responses_payload,
            provider_latency_ms=provider_latency_ms,
        )

    def _shape_prompt(
        self, messages: list[Message]
    ) -> tuple[str, str | None, list[dict[str, str]]]:
        """Convert the canonical message list into legacy prompt fields.

        Args:
            messages: Ordered conversation history destined for the provider.

        Returns:
            Tuple containing the user prompt, optional system prompt, and a
            list of prior conversational turns represented as role/content
            dictionaries.
        """
        if not messages:
            return "", None, []
        final = messages[-1]
        system = next((m.content for m in messages if m.role == "system"), None)
        ctx = [{"role": m.role, "content": m.content} for m in messages[:-1] if m.role != "system"]
        return final.content, system, ctx


class _StreamingWrapper:
    """Iterator wrapper capturing the terminal ChatResponse from providers."""

    def __init__(self, generator: Iterable[str]) -> None:
        """Initialize the wrapper with the provider generator."""
        self._generator = iter(generator)
        self.final_response: Optional[object] = None

    def __iter__(self) -> Iterable[str]:
        while True:
            try:
                chunk = next(self._generator)
            except StopIteration as stop:
                self.final_response = stop.value if hasattr(stop, "value") else None
                break
            yield chunk
