"""CompletionPipeline — structured request processing for /v1/chat/completions.

Breaks the monolithic chat_completions handler (CC=12, fan=22) into
discrete pipeline stages, each with CC ≤ 4.
"""

from __future__ import annotations

import json
import time
from dataclasses import dataclass, field
from typing import Any

import litellm
import structlog
from fastapi import HTTPException
from fastapi.responses import JSONResponse, StreamingResponse

from proxym.config import get_settings
from proxym.context import detect_project, ProjectContext
from proxym.router import AnalysisResult, analyze_request
from proxym.router.strategy import Router, RoutingDecision

from proxym.api._state import (
    get_context_payload,
    get_mcp_router,
    get_router,
    get_session_store,
)

logger = structlog.get_logger()


# ── Constants ────────────────────────────────────────────────

_TIER_ALIASES: dict[str, str] = {
    "cheap": "google/gemini-2.5-flash",
    "balanced": "google/gemini-3-flash-preview",
    "premium": "anthropic/opus-4.6",
    "free": "google/gemini-2.5-flash",
    "local": "ollama/qwen2.5-coder-3b",
}

_PASSTHROUGH_PARAMS = (
    "temperature", "max_tokens", "top_p", "stop", "tools",
    "tool_choice", "response_format", "seed",
)


# ── Provider key injection ───────────────────────────────────

def _set_provider_keys(kwargs: dict, model) -> None:
    """Inject the correct API key for a model's provider."""
    settings = get_settings()
    provider_keys = {
        "anthropic": ("api_key", settings.anthropic_api_key),
        "openai": ("api_key", settings.openai_api_key),
        "openrouter": ("api_key", settings.openrouter_api_key),
        "google": ("api_key", settings.google_ai_key),
        "deepseek": ("api_key", settings.deepseek_api_key),
        "groq": ("api_key", settings.groq_api_key),
        "together": ("api_key", settings.together_api_key),
        "mistral": ("api_key", settings.mistral_api_key),
        "fireworks": ("api_key", settings.fireworks_api_key),
        "cerebras": ("api_key", settings.cerebras_api_key),
        "ollama": ("api_base", settings.ollama_base_url),
    }
    entry = provider_keys.get(model.provider)
    if entry:
        key_name, key_value = entry
        if key_value:
            kwargs[key_name] = key_value


def _calculate_cost(model, prompt_tokens: int, completion_tokens: int) -> float:
    """Calculate actual cost from token counts."""
    input_cost = (prompt_tokens / 1_000_000) * model.input_cost
    output_cost = (completion_tokens / 1_000_000) * model.output_cost
    return input_cost + output_cost


# ── Pipeline ─────────────────────────────────────────────────

class CompletionPipeline:
    """Structured pipeline for processing a chat completion request.

    Stages:
      1. build_context  — inject delta context, detect project
      2. route          — analyze request + pick model
      3. inject_tools   — optionally inject MCP tools
      4. touch_session  — register/update session
      5. execute        — call LiteLLM with fallback chain
    """

    def __init__(self, body: dict, task_tier: str | None, headers: dict[str, str]):
        self.body = body
        self.task_tier = task_tier
        self.headers = {k.lower(): v for k, v in headers.items()}
        self.messages: list[dict] = body.get("messages", [])
        self.requested_model: str | None = body.get("model")
        self.stream: bool = body.get("stream", False)

        # Populated during pipeline stages
        self.project_ctx: ProjectContext | None = None
        self.analysis: AnalysisResult | None = None
        self.decision: RoutingDecision | None = None
        self.rtr: Router | None = None

    async def execute(self) -> StreamingResponse | JSONResponse:
        """Run the full pipeline and return the response."""
        self._build_context()
        self._route()
        self._inject_tools()
        self._touch_session()
        return await self._run()

    def _build_context(self) -> None:
        """Stage 1: Inject delta context and detect project."""
        if self.headers.get("x-inject-context", "").lower() == "true":
            _inject_context(self.messages)
        self.project_ctx = detect_project(self.messages, headers=self.headers)

    def _route(self) -> None:
        """Stage 2: Analyze request and pick model."""
        self.analysis = analyze_request(
            self.messages,
            explicit_tier=self.task_tier,
            explicit_model=(
                self.requested_model
                if self.requested_model and "/" in self.requested_model
                else None
            ),
        )
        self.rtr = get_router()
        explicit = _resolve_model_alias(self.requested_model)
        self.decision = self.rtr.route(self.analysis, explicit_model=explicit)

    def _inject_tools(self) -> None:
        """Stage 3: Inject MCP tools if enabled and not already provided."""
        if self.headers.get("x-enable-tools", "").lower() != "true":
            return
        if "tools" in self.body:
            return
        mcp_router = get_mcp_router()
        servers = mcp_router.select_servers(self.project_ctx, self.analysis)
        if not servers:
            return
        mcp_tools = mcp_router.get_tools_for_servers(servers)
        if mcp_tools:
            self.body["tools"] = mcp_tools
            self.project_ctx.active_mcp_servers = [s.name for s in servers]

    def _touch_session(self) -> None:
        """Stage 4: Register or update agent session."""
        session_id = self.headers.get("x-session-id", "")
        if not session_id and not self.project_ctx.project_id:
            return
        store = get_session_store()
        session = store.get_or_create(
            session_id=session_id, project_id=self.project_ctx.project_id,
        )
        self.project_ctx.session_id = session.session_id

    async def _run(self) -> StreamingResponse | JSONResponse:
        """Stage 5: Build kwargs and execute with fallback chain."""
        kwargs = self._build_kwargs()
        return await self._execute_with_fallbacks(kwargs)

    def _build_kwargs(self) -> dict:
        """Build LiteLLM completion kwargs from request body and routing decision."""
        kwargs = {
            "model": self.decision.model.litellm_model,
            "messages": self.messages,
            "stream": self.stream,
        }
        for key in _PASSTHROUGH_PARAMS:
            if key in self.body:
                kwargs[key] = self.body[key]
        _set_provider_keys(kwargs, self.decision.model)
        return kwargs

    async def _execute_with_fallbacks(self, kwargs: dict) -> StreamingResponse | JSONResponse:
        """Try primary model + fallbacks, return first successful response."""
        t0 = time.monotonic()
        all_models = [self.decision.model] + self.decision.fallbacks

        last_error = None
        for i, model_spec in enumerate(all_models):
            try:
                kwargs["model"] = model_spec.litellm_model
                _set_provider_keys(kwargs, model_spec)

                if self.stream:
                    return self._make_stream_response(kwargs, model_spec, t0)
                else:
                    return await self._make_non_stream_response(kwargs, model_spec, t0, i)
            except Exception as e:
                last_error = e
                logger.warning(
                    "model_failed",
                    model=model_spec.id,
                    error=str(e)[:200],
                    fallback_index=i,
                )
                continue

        raise HTTPException(
            status_code=502,
            detail=f"All models failed. Last error: {str(last_error)[:200]}",
        )

    def _make_stream_response(self, kwargs, model_spec, t0) -> StreamingResponse:
        """Build streaming response."""
        response = litellm.acompletion(**kwargs)
        return StreamingResponse(
            self._stream_wrapper(response, model_spec, t0),
            media_type="text/event-stream",
            headers={
                "X-Model-Id": model_spec.id,
                "X-Task-Tier": self.analysis.tier.value,
                "X-Routing-Reason": self.decision.reason[:100],
            },
        )

    async def _make_non_stream_response(self, kwargs, model_spec, t0, fallback_index) -> JSONResponse:
        """Execute non-streaming completion, record cost, return JSONResponse."""
        response = await litellm.acompletion(**kwargs)
        elapsed = time.monotonic() - t0
        self.rtr.record_latency(model_spec.id, elapsed)

        resp_data = response.model_dump() if hasattr(response, "model_dump") else dict(response)

        usage = resp_data.get("usage") or {}
        prompt_tokens = usage.get("prompt_tokens", 0) or 0
        completion_tokens = usage.get("completion_tokens", 0) or 0
        cost = _calculate_cost(model_spec, prompt_tokens, completion_tokens)
        self.rtr.record_cost(cost)

        resp_data["_proxy"] = {
            "model_id": model_spec.id,
            "tier": self.analysis.tier.value,
            "cost_usd": round(cost, 6),
            "routing_reason": self.decision.reason,
            "elapsed_ms": round(elapsed * 1000, 1),
            "fallback_index": fallback_index,
        }

        return JSONResponse(
            content=resp_data,
            headers={
                "X-Model-Id": model_spec.id,
                "X-Task-Tier": self.analysis.tier.value,
                "X-Cost-Usd": f"{cost:.6f}",
                "X-Routing-Reason": self.decision.reason[:100],
            },
        )

    async def _stream_wrapper(self, response_coro, model_spec, t0):
        """Wrap streaming response to track cost after completion."""
        response = await response_coro
        prompt_tokens = 0
        completion_tokens = 0

        async for chunk in response:
            if hasattr(chunk, "usage") and chunk.usage:
                prompt_tokens = getattr(chunk.usage, "prompt_tokens", 0) or 0
                completion_tokens = getattr(chunk.usage, "completion_tokens", 0) or 0

            yield f"data: {json.dumps(chunk.model_dump() if hasattr(chunk, 'model_dump') else dict(chunk))}\n\n"

        elapsed = time.monotonic() - t0
        self.rtr.record_latency(model_spec.id, elapsed)
        cost = _calculate_cost(model_spec, prompt_tokens, completion_tokens)
        self.rtr.record_cost(cost)

        yield "data: [DONE]\n\n"


# ── Module-level helpers (used by pipeline) ──────────────────

def _inject_context(messages: list[dict]) -> None:
    """Inject delta context payload into messages (mutates in place)."""
    payload = get_context_payload()
    if not payload:
        return
    if messages and messages[0].get("role") == "system":
        messages[0]["content"] += f"\n\n{payload}"
    else:
        messages.insert(0, {
            "role": "system",
            "content": f"Current codebase context:\n{payload}",
        })


def _resolve_model_alias(requested_model: str | None) -> str | None:
    """Resolve tier aliases (cheap, balanced, etc.) to model IDs."""
    if not requested_model:
        return None
    return _TIER_ALIASES.get(requested_model, requested_model)
