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

from config import LLAMA_SWAP_URL
from services.gateway_stream import record_stream_chunk, record_usage
from services.router_logic import LANES, choose_for_lane
from services.routing_policy import load_policy

router = APIRouter(prefix="/v1")

# Virtuelle Lanes, die der Gateway zusätzlich zu den echten Modellen als „Modell" anbietet.
_LANE_LABELS = {"coding": "Coding (Router → coder/heavy/fast)", "chat": "Chat (Router → fast/heavy)"}


@router.get("/models")
async def models():
    async with httpx.AsyncClient(timeout=10) as c:
        r = await c.get(f"{LLAMA_SWAP_URL}/v1/models")
        data = r.json()
    # Lanes ganz oben einblenden, damit IDEs einfach „coding"/„chat" wählen können.
    lanes = [{"id": lane, "object": "model", "owned_by": "mc2-router",
              "description": _LANE_LABELS.get(lane, lane)} for lane in LANES]
    if isinstance(data, dict) and isinstance(data.get("data"), list):
        data["data"] = lanes + data["data"]
    return JSONResponse(data, status_code=r.status_code)


async def _proxy(path: str, request: Request):
    body = await request.json()
    requested = str(body.get("model") or "auto")
    if requested.lower() in ("auto", "chat", "coding"):
        lane = requested.lower()
        alias, reason = choose_for_lane(lane, body)
        body["model"] = alias
        routed = {"x-mc-routed-to": alias, "x-mc-route-reason": reason, "x-mc-lane": lane}
    else:
        alias = requested
        routed = {"x-mc-routed-to": requested}
    # fast-Spur: Thinking aus für flotte Antworten (sofern Client es nicht selbst setzt).
    pol = load_policy()
    if pol["fast_no_think"] and alias == pol["fast"] and "chat_template_kwargs" not in body:
        body["chat_template_kwargs"] = {"enable_thinking": False}
    url = f"{LLAMA_SWAP_URL}{path}"

    if body.get("stream"):
        async def gen():
            async with httpx.AsyncClient(timeout=None) as c:
                async with c.stream("POST", url, json=body) as r:
                    async for chunk in r.aiter_raw():
                        record_stream_chunk(chunk, alias)
                        yield chunk
        return StreamingResponse(gen(), media_type="text/event-stream", headers=routed)

    async with httpx.AsyncClient(timeout=600) as c:
        r = await c.post(url, json=body)
        resp_json = r.json()
        record_usage(resp_json.get("usage") if isinstance(resp_json, dict) else None, alias)
        return JSONResponse(resp_json, status_code=r.status_code, headers=routed)


@router.post("/chat/completions")
async def chat_completions(request: Request):
    return await _proxy("/v1/chat/completions", request)


@router.post("/completions")
async def completions(request: Request):
    return await _proxy("/v1/completions", request)