from __future__ import annotations

import logging
import time
from dataclasses import dataclass, field
from typing import Callable, Iterable, Optional, Union

_LOG = logging.getLogger(__name__)

from ..domain.model import (
    RBC,
    AdmissionDecision,
    ChatResponse,
    GatewayRequest,
    Message,
    RouterInfo,
    StreamingResponse,
)
from .container import Services

PipelineResult = Union[ChatResponse, StreamingResponse]
Step = Callable[[Services, "Context"], Optional[PipelineResult]]


@dataclass
class Context:
    req: GatewayRequest
    request_id: str
    canonical_model: str
    plan: dict[str, object]
    price_book_version: str
    provider_hint: Optional[str] = None
    alias_overrides: dict[str, str] = field(default_factory=dict)
    prompt_messages: list[Message] = field(default_factory=list)
    prompt_estimate: int = 0
    prompt_multiplier: int = 1
    prompt_budget: int = 0
    completion_budget: int = 0
    temperature_hint: Optional[float] = None
    admission: Optional[AdmissionDecision] = None
    rbc: Optional[RBC] = None
    rbc_token: Optional[str] = None
    chosen_provider: Optional[str] = None
    provider_model: Optional[str] = None
    router_info: RouterInfo = field(default_factory=RouterInfo)
    api_key_override: Optional[str] = None
    billing_mode: str = "resell"
    fbr_id: Optional[str] = None
    api_key_id: Optional[str] = None
    provider_metadata: dict[str, object] = field(default_factory=dict)
    finalizers: list[Callable[[], None]] = field(default_factory=list)
    step_timings_ms: dict[str, float] = field(default_factory=dict)
    pipeline_start_perf: float = 0.0


class ChatCompletionPipeline:
    def __init__(self, services: Services, steps: list[Step]) -> None:
        self._services = services
        self._steps = steps

    def run(self, ctx: Context) -> PipelineResult:
        result: Optional[PipelineResult] = None
        ctx.pipeline_start_perf = time.perf_counter()
        try:
            for step in self._steps:
                t0 = time.perf_counter()
                out = step(self._services, ctx)
                ctx.step_timings_ms[step.__name__] = (
                    (time.perf_counter() - t0) * 1000.0
                )
                if out is not None:
                    result = out
                    break
            if result is None:
                raise RuntimeError("Pipeline completed with no response.")
        finally:
            ctx.step_timings_ms["total_ms"] = (
                (time.perf_counter() - ctx.pipeline_start_perf) * 1000.0
            )
            if result is None or not isinstance(result, StreamingResponse):
                self._drain_finalizers(ctx)

        if isinstance(result, StreamingResponse):
            if result.stream is None:
                raise RuntimeError("streaming response missing stream generator")
            result.stream = self._wrap_stream(result.stream, ctx)
            return result
        return result

    def _drain_finalizers(self, ctx: Context) -> None:
        while ctx.finalizers:
            finalize = ctx.finalizers.pop()
            try:
                finalize()
            except Exception:
                _LOG.warning(
                    "Pipeline finalizer %r failed",
                    getattr(finalize, "__qualname__", finalize),
                    exc_info=True,
                )

    def _wrap_stream(self, stream: Iterable[str], ctx: Context) -> Iterable[str]:
        def generator() -> Iterable[str]:
            try:
                for chunk in stream:
                    yield chunk
            finally:
                self._drain_finalizers(ctx)

        return generator()
