"""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

try:
    import litellm
    LITELLM_AVAILABLE = True
except ImportError:
    # llx provides its own completion interface
    LITELLM_AVAILABLE = False
    litellm = None

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": "openrouter/step-3.5-flash:free",
    "balanced": "openrouter/step-3.5-flash:free",
    "premium": "anthropic/opus-4.6",
    "free": "openrouter/step-3.5-flash:free",
    "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."""
        if not LITELLM_AVAILABLE:
            # Use llx streaming if available
            try:
                from llx.routing.client import LlxClient, ChatMessage
                from llx.config import LlxConfig
                messages = kwargs.get("messages", [])
                client = LlxClient(LlxConfig.load())
                response = client.stream_chat([ChatMessage(role=m["role"], content=m["content"]) for m in messages])
                return StreamingResponse(
                    self._stream_wrapper_llx(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],
                    },
                )
            except ImportError:
                raise HTTPException(status_code=502, detail="LLX not available for streaming")
        
        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."""
        if not LITELLM_AVAILABLE:
            # Use llx if available
            try:
                from llx.routing.client import LlxClient, ChatMessage
                from llx.config import LlxConfig
                messages = kwargs.get("messages", [])
                client = LlxClient(LlxConfig.load())
                response = client.chat([ChatMessage(role=m["role"], content=m["content"]) for m in messages])
                
                # Convert llx response to litellm-like format
                elapsed = time.monotonic() - t0
                self.rtr.record_latency(model_spec.id, elapsed)
                
                resp_data = {
                    "id": f"llx-{response.id}",
                    "object": "chat.completion",
                    "created": int(time.time()),
                    "model": model_spec.litellm_model,
                    "choices": [{
                        "index": 0,
                        "message": {
                            "role": "assistant",
                            "content": response.content,
                        },
                        "finish_reason": "stop",
                    }],
                    "usage": {
                        "prompt_tokens": response.usage.prompt_tokens if hasattr(response, 'usage') else 0,
                        "completion_tokens": response.usage.completion_tokens if hasattr(response, 'usage') else 0,
                        "total_tokens": response.usage.total_tokens if hasattr(response, 'usage') else 0,
                    },
                }
                
                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-Routing-Reason": self.decision.reason[:100],
                        "X-Cost-USD": f"{cost:.6f}",
                    },
                )
            except ImportError:
                raise HTTPException(status_code=502, detail="LLX not available")
        
        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"

    async def _stream_wrapper_llx(self, response_iter, model_spec, t0):
        """Wrap llx streaming response to track cost after completion."""
        prompt_tokens = 0
        completion_tokens = 0
        
        for chunk in response_iter:
            # Convert llx chunk to OpenAI-compatible format
            chunk_data = {
                "id": f"llx-{chunk.get('id', 'unknown')}",
                "object": "chat.completion.chunk",
                "created": int(time.time()),
                "model": model_spec.litellm_model,
                "choices": [{
                    "index": 0,
                    "delta": {"content": chunk.get("content", "")},
                    "finish_reason": chunk.get("finish_reason", None),
                }],
            }
            
            yield f"data: {json.dumps(chunk_data)}\n\n"
            
            # Track tokens if available
            if "usage" in chunk:
                prompt_tokens += chunk["usage"].get("prompt_tokens", 0)
                completion_tokens += chunk["usage"].get("completion_tokens", 0)

        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)
