"""Chat completions endpoint + helpers — extracted from main.py."""

from __future__ import annotations

import json
import time

import litellm
import structlog
from fastapi import APIRouter, Header, HTTPException, Request
from fastapi.responses import JSONResponse, StreamingResponse

from proxym.config import get_settings
from proxym.context import detect_project
from proxym.context.tool_selector import select_tools
from proxym.providers import get_model
from proxym.router import AnalysisResult, analyze_request
from proxym.router.strategy import RoutingDecision

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

logger = structlog.get_logger()

router = APIRouter()


# ── Main completion endpoint (OpenAI-compatible) ─────────────

@router.post("/v1/chat/completions")
async def chat_completions(
    request: Request,
    x_task_tier: str | None = Header(None, alias="X-Task-Tier"),
):
    """OpenAI-compatible chat completions with intelligent routing.

    Extra headers:
      X-Task-Tier: trivial|operational|standard|complex|deep
      X-Inject-Context: true  (inject latest delta context)
      X-Working-Directory: /path/to/project  (project detection)
      X-Session-Id: abc123  (session continuity)
      X-Enable-Tools: true  (inject MCP tools based on context)
    """
    body = await request.json()
    messages: list[dict] = body.get("messages", [])
    requested_model: str | None = body.get("model")
    stream: bool = body.get("stream", False)

    # 1. Inject delta context if requested
    if request.headers.get("X-Inject-Context", "").lower() == "true":
        _inject_context(messages)

    # 2. Detect project context
    headers = {k.lower(): v for k, v in request.headers.items()}
    project_ctx = detect_project(messages, headers=headers)

    # 3. Analyze + route
    analysis = analyze_request(
        messages,
        explicit_tier=x_task_tier,
        explicit_model=requested_model if requested_model and "/" in requested_model else None,
    )
    rtr = get_router()
    explicit = _resolve_model_alias(requested_model)
    decision = rtr.route(analysis, explicit_model=explicit)

    # 4. Inject MCP tools if enabled and not already provided
    if headers.get("x-enable-tools", "").lower() == "true" and "tools" not in body:
        mcp_router = get_mcp_router()
        servers = mcp_router.select_servers(project_ctx, analysis)
        if servers:
            mcp_tools = mcp_router.get_tools_for_servers(servers)
            if mcp_tools:
                body["tools"] = mcp_tools
                project_ctx.active_mcp_servers = [s.name for s in servers]

    # 5. Touch session
    session_id = headers.get("x-session-id", "")
    if session_id or project_ctx.project_id:
        store = get_session_store()
        session = store.get_or_create(
            session_id=session_id, project_id=project_ctx.project_id,
        )
        project_ctx.session_id = session.session_id

    # 6. Build kwargs + execute with fallback chain
    litellm_kwargs = _build_litellm_kwargs(body, messages, stream, decision)
    return await _execute_with_fallbacks(
        litellm_kwargs, decision, analysis, rtr, stream,
    )


# ── Helpers ──────────────────────────────────────────────────

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


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


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)


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


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


async def _execute_with_fallbacks(
    litellm_kwargs: dict,
    decision: RoutingDecision,
    analysis: AnalysisResult,
    rtr,
    stream: bool,
) -> StreamingResponse | JSONResponse:
    """Try primary model + fallbacks, return first successful response."""
    t0 = time.monotonic()
    all_models = [decision.model] + decision.fallbacks

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

            if stream:
                return await _handle_stream(
                    litellm_kwargs, model_spec, analysis, rtr, t0, decision,
                )
            else:
                return await _handle_non_stream(
                    litellm_kwargs, model_spec, analysis, rtr, t0, decision, 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]}",
    )


async def _handle_stream(
    litellm_kwargs, model_spec, analysis, rtr, t0, decision,
) -> StreamingResponse:
    """Execute streaming completion and return StreamingResponse."""
    response = await litellm.acompletion(**litellm_kwargs)
    return StreamingResponse(
        _stream_wrapper(response, model_spec, rtr, t0, decision),
        media_type="text/event-stream",
        headers={
            "X-Model-Id": model_spec.id,
            "X-Task-Tier": analysis.tier.value,
            "X-Routing-Reason": decision.reason[:100],
        },
    )


async def _handle_non_stream(
    litellm_kwargs, model_spec, analysis, rtr, t0, decision, fallback_index,
) -> JSONResponse:
    """Execute non-streaming completion, record cost, return JSONResponse."""
    response = await litellm.acompletion(**litellm_kwargs)
    elapsed = time.monotonic() - t0
    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)
    rtr.record_cost(cost)

    resp_data["_proxy"] = {
        "model_id": model_spec.id,
        "tier": analysis.tier.value,
        "cost_usd": round(cost, 6),
        "routing_reason": 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": analysis.tier.value,
            "X-Cost-Usd": f"{cost:.6f}",
            "X-Routing-Reason": decision.reason[:100],
        },
    )


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

    async for chunk in response:
        # Track tokens from stream
        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"

    # Record cost after stream completes
    elapsed = time.monotonic() - t0
    rtr.record_latency(model_spec.id, elapsed)
    cost = _calculate_cost(model_spec, prompt_tokens, completion_tokens)
    rtr.record_cost(cost)

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


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


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
