"""OpenAI-compatible endpoints. Requests are validated, the model is loaded on demand, then the
request is proxied to the local worker. Usage/timings are recorded from the worker's response."""

from __future__ import annotations

import asyncio
import json
import logging
import time
from typing import Any

import httpx
from fastapi import APIRouter, Depends, Request
from fastapi.responses import JSONResponse, StreamingResponse

from ..auth import Principal, require_inference
from ..errors import APIError, GenerationTimeout, WorkerCrashed
from ..events import bus
from ..manager import ModelManager

log = logging.getLogger("llm_api.openai")
router = APIRouter()


def _manager(request: Request) -> ModelManager:
    return request.app.state.manager


async def _parse_body(request: Request) -> dict:
    settings = request.app.state.settings
    raw = await request.body()
    if len(raw) > settings.max_body_bytes:
        raise APIError("Request body too large.", status_code=413, code="BODY_TOO_LARGE")
    try:
        body = json.loads(raw or b"{}")
    except json.JSONDecodeError:
        raise APIError("Request body is not valid JSON.")
    if not isinstance(body, dict):
        raise APIError("Request body must be a JSON object.")
    return body


def _validate_sampling(body: dict) -> None:
    def num(k, lo, hi):
        v = body.get(k)
        if v is None:
            return
        if not isinstance(v, (int, float)) or v < lo or v > hi:
            raise APIError(f"{k} must be a number between {lo} and {hi}.", param=k)
    num("temperature", 0, 2)
    num("top_p", 0, 1)
    num("presence_penalty", -2, 2)
    num("frequency_penalty", -2, 2)
    mt = body.get("max_tokens", body.get("max_completion_tokens"))
    if mt is not None and (not isinstance(mt, int) or mt < 1 or mt > 200000):
        raise APIError("max_tokens must be a positive integer.", param="max_tokens")
    n = body.get("n")
    if n not in (None, 1):
        raise APIError("Only n=1 is supported.", param="n")
    stop = body.get("stop")
    if stop is not None and not isinstance(stop, (str, list)):
        raise APIError("stop must be a string or an array of strings.", param="stop")
    if isinstance(stop, list) and len(stop) > 8:
        raise APIError("At most 8 stop sequences are supported.", param="stop")


async def _pick_model(request: Request, body: dict, endpoint: str) -> tuple[str, dict]:
    """Resolve model name (alias/auto) -> registry row."""
    manager = _manager(request)
    name = body.get("model")
    if not name or not isinstance(name, str):
        # default model setting
        name = await request.app.state.db.get_setting("default_model")
        if not name:
            cur = manager.current_model()
            if cur:
                name = cur.model["id"]
        if not name:
            raise APIError("model is required.", param="model")
    if name == "auto":
        from ..routing import choose_auto
        name = await choose_auto(request.app.state, body, endpoint)
    model = await manager.registry.resolve(name)
    if not model:
        from ..errors import ModelNotFound
        raise ModelNotFound(f"The model '{name}' does not exist. Use GET /v1/models to list available models.")
    return name, model


async def _record(app_state, *, model_id: str | None, requested: str, endpoint: str, principal: Principal,
                  stream: bool, usage: dict | None, timings: dict | None, status: int, error_code: str | None,
                  started: float, load_wait_ms: float, prompt: Any = None, completion: str | None = None) -> None:
    db = app_state.db
    settings = app_state.settings
    log_prompts = bool(await db.get_setting("log_prompts", settings.log_prompts))
    usage = usage or {}
    timings = timings or {}
    tps = timings.get("generation_tps") or timings.get("predicted_per_second")
    ttft = timings.get("ttft_ms") or timings.get("prompt_ms")
    total_ms = (time.time() - started) * 1000
    ctoks = usage.get("completion_tokens") or timings.get("predicted_n")
    ptoks = usage.get("prompt_tokens") or timings.get("prompt_n")
    await db.execute(
        "INSERT INTO inference_requests(created_at, model_id, requested_model, endpoint, api_key_id, stream, prompt_tokens, "
        "completion_tokens, ttft_ms, total_ms, tps, load_wait_ms, status, error_code, prompt, completion) "
        "VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
        (started, model_id, requested, endpoint, principal.id if principal.kind == "api_key" else None, int(stream),
         ptoks, ctoks, ttft, round(total_ms, 1), tps, round(load_wait_ms, 1), status, error_code,
         json.dumps(prompt, ensure_ascii=False)[:20000] if (log_prompts and prompt is not None) else None,
         (completion or "")[:20000] if log_prompts else None))
    if model_id and status < 400:
        await app_state.manager.registry.record_usage(model_id, int(ctoks or 0), tps, ttft)
        app_state.manager.stats["tokens"] += int(ctoks or 0)
    bus.publish("request", {"model_id": model_id, "endpoint": endpoint, "status": status, "tps": tps, "ttft_ms": ttft,
                            "completion_tokens": ctoks, "prompt_tokens": ptoks, "total_ms": round(total_ms), "stream": stream})


def _extract_timings(data: dict) -> dict:
    t = data.get("timings") or {}
    if "predicted_per_second" in t and "generation_tps" not in t:  # llama.cpp
        t = {"generation_tps": t.get("predicted_per_second"), "prompt_tps": t.get("prompt_per_second"),
             "ttft_ms": t.get("prompt_ms"), "prompt_ms": t.get("prompt_ms"), "generation_ms": t.get("predicted_ms"),
             "predicted_n": t.get("predicted_n"), "prompt_n": t.get("prompt_n"), "cached_tokens": t.get("cache_n")}
    return t


async def _proxy(request: Request, endpoint: str, body: dict, principal: Principal, *, kind: str):
    """Common path for chat/completions/embeddings/rerank."""
    app_state = request.app.state
    manager: ModelManager = app_state.manager
    started = time.time()
    requested, model = await _pick_model(request, body, endpoint)
    if kind == "chat" or kind == "completion":
        if model.get("embedding") or model.get("reranker"):
            raise APIError(f"Model '{model['id']}' is an {'embedding' if model.get('embedding') else 'reranking'} model and cannot generate text.",
                           code="WRONG_MODEL_TYPE")
    if kind == "embeddings" and not model.get("embedding"):
        # Allow causal LMs served by MLX to embed (last-token pooling) only if explicitly tagged; otherwise reject
        if model["runtime"] != "mlx":
            raise APIError(f"Model '{model['id']}' is not an embedding model.", code="WRONG_MODEL_TYPE")
    if kind == "rerank" and not model.get("reranker"):
        raise APIError(f"Model '{model['id']}' is not a reranking model.", code="WRONG_MODEL_TYPE")
    stream = bool(body.get("stream")) and kind in ("chat", "completion")
    t_load = time.time()
    try:
        lm = await manager.ensure_loaded(model["id"], reason=f"{endpoint}")
    except APIError as e:
        await _record(app_state, model_id=model["id"], requested=requested, endpoint=endpoint, principal=principal,
                      stream=stream, usage=None, timings=None, status=e.status_code, error_code=e.code, started=started,
                      load_wait_ms=(time.time() - t_load) * 1000)
        raise
    load_wait_ms = (time.time() - t_load) * 1000
    body = dict(body)
    body["model"] = model["id"]
    if stream and kind == "chat":
        so = dict(body.get("stream_options") or {})
        so["include_usage"] = True
        body["stream_options"] = so
    url = f"{lm.handle.base_url}{endpoint}"
    timeout = httpx.Timeout(connect=10.0, read=float(app_state.settings.generation_timeout_seconds), write=60.0, pool=10.0)
    prompt_for_log = body.get("messages") or body.get("prompt") or body.get("input")

    async with manager.use(lm):
        if not stream:
            try:
                r = await manager.client.post(url, json=body, timeout=timeout)
            except httpx.ReadTimeout:
                raise GenerationTimeout("The model did not finish generating in time.")
            except httpx.HTTPError as e:
                if not lm.handle.alive():
                    raise WorkerCrashed(f"The inference worker for '{model['id']}' crashed during the request.")
                raise APIError(f"Worker connection error: {e}", status_code=502, code="WORKER_UNREACHABLE", error_type="runtime_error")
            try:
                data = r.json()
            except Exception:
                data = {"error": {"message": r.text[:500], "type": "runtime_error", "code": "WORKER_BAD_RESPONSE"}}
            if r.status_code >= 400:
                err = data.get("error") if isinstance(data, dict) else None
                code = (err or {}).get("code") if isinstance(err, dict) else None
                msg = (err or {}).get("message") if isinstance(err, dict) else str(err)
                await _record(app_state, model_id=model["id"], requested=requested, endpoint=endpoint, principal=principal,
                              stream=False, usage=None, timings=None, status=r.status_code, error_code=str(code or "WORKER_ERROR"),
                              started=started, load_wait_ms=load_wait_ms, prompt=prompt_for_log)
                # normalise llama.cpp error shape
                return JSONResponse(status_code=r.status_code, content={"error": {"message": msg or "worker error",
                                                                                 "type": (err or {}).get("type", "runtime_error") if isinstance(err, dict) else "runtime_error",
                                                                                 "code": code or "WORKER_ERROR"}})
            data["model"] = requested if requested != "auto" else model["id"]
            timings = _extract_timings(data)
            if timings:
                data["timings"] = timings
            completion = None
            if kind == "chat":
                completion = ((data.get("choices") or [{}])[0].get("message") or {}).get("content")
            elif kind == "completion":
                completion = (data.get("choices") or [{}])[0].get("text")
            await _record(app_state, model_id=model["id"], requested=requested, endpoint=endpoint, principal=principal,
                          stream=False, usage=data.get("usage"), timings=timings, status=200, error_code=None,
                          started=started, load_wait_ms=load_wait_ms, prompt=prompt_for_log, completion=completion)
            return JSONResponse(content=data)

        # ---- streaming ----------------------------------------------------
        req = manager.client.build_request("POST", url, json=body, timeout=timeout)
        try:
            resp = await manager.client.send(req, stream=True)
        except httpx.HTTPError as e:
            raise APIError(f"Worker connection error: {e}", status_code=502, code="WORKER_UNREACHABLE", error_type="runtime_error")
        if resp.status_code >= 400:
            raw = await resp.aread()
            await resp.aclose()
            try:
                data = json.loads(raw)
            except Exception:
                data = {"error": {"message": raw.decode(errors="replace")[:500], "type": "runtime_error", "code": "WORKER_ERROR"}}
            err = data.get("error") or {}
            await _record(app_state, model_id=model["id"], requested=requested, endpoint=endpoint, principal=principal,
                          stream=True, usage=None, timings=None, status=resp.status_code, error_code=str(err.get("code") or "WORKER_ERROR"),
                          started=started, load_wait_ms=load_wait_ms, prompt=prompt_for_log)
            return JSONResponse(status_code=resp.status_code, content={"error": {"message": err.get("message", "worker error"),
                                                                                "type": err.get("type", "runtime_error"),
                                                                                "code": err.get("code", "WORKER_ERROR")}})

        lm.in_flight += 1  # held until the generator finishes

        async def gen():
            usage = None
            timings: dict = {}
            completion_parts: list[str] = []
            status = 200
            error_code = None
            try:
                async for line in resp.aiter_lines():
                    if not line:
                        continue
                    if line.startswith("data: "):
                        payload = line[6:]
                        if payload.strip() == "[DONE]":
                            yield b"data: [DONE]\n\n"
                            continue
                        try:
                            obj = json.loads(payload)
                        except json.JSONDecodeError:
                            yield (line + "\n\n").encode()
                            continue
                        if "error" in obj and "choices" not in obj:
                            status = 500
                            error_code = (obj["error"] or {}).get("code", "GENERATION_FAILED")
                            yield (f"data: {json.dumps(obj)}\n\n").encode()
                            continue
                        obj["model"] = requested if requested != "auto" else model["id"]
                        if obj.get("usage"):
                            usage = obj["usage"]
                        if obj.get("timings"):
                            timings = _extract_timings(obj)
                            obj["timings"] = timings
                        for ch in obj.get("choices") or []:
                            d = ch.get("delta") or {}
                            if d.get("content"):
                                completion_parts.append(d["content"])
                            if ch.get("text"):
                                completion_parts.append(ch["text"])
                        yield (f"data: {json.dumps(obj, ensure_ascii=False)}\n\n").encode()
                    else:
                        yield (line + "\n").encode()
            except (httpx.ReadTimeout, asyncio.TimeoutError):
                status, error_code = 504, "GENERATION_TIMEOUT"
                yield (f"data: {json.dumps({'error': {'message': 'generation timed out', 'type': 'runtime_error', 'code': 'GENERATION_TIMEOUT'}})}\n\n").encode()
            except httpx.HTTPError:
                status, error_code = 502, "WORKER_CRASHED" if not lm.handle.alive() else "WORKER_STREAM_ERROR"
                yield (f"data: {json.dumps({'error': {'message': 'worker stream interrupted', 'type': 'runtime_error', 'code': error_code}})}\n\n").encode()
            finally:
                await resp.aclose()
                lm.in_flight = max(0, lm.in_flight - 1)
                lm.last_used = time.time()
                try:
                    await _record(app_state, model_id=model["id"], requested=requested, endpoint=endpoint, principal=principal,
                                  stream=True, usage=usage, timings=timings, status=status, error_code=error_code,
                                  started=started, load_wait_ms=load_wait_ms, prompt=prompt_for_log,
                                  completion="".join(completion_parts))
                except Exception:
                    log.exception("record failed")

        return StreamingResponse(gen(), media_type="text/event-stream",
                                 headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no", "Connection": "keep-alive"})


# ---------------------------------------------------------------------------


@router.get("/v1/models")
async def list_models(request: Request, principal: Principal = Depends(require_inference)):
    manager = _manager(request)
    models = await manager.registry.list_models()
    aliases = await manager.registry.aliases()
    out = []
    for m in models:
        if not m.get("enabled"):
            continue
        st = manager.status_of(m["id"])
        out.append({
            "id": m["id"], "object": "model", "created": int(m["created_at"]), "owned_by": m.get("provider") or "local",
            "root": m["id"], "parent": None,
            # llm-api extensions (ignored by SDKs)
            "name": m["name"], "family": m["family"], "runtime": m["runtime"], "quantization": m["quantization"],
            "parameter_count": m["parameter_count"], "estimated_ram_gb": m["estimated_ram_gb"],
            "context": m["recommended_context"], "max_context": m["max_context"], "task": m["task"],
            "compatibility": m["compatibility_status"], "status": st, "loaded": st == "ready",
            "capabilities": {"vision": m["vision"], "embedding": m["embedding"], "reranking": m["reranker"],
                             "tools": m["tools"], "thinking": m["thinking"]},
        })
    for alias, mid in aliases.items():
        out.append({"id": alias, "object": "model", "created": 0, "owned_by": "alias", "root": mid, "parent": mid,
                    "alias_of": mid})
    return {"object": "list", "data": out}


@router.get("/v1/models/{model_id:path}")
async def get_model(model_id: str, request: Request, principal: Principal = Depends(require_inference)):
    manager = _manager(request)
    m = await manager.registry.resolve(model_id)
    if not m:
        from ..errors import ModelNotFound
        raise ModelNotFound(f"The model '{model_id}' does not exist.")
    return {"id": m["id"], "object": "model", "created": int(m["created_at"]), "owned_by": m.get("provider") or "local",
            "status": manager.status_of(m["id"])}


@router.post("/v1/chat/completions")
async def chat_completions(request: Request, principal: Principal = Depends(require_inference)):
    body = await _parse_body(request)
    msgs = body.get("messages")
    if not isinstance(msgs, list) or not msgs:
        raise APIError("messages must be a non-empty array.", param="messages")
    for m in msgs:
        if not isinstance(m, dict) or "role" not in m:
            raise APIError("Each message needs a role.", param="messages")
    _validate_sampling(body)
    return await _proxy(request, "/v1/chat/completions", body, principal, kind="chat")


@router.post("/v1/completions")
async def completions(request: Request, principal: Principal = Depends(require_inference)):
    body = await _parse_body(request)
    if "prompt" not in body:
        raise APIError("prompt is required.", param="prompt")
    _validate_sampling(body)
    return await _proxy(request, "/v1/completions", body, principal, kind="completion")


@router.post("/v1/embeddings")
async def embeddings(request: Request, principal: Principal = Depends(require_inference)):
    body = await _parse_body(request)
    if "input" not in body:
        raise APIError("input is required.", param="input")
    inp = body["input"]
    n = len(inp) if isinstance(inp, list) else 1
    if n > 256:
        raise APIError("At most 256 inputs per request.", param="input")
    return await _proxy(request, "/v1/embeddings", body, principal, kind="embeddings")


@router.post("/v1/rerank")
async def rerank(request: Request, principal: Principal = Depends(require_inference)):
    body = await _parse_body(request)
    if not body.get("query") or not isinstance(body.get("documents"), list):
        raise APIError("query and documents are required.")
    if len(body["documents"]) > 200:
        raise APIError("At most 200 documents per request.", param="documents")
    return await _proxy(request, "/v1/rerank", body, principal, kind="rerank")