"""Chat completions endpoint — delegates to CompletionPipeline.

The endpoint is a thin wrapper; all logic lives in _pipeline.py.
"""

from __future__ import annotations

import json

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

from proxym.api._pipeline import CompletionPipeline
from proxym.api._state import get_router
from proxym.config import get_settings
from proxym.providers import MODELS, TaskTier
from proxym.observability import log_call
from proxym.router import analyze_request
from proxym.router.strategy import CostLedger


def _inject_provider_key(kwargs: dict, provider: str) -> None:
    """Inject the correct API key for a provider into litellm kwargs."""
    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(provider)
    if entry:
        key_name, key_value = entry
        if key_value:
            kwargs[key_name] = key_value

router = APIRouter()
ledger = CostLedger()


def _resolve_model(model: str | None) -> str | None:
    if not model:
        return model
    return get_settings().resolve_model_alias(model)


def _analyze_request(messages, x_task_tier: str | None = None):
    return analyze_request(messages, explicit_tier=x_task_tier)


def _find_model_spec(name: str):
    """Find a ModelSpec by id, litellm_model, or partial match."""
    if name in MODELS:
        return MODELS[name]
    for spec in MODELS.values():
        if spec.litellm_model == name or spec.litellm_model.endswith("/" + name):
            return spec
    return None


def _route_to_model(messages, model: str | None = None, x_task_tier: str | None = None):
    resolved = _resolve_model(model)
    if resolved:
        spec = _find_model_spec(resolved)
        if spec:
            return spec.litellm_model, spec.provider
    analysis = _analyze_request(messages, x_task_tier)
    if isinstance(analysis, str):
        tier_map = {
            "simple": TaskTier.TRIVIAL,
            "trivial": TaskTier.TRIVIAL,
            "complex": TaskTier.COMPLEX,
            "deep": TaskTier.DEEP_REASONING,
            "standard": TaskTier.STANDARD,
        }
        analysis = analyze_request(messages, explicit_tier=tier_map.get(analysis, TaskTier.STANDARD))
    decision = get_router().route(analysis)
    return decision.model.litellm_model, decision.model.provider


async def completion(**kwargs):
    if kwargs.get("stream"):
        from proxym.main import litellm as main_litellm
        return main_litellm.acompletion(**kwargs)
    from proxym.main import litellm as main_litellm
    return await main_litellm.acompletion(**kwargs)


def _usage_tokens(resp) -> tuple[int, int]:
    usage = getattr(resp, "usage", None)
    if usage is None and isinstance(resp, dict):
        usage = resp.get("usage")
    if isinstance(usage, dict):
        prompt_tokens = usage.get("prompt_tokens", 0)
        completion_tokens = usage.get("completion_tokens", 0)
    elif usage is not None:
        prompt_tokens = getattr(usage, "prompt_tokens", 0)
        completion_tokens = getattr(usage, "completion_tokens", 0)
    else:
        prompt_tokens, completion_tokens = 0, 0
    # Guard against non-int values (e.g. MagicMock)
    if not isinstance(prompt_tokens, (int, float)):
        prompt_tokens = 0
    if not isinstance(completion_tokens, (int, float)):
        completion_tokens = 0
    return int(prompt_tokens) or 0, int(completion_tokens) or 0


def _response_to_dict(resp):
    if isinstance(resp, dict):
        return resp
    if hasattr(resp, "model_dump"):
        dumped = resp.model_dump()
        if isinstance(dumped, dict):
            return dumped
    if hasattr(resp, "choices"):
        message = getattr(resp.choices[0], "message", None)
        raw_content = getattr(message, "content", "")
        content = raw_content if isinstance(raw_content, str) else str(raw_content)
        tool_calls = getattr(message, "tool_calls", None)
        if not isinstance(tool_calls, list):
            tool_calls = None
        return {
            "choices": [
                {
                    "message": {
                        "content": content,
                        "tool_calls": tool_calls,
                    }
                }
            ],
            "usage": {
                "prompt_tokens": _usage_tokens(resp)[0],
                "completion_tokens": _usage_tokens(resp)[1],
            },
            "model": getattr(resp, "model", None),
        }
    return dict(resp)


@router.post("/v1/chat/completions")
@log_call(level="info", slow_threshold_ms=5000)
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)
    """
    try:
        body = await request.json()
    except json.JSONDecodeError as exc:
        raise HTTPException(status_code=400, detail="invalid json") from exc

    if not body.get("messages"):
        raise HTTPException(status_code=422, detail="messages must not be empty")

    if request.headers.get("x-enable-tools", "").lower() == "true" or request.headers.get("x-inject-context", "").lower() == "true":
        pipeline = CompletionPipeline(body, x_task_tier, headers=dict(request.headers))
        response = await pipeline.execute()
        if isinstance(response, JSONResponse) and "X-Cost-Usd" in response.headers:
            response.headers["X-Proxym-Cost"] = response.headers["X-Cost-Usd"]
        return response

    model, provider = _route_to_model(body.get("messages", []), body.get("model"), x_task_tier)
    analysis = _analyze_request(body.get("messages", []), x_task_tier)
    if isinstance(analysis, str):
        analysis = analyze_request(body.get("messages", []), explicit_tier=analysis)

    kwargs = {k: v for k, v in body.items() if k != "model"}
    kwargs["model"] = model
    _inject_provider_key(kwargs, provider)

    settings = get_settings()
    fallback_model = settings.resolve_model_alias("cheap")

    try:
        resp = await _try_providers(kwargs, body, fallback_model, settings)
        if body.get("stream"):
            async def _stream():
                async for chunk in resp:
                    payload = chunk.model_dump() if hasattr(chunk, "model_dump") else dict(chunk)
                    yield f"data: {json.dumps(payload)}\n\n"
                yield "data: [DONE]\n\n"
            return StreamingResponse(_stream(), media_type="text/event-stream")
    except HTTPException:
        raise
    except Exception as exc:
        raise HTTPException(status_code=502, detail=str(exc)) from exc

    return _format_response(resp, body, model, analysis, ledger)


async def _try_providers(kwargs: dict, body: dict, fallback_model: str, settings) -> Any:
    """Try completion with primary model, fallback on error."""
    try:
        if body.get("stream"):
            return await completion(**kwargs)
        return await completion(**kwargs)
    except httpx.HTTPStatusError as exc:
        return await _handle_provider_error(exc, kwargs, body, fallback_model, settings)


async def _handle_provider_error(exc: httpx.HTTPStatusError, kwargs: dict, body: dict, fallback_model: str, settings):
    """Handle provider error with fallback logic."""
    status_code = getattr(exc.response, "status_code", 500)
    if status_code not in (429, 500):
        raise HTTPException(status_code=502, detail=str(exc)) from exc

    model = kwargs.get("model", "")
    fb = fallback_model if fallback_model != model else settings.resolve_model_alias("balanced")
    fb_spec = _find_model_spec(fb)
    kwargs["model"] = fb_spec.litellm_model if fb_spec else fb
    if fb_spec:
        _inject_provider_key(kwargs, fb_spec.provider)
    try:
        return await completion(**kwargs)
    except Exception as inner:
        raise HTTPException(status_code=503, detail=str(inner)) from inner


def _format_response(resp, body: dict, model: str, analysis, ledger) -> JSONResponse:
    """Format completion response with metadata."""
    prompt_tokens, completion_tokens = _usage_tokens(resp)
    cost = ((prompt_tokens + completion_tokens) / 1_000_000) if (prompt_tokens or completion_tokens) else 0.0
    ledger.record(cost)
    data = _response_to_dict(resp)
    resolved_model = _resolve_model(body.get("model")) or model
    spec = _find_model_spec(model) or _find_model_spec(resolved_model)
    model_id = spec.id if spec else resolved_model
    data["_proxy"] = {
        "model_id": model_id,
        "tier": analysis.tier.value,
        "cost_usd": round(cost, 6),
        "routing_reason": getattr(analysis, "reasoning", "legacy-route"),
        "elapsed_ms": 0.0,
        "fallback_index": 0,
    }
    return JSONResponse(
        content=data,
        headers={
            "X-Proxym-Cost": f"{cost:.6f}",
            "X-Cost-Usd": f"{cost:.6f}",
            "X-Task-Tier": analysis.tier.value,
            "X-Routing-Reason": getattr(analysis, "reasoning", "legacy-route")[:100],
            "X-Model-Id": str(model_id),
        },
    )
