"""
FastAPI AI router gateway — Zed entrypoint with tier/lane orchestration.
"""

from __future__ import annotations

import hashlib
import json
import logging
import os
import time
from typing import Any

import httpx
from fastapi import FastAPI, Header, HTTPException, Request, Response
from fastapi.responses import JSONResponse, StreamingResponse

from metrics import CLASSIFY, DURATION, ESCALATIONS, REQUESTS, metrics_payload
from orchestrator import Orchestrator, SessionStore, Tier
from rules_loader import reload_configs

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("ai-router")

LITELLM_URL = os.environ.get("LITELLM_INTERNAL_URL", "http://litellm:4000").rstrip("/")
LITELLM_KEY = os.environ.get("LITELLM_MASTER_KEY", "")
ROUTER_API_KEY = os.environ.get("ROUTER_API_KEY", "")
DEFAULT_MODEL = os.environ.get("DEFAULT_LITELLM_MODEL", "smart-router")
REDIS_URL = os.environ.get("REDIS_URL", "redis://redis:6379/0")

session_store = SessionStore(REDIS_URL or None)
orchestrator = Orchestrator(session_store)

app = FastAPI(title="EventHub AI Router", version="2.0.0")


def _extract_text(messages: list[dict[str, Any]]) -> str:
    parts: list[str] = []
    for msg in messages:
        content = msg.get("content", "")
        if isinstance(content, str):
            parts.append(content)
        elif isinstance(content, list):
            for block in content:
                if isinstance(block, dict) and block.get("type") == "text":
                    parts.append(str(block.get("text", "")))
    return "\n".join(parts)


def _has_image(messages: list[dict[str, Any]]) -> bool:
    for msg in messages:
        content = msg.get("content", "")
        if not isinstance(content, list):
            continue
        for block in content:
            if isinstance(block, dict) and block.get("type") in ("image_url", "image"):
                return True
    return False


def _estimate_tokens(text: str) -> int:
    return max(1, len(text) // 4)


def _auth_or_403(authorization: str | None) -> None:
    if not ROUTER_API_KEY:
        return
    if not authorization or not authorization.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Missing Bearer token")
    token = authorization.removeprefix("Bearer ").strip()
    if token != ROUTER_API_KEY:
        raise HTTPException(status_code=403, detail="Invalid API key")


def _quality_mode(header: str | None, body: dict[str, Any]) -> str | None:
    if header:
        mode = header.strip().lower()
        if mode in ("auto", "economy", "balanced", "max"):
            return mode
    meta = body.get("metadata") or {}
    if isinstance(meta, dict):
        mode = str(meta.get("quality_mode", "")).lower()
        if mode in ("auto", "economy", "balanced", "max"):
            return mode
    return None


def _session_id(body: dict[str, Any]) -> str:
    meta = body.get("metadata") or {}
    if isinstance(meta, dict) and meta.get("session_id"):
        return str(meta["session_id"])
    if body.get("user"):
        return str(body["user"])
    return ""


def _router_meta(decision, *, requested: str) -> dict[str, Any]:
    return {
        "tier": decision.tier.value,
        "lane": decision.lane,
        "model": decision.model,
        "escalation_level": decision.escalation_level,
        "quality_mode": decision.quality_mode,
        "confidence": round(decision.confidence, 3),
        "requested_model": requested,
        "delegated_internal": decision.delegated_internal,
    }


@app.get("/health")
async def health() -> dict[str, str]:
    return {"status": "ok"}


@app.get("/metrics")
async def metrics() -> Response:
    body, content_type = metrics_payload()
    return Response(content=body, media_type=content_type)


@app.post("/admin/reload-config")
async def admin_reload(authorization: str | None = Header(default=None)) -> dict[str, str]:
    _auth_or_403(authorization)
    reload_configs()
    global orchestrator  # noqa: PLW0603
    orchestrator = Orchestrator(session_store)
    return {"status": "reloaded"}


@app.get("/v1/models")
async def list_models(authorization: str | None = Header(default=None)) -> dict[str, Any]:
    _auth_or_403(authorization)
    async with httpx.AsyncClient(timeout=30.0) as client:
        resp = await client.get(
            f"{LITELLM_URL}/v1/models",
            headers={"Authorization": f"Bearer {LITELLM_KEY}"},
        )
    if resp.status_code >= 400:
        raise HTTPException(status_code=resp.status_code, detail=resp.text)
    return resp.json()


@app.post("/classify")
async def classify_debug(
    request: Request,
    authorization: str | None = Header(default=None),
    x_ai_quality: str | None = Header(default=None, alias="X-AI-Quality"),
) -> dict[str, Any]:
    _auth_or_403(authorization)
    body = await request.json()
    messages = body.get("messages") or []
    text = _extract_text(messages)
    decision = orchestrator.resolve(
        messages,
        quality_mode=_quality_mode(x_ai_quality, body),
        session_id=_session_id(body),
        text=text,
        has_image=_has_image(messages),
        token_estimate=_estimate_tokens(text),
    )
    CLASSIFY.labels(tier=decision.tier.value).inc()
    return {
        "tier": decision.tier.value,
        "lane": decision.lane,
        "model": decision.model,
        "x_router_meta": _router_meta(decision, requested=body.get("model", DEFAULT_MODEL)),
    }


async def _forward_litellm(
    forward: dict[str, Any],
    *,
    stream: bool,
    decision_meta: dict[str, Any],
) -> Any:
    headers = {
        "Authorization": f"Bearer {LITELLM_KEY}",
        "Content-Type": "application/json",
    }
    start = time.perf_counter()
    tier = decision_meta.get("tier", "UNKNOWN")
    model = forward.get("model", DEFAULT_MODEL)

    async with httpx.AsyncClient(timeout=180.0) as client:
        if stream:
            req = client.build_request(
                "POST",
                f"{LITELLM_URL}/v1/chat/completions",
                headers=headers,
                content=json.dumps(forward),
            )
            resp = await client.send(req, stream=True)

            async def event_stream():
                async for chunk in resp.aiter_bytes():
                    yield chunk

            REQUESTS.labels(tier=tier, lane=decision_meta["lane"], model=model, status=str(resp.status_code)).inc()
            DURATION.labels(tier=tier, model=model).observe(time.perf_counter() - start)
            return StreamingResponse(
                event_stream(),
                status_code=resp.status_code,
                media_type=resp.headers.get("content-type", "text/event-stream"),
                headers={"X-Router-Meta": json.dumps(decision_meta)},
            )

        resp = await client.post(
            f"{LITELLM_URL}/v1/chat/completions",
            headers=headers,
            json=forward,
        )

    status = str(resp.status_code)
    REQUESTS.labels(tier=tier, lane=decision_meta["lane"], model=model, status=status).inc()
    DURATION.labels(tier=tier, model=model).observe(time.perf_counter() - start)

    if resp.status_code >= 400:
        return JSONResponse(status_code=resp.status_code, content=resp.json())

    data = resp.json()
    if isinstance(data, dict):
        data["x_router_meta"] = decision_meta
    return data


@app.post("/v1/chat/completions")
async def chat_completions(
    request: Request,
    authorization: str | None = Header(default=None),
    x_ai_quality: str | None = Header(default=None, alias="X-AI-Quality"),
) -> Any:
    _auth_or_403(authorization)

    body: dict[str, Any] = await request.json()
    messages = body.get("messages") or []
    if not isinstance(messages, list):
        raise HTTPException(status_code=400, detail="messages must be a list")

    text = _extract_text(messages)
    has_image = _has_image(messages)
    session_id = _session_id(body)
    prompt_hash = hashlib.sha256(text.encode()).hexdigest()[:16]

    requested_model = body.get("model", DEFAULT_MODEL)
    if requested_model in ("smart-router", "auto", ""):
        decision = orchestrator.resolve(
            messages,
            quality_mode=_quality_mode(x_ai_quality, body),
            session_id=session_id,
            text=text,
            has_image=has_image,
            token_estimate=_estimate_tokens(text),
        )
        target_model = decision.model
        meta = _router_meta(decision, requested=requested_model)
    else:
        target_model = requested_model
        meta = {"selected_model": target_model, "requested_model": requested_model}

    forward = dict(body)
    forward["model"] = target_model
    forward.setdefault("metadata", {})
    if isinstance(forward["metadata"], dict):
        forward["metadata"]["semantic_tier"] = meta.get("tier")
        forward["metadata"]["session_id"] = session_id or forward["metadata"].get("session_id")

    log.info(
        "route tier=%s lane=%s model=%s requested=%s",
        meta.get("tier"),
        meta.get("lane"),
        target_model,
        requested_model,
    )

    stream = bool(body.get("stream", False))
    result = await _forward_litellm(forward, stream=stream, decision_meta=meta)

    escalate = isinstance(result, JSONResponse) and result.status_code >= 429
    orchestrator.after_request(
        session_id,
        prompt_hash=prompt_hash,
        success=not escalate,
        escalate=escalate,
    )
    if escalate and meta.get("lane") != "C":
        ESCALATIONS.labels(from_lane=meta.get("lane", "?"), to_lane="next", reason="http_error").inc()

    return result