gen.py

import os
import base64
import random
import httpx
from urllib.parse import quote
from fastapi import APIRouter, Request, HTTPException, Header
from fastapi.responses import Response, JSONResponse, StreamingResponse
import re
from typing import Optional, Any
import json
from helper.assets import (
    save_base64_image,
    cleanup_image,
    is_base64_image,
)
import asyncio
from helper.ratelimit import (
    enforce_rate_limit,
    resolve_rate_limit_identity,
    check_audio_rate_limit,
    check_video_rate_limit,
    check_image_rate_limit,
    MAX_CHAT_PROMPT_BYTES,
    MAX_CHAT_PROMPT_CHARS,
    MAX_GROQ_PROMPT_BYTES,
    MAX_GROQ_PROMPT_CHARS,
    MAX_MEDIA_PROMPT_BYTES,
    MAX_MEDIA_PROMPT_CHARS,
    extract_user_text,
    calculate_messages_size,
    normalize_prompt_value,
    enforce_prompt_size,
    resolve_bound_subject,
    get_usage_snapshot_for_subject,
)
from helper.keywords import *
from uuid import uuid4
from time import time
from typing import Dict, List, Optional, Tuple

router = APIRouter(prefix="/gen")

PKEY = os.getenv("POLLINATIONS_KEY", "")
PKEY2 = os.getenv("POLLINATIONS2_KEY", "")
PKEY3 = os.getenv("POLLINATIONS3_KEY", "")
AIRFORCE_KEY = os.getenv("AIRFORCE")
AIRFORCE_VIDEO_MODEL = "grok-imagine-video"
AIRFORCE_API_URL = "https://api.airforce/v1/images/generations"

valid_ratios = {"3:2", "2:3", "1:1", "", None}
ratios = {"3:2", "2:3", "1:1"}
valid_modes = {"normal", "fun", "", None}
modes = {"normal", "fun"}

MODEL_MAP = {
    "llama-3.1-8b-instant": "Meta Llama 3.1 8B Instant",
    "gpt-4o-mini": "OpenAI GPT 4o Mini",
    "nemotron-3-super": "NVIDIA Nemotron 3 Super",
    "openai/gpt-oss-120b": "OpenAI GPT-OSS 120B",
    "openai/gpt-oss-20b": "OpenAI GPT-OSS 20B",
    "qwen-3-235b-a22b-instruct-2507": "Qwen3 Instruct",
    "llama-3.3-70b-versatile": "Meta Llama 3.3 70B Versatile",
    "meta-llama/llama-4-scout-17b-16e-instruct": "Meta Llama 4 Scout",
}

FALLBACK_MODEL = "meta-llama/llama-4-scout-17b-16e-instruct"
FALLBACK_PROVIDER = "groq"

# Header that API-key authenticated clients send so we know to stream
# thinking tokens back to them.
API_KEY_HEADER = "x-api-key"


# ──────────────────────────────────────────────
# CENTRAL ROUTING LOGIC
# ──────────────────────────────────────────────

def route_chat(
    messages: List[Dict[str, Any]],
    uses_tools: bool = False,
) -> Tuple[str, str]:
    """
    Inspect messages and return (chosen_model, provider).

    This is the single source of truth for model selection.
    No API calls, no side-effects — pure routing logic.
    """
    total_chars, total_bytes = calculate_messages_size(messages)
    prompt_text = extract_user_text(messages)

    long_context   = is_long_context(messages)
    code_present   = contains_code(prompt_text)
    math_heavy     = is_math_heavy(prompt_text)
    structured_task = is_structured_task(prompt_text)
    multi_q        = multiple_questions(prompt_text)
    code_heavy     = is_code_heavy(prompt_text, code_present, long_context)
    has_images     = contains_images(messages)

    score = 0
    if long_context:    score += 3
    if math_heavy:      score += 3
    if structured_task: score += 2
    if code_present:    score += 2
    if multi_q:         score += 1
    for kw in REASONING_KEYWORDS:
        if kw in prompt_text:
            score += 1
    score = min(score, 10)

    # ── multimodal fast-path ──────────────────
    if has_images:
        return "gpt-4o-mini", "navy vision"

    # ── tool-use branch ──────────────────────
    if uses_tools:
        if long_context:
            return "nemotron-3-super", "navy"
        if score >= 6:
            return "nemotron-3-super", "navy"
        if score >= 4:
            return "openai/gpt-oss-120b", "groq"
        return "openai/gpt-oss-20b", "groq"

    # ── code branch ──────────────────────────
    if code_present:
        if code_heavy and score >= 6:
            return "o3-mini", "navy"
        if score >= 4:
            return "llama-3.3-70b-versatile", "groq"

    # ── general reasoning branch ─────────────
    if score >= 6:
        return "sonar", "navy"
    if score >= 4:
        return "meta-llama/llama-4-scout-17b-16e-instruct", "groq"

    # ── default ──────────────────────────────
    chosen_model, provider = "llama-3.1-8b-instant", "groq"

    # Groq context-size guard — promote to navy if too large
    if provider == "groq" and (
        total_chars > MAX_GROQ_PROMPT_CHARS or total_bytes > MAX_GROQ_PROMPT_BYTES
    ):
        return "gpt-4o-mini", "navy"

    return chosen_model, provider


def _log_routing(
    chosen_model: str,
    provider: str,
    messages: List[Dict[str, Any]],
    uses_tools: bool,
) -> None:
    prompt_text = extract_user_text(messages)
    long_context    = is_long_context(messages)
    code_present    = contains_code(prompt_text)
    math_heavy      = is_math_heavy(prompt_text)
    structured_task = is_structured_task(prompt_text)
    multi_q         = multiple_questions(prompt_text)
    has_images      = contains_images(messages)
    print(
        f"\n[ADVANCED ROUTER]\n"
        f"  Uses tools:    {uses_tools}\n"
        f"  Long context:  {long_context}\n"
        f"  Code present:  {code_present}\n"
        f"  Math heavy:    {math_heavy}\n"
        f"  Structured:    {structured_task}\n"
        f"  Multi-question:{multi_q}\n"
        f"  Has images:    {has_images}\n"
        f"  → Selected:    {chosen_model} ({provider})\n"
    )


# ──────────────────────────────────────────────
# CENTRAL HTTP CALL
# ──────────────────────────────────────────────

def _get_provider_url_and_key(provider: str) -> Tuple[str, str]:
    """Return (url, api_key) for the given provider, raising on misconfiguration."""
    if provider == "groq":
        keys = [k.strip() for k in os.getenv("GROQ_KEY", "").split(",") if k.strip()]
        if not keys:
            raise HTTPException(500, "Missing GROQ_KEY(s)")
        return "https://api.groq.com/openai/v1/chat/completions", random.choice(keys)

    if provider == "cerebras":
        keys = [k.strip() for k in os.getenv("CER_KEY", "").split(",") if k.strip()]
        if not keys:
            raise HTTPException(500, "Missing CER_KEY(s)")
        return "https://api.cerebras.ai/v1/chat/completions", random.choice(keys)

    if provider == "navy vision":
        keys = [k.strip() for k in os.getenv("NAVY_KEY", "").split(",") if k.strip()]
        if not keys:
            raise HTTPException(500, "Missing NAVY_KEY(s)")
        return "https://api.navy/v1/chat/completions", random.choice(keys)

    if provider == "navy":
        keys = [k.strip() for k in os.getenv("NAVY_TEXT_ONLY", "").split(",") if k.strip()]
        if not keys:
            raise HTTPException(500, "Missing NAVY_TEXT_ONLY key(s)")
        return "https://api.navy/v1/chat/completions", random.choice(keys)

    raise HTTPException(500, f"Unknown provider: {provider!r}")


async def call_chat_completions(
    messages: List[Dict[str, Any]],
    model: str,
    provider: str,
    extra_body: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """
    Resilient chat-completions call designed to survive Cloudflare 524 timeouts.

    Strategy:
      1. Ask the upstream for a *streaming* response so bytes arrive before
         Cloudflare's ~100 s idle timeout fires.
      2. Accumulate the stream into a single synthetic non-streaming payload
         so callers don't need to change.
      3. Retry up to 2 times (with a short back-off) on 502/503/524.
      4. On exhausted retries fall through to the Groq fallback.
    """
    url, api_key = _get_provider_url_and_key(provider)
    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}

    # Always request streaming upstream — we reassemble below.
    body: Dict[str, Any] = {"model": model, "messages": messages, "stream": True}
    if extra_body:
        body.update(extra_body)
        body["stream"] = True  # force streaming even if caller passed stream=False

    TRANSIENT = {502, 503, 524, 429}
    MAX_ATTEMPTS = 3

    last_exc: Optional[Exception] = None

    for attempt in range(MAX_ATTEMPTS):
        if attempt:
            await asyncio.sleep(2 ** attempt)  # 2 s, 4 s

        try:
            async with httpx.AsyncClient(timeout=httpx.Timeout(300.0, read=300.0)) as client:
                async with client.stream("POST", url, json=body, headers=headers) as r:
                    # Transient upstream error — retry.
                    if r.status_code in TRANSIENT:
                        body_bytes = await r.aread()
                        last_exc = HTTPException(
                            status_code=r.status_code,
                            detail=body_bytes.decode("utf-8", errors="replace")[:500],
                        )
                        print(f"[call_chat_completions] attempt {attempt+1} got {r.status_code}, retrying…")
                        continue

                    if r.status_code != 200:
                        body_bytes = await r.aread()
                        raise HTTPException(
                            status_code=r.status_code,
                            detail=body_bytes.decode("utf-8", errors="replace")[:1000],
                        )

                    # ── Reassemble streaming SSE into a single response object ──
                    accumulated_content = ""
                    accumulated_reasoning = ""
                    tool_calls_map: Dict[int, Dict[str, Any]] = {}
                    usage: Dict[str, Any] = {}
                    finish_reason: Optional[str] = None
                    resp_id = ""
                    resp_model = model

                    async for line in r.aiter_lines():
                        if not line or not line.startswith("data:"):
                            continue
                        raw = line[5:].strip()
                        if raw == "[DONE]":
                            break
                        try:
                            obj = json.loads(raw)
                        except Exception:
                            continue

                        if not isinstance(obj, dict):
                            continue

                        resp_id = resp_id or obj.get("id", "")
                        resp_model = obj.get("model", resp_model)

                        if "usage" in obj and obj["usage"]:
                            usage = obj["usage"]

                        choices = obj.get("choices") or []
                        if not choices:
                            continue

                        choice = choices[0]
                        finish_reason = choice.get("finish_reason") or finish_reason
                        delta = choice.get("delta") or {}

                        # Accumulate text content.
                        dc = delta.get("content")
                        if dc:
                            accumulated_content += dc

                        # Accumulate reasoning / thinking tokens.
                        dr = delta.get("reasoning_content") or delta.get("reasoning")
                        if dr:
                            accumulated_reasoning += dr

                        # Accumulate tool-call argument chunks (streamed as fragments).
                        for tc_delta in (delta.get("tool_calls") or []):
                            idx = tc_delta.get("index", 0)
                            if idx not in tool_calls_map:
                                tool_calls_map[idx] = {
                                    "id": tc_delta.get("id", ""),
                                    "type": tc_delta.get("type", "function"),
                                    "function": {"name": "", "arguments": ""},
                                }
                            existing = tool_calls_map[idx]
                            if tc_delta.get("id"):
                                existing["id"] = tc_delta["id"]
                            fn_delta = tc_delta.get("function") or {}
                            if fn_delta.get("name"):
                                existing["function"]["name"] += fn_delta["name"]
                            if fn_delta.get("arguments"):
                                existing["function"]["arguments"] += fn_delta["arguments"]

            # Reassemble into a standard non-streaming response shape.
            tool_calls_list = [tool_calls_map[i] for i in sorted(tool_calls_map)]

            message: Dict[str, Any] = {"role": "assistant", "content": accumulated_content}
            if accumulated_reasoning:
                message["reasoning_content"] = accumulated_reasoning
            if tool_calls_list:
                message["tool_calls"] = tool_calls_list

            return {
                "id": resp_id,
                "object": "chat.completion",
                "model": resp_model,
                "choices": [
                    {
                        "index": 0,
                        "message": message,
                        "finish_reason": finish_reason or "stop",
                    }
                ],
                "usage": usage,
            }

        except HTTPException:
            raise
        except (httpx.RemoteProtocolError, httpx.ReadError, httpx.ConnectError) as exc:
            last_exc = exc
            print(f"[call_chat_completions] attempt {attempt+1} network error: {exc}, retrying…")
            continue

    # All attempts exhausted — fall back to Groq.
    print(f"[call_chat_completions] all attempts failed ({last_exc}), falling back to Groq")
    fb_url, fb_key = _get_provider_url_and_key(FALLBACK_PROVIDER)
    fb_headers = {"Authorization": f"Bearer {fb_key}", "Content-Type": "application/json"}
    fallback_body = {
        "model": FALLBACK_MODEL,
        "messages": messages,
        "stream": False,
    }
    if extra_body:
        # Forward tools/tool_choice but not stream override.
        for k in ("tools", "tool_choice"):
            if k in extra_body:
                fallback_body[k] = extra_body[k]

    async with httpx.AsyncClient(timeout=httpx.Timeout(120.0)) as client:
        fb_r = await client.post(fb_url, json=fallback_body, headers=fb_headers)

    if fb_r.status_code != 200:
        raise HTTPException(
            status_code=fb_r.status_code,
            detail=f"Primary and fallback both failed. Fallback: {fb_r.text[:500]}",
        )
    return fb_r.json()


def _extract_text_from_response(data: Dict[str, Any]) -> str:
    try:
        return data["choices"][0]["message"]["content"] or ""
    except Exception:
        return ""


def _extract_usage(data: Dict[str, Any]) -> Tuple[int, int]:
    usage = data.get("usage", {})
    input_tok  = usage.get("prompt_tokens")     or usage.get("input_tokens",  0)
    output_tok = usage.get("completion_tokens") or usage.get("output_tokens", 0)
    return input_tok, output_tok


# ──────────────────────────────────────────────
# HELPER: image generation
# ──────────────────────────────────────────────

def is_cinematic_image_prompt(prompt: str) -> bool:
    for kw in CREATIVE_KEYWORDS:
        if kw in prompt.lower():
            return True
    return False


def _is_api_key_request(request: Request) -> bool:
    """
    Return True when the caller authenticated with an API key rather than a
    session cookie / browser auth.  We use this to decide whether to forward
    think-tag / reasoning_content tokens to the client.
    """
    return bool(
        request.headers.get(API_KEY_HEADER)
        or request.headers.get("authorization", "").lower().startswith("bearer ")
    )


def _inject_reasoning_into_chunk(obj: Dict[str, Any]) -> Dict[str, Any]:
    """
    Some navy models return thinking tokens in a non-standard
    ``reasoning_content`` field inside each delta.  When that field is
    present we wrap it in <think>…</think> and prepend it to the regular
    ``content`` delta so that every SSE-speaking client sees a single,
    unified text stream.

    The original ``reasoning_content`` field is preserved so clients that
    know about it can still use it directly.
    """
    try:
        delta = obj["choices"][0]["delta"]
    except (KeyError, IndexError, TypeError):
        return obj

    reasoning = delta.get("reasoning_content") or delta.get("reasoning") or ""
    content   = delta.get("content") or ""

    if reasoning and isinstance(reasoning, str):
        # Wrap in <think> tags and prepend to the visible content delta.
        wrapped = f"<think>{reasoning}</think>"
        delta["content"] = wrapped + content
        # Keep the raw field so native clients can parse it too.
        delta["reasoning_content"] = reasoning
        obj["choices"][0]["delta"] = delta

    return obj


def _normalize_usage_block(obj: Dict[str, Any]) -> Dict[str, Any]:
    """Rewrite the usage block to a canonical shape (in-place, returns obj)."""
    if "usage" not in obj or not isinstance(obj.get("usage"), dict):
        return obj
    u = obj["usage"]
    input_tok  = u.get("prompt_tokens")     or u.get("input_tokens",  0)
    output_tok = u.get("completion_tokens") or u.get("output_tokens", 0)
    obj["usage"] = {
        "prompt_tokens":     input_tok,
        "completion_tokens": output_tok,
        "total_tokens":      input_tok + output_tok,
        "input_tokens":      input_tok,
        "output_tokens":     output_tok,
    }
    return obj


# ──────────────────────────────────────────────
# IMAGE GENERATION
# ──────────────────────────────────────────────

@router.post("/image")
@router.get("/image/{prompt}")
async def generate_image(
    request: Request,
    prompt: str = None,
    authorization: str = Header(None),
    x_client_id: str = Header(None),
):
    timeout = httpx.Timeout(300.0, read=300.0)

    if prompt is None:
        payload = await request.json()
        prompt = payload.get("prompt")
        mode = payload.get("mode")
        image_urls = payload.get("image_urls")
    else:
        mode = request.query_params.get("mode")
        image_urls = request.query_params.getlist("image_urls")

    prompt = normalize_prompt_value(prompt, "prompt")
    enforce_prompt_size(prompt, MAX_MEDIA_PROMPT_CHARS, MAX_MEDIA_PROMPT_BYTES, "Image prompt")
    await check_image_rate_limit(request, authorization, x_client_id)

    chosen_model = "zimage"
    if is_cinematic_image_prompt(prompt):
        chosen_model = "flux"

    if isinstance(mode, str):
        m = mode.strip().lower()
        if m == "fantasy":
            chosen_model = "flux"
        elif m == "realistic":
            chosen_model = "zimage"

    has_input_image = bool(image_urls)
    temp_assets = []

    if has_input_image:
        chosen_model = "klein"

    params = {"model": chosen_model, "key": PKEY2}

    if has_input_image:
        processed = []
        for img in image_urls[:2]:
            if is_base64_image(img):
                image_id = save_base64_image(img)
                temp_assets.append(image_id)
                served = f"{request.base_url}asset-cdn/assets/{image_id}"
                processed.append(served)
            else:
                processed.append(img)
        params["image"] = "|".join(processed)

    encoded_prompt = quote(prompt, safe="")
    query = "&".join(f"{k}={quote(str(v), safe='')}" for k, v in params.items())
    url = f"https://gen.pollinations.ai/image/{encoded_prompt}?{query}"

    try:
        async with httpx.AsyncClient(timeout=timeout) as client:
            resp = await client.get(url)
    finally:
        for aid in temp_assets:
            cleanup_image(aid)

    if resp.status_code != 200:
        raise HTTPException(500, f"Pollinations error: {resp.status_code}")

    return Response(content=resp.content, media_type="image/jpeg")


# ──────────────────────────────────────────────
# SFX GENERATION
# ──────────────────────────────────────────────

@router.get("/sfx/{prompt}")
@router.post("/sfx")
async def gensfx(
    request: Request,
    prompt: str = None,
    authorization: str = Header(None),
    x_client_id: str = Header(None),
):
    if prompt is None:
        payload = await request.json()
        prompt = payload.get("prompt")

    prompt = normalize_prompt_value(prompt, "prompt")
    enforce_prompt_size(prompt, MAX_MEDIA_PROMPT_CHARS, MAX_MEDIA_PROMPT_BYTES, "Audio prompt")
    await check_audio_rate_limit(request, authorization, x_client_id)

    url = f"https://gen.pollinations.ai/audio/{prompt}?model=acestep&key={PKEY}"

    async with httpx.AsyncClient(timeout=None) as client:
        resp = await client.get(url)

    if resp.status_code != 200:
        return JSONResponse(
            status_code=resp.status_code,
            content={"success": False, "error": "Upstream music/sfx generation failed"},
        )

    return Response(resp.content, media_type="audio/mpeg")


# ──────────────────────────────────────────────
# TTS GENERATION
# ──────────────────────────────────────────────

@router.get("/tts/{prompt}")
@router.post("/tts")
async def gentts(
    request: Request,
    prompt: str = None,
    authorization: str = Header(None),
    x_client_id: str = Header(None),
):
    if prompt is None:
        payload = await request.json()
        prompt = payload.get("prompt")

    prompt = normalize_prompt_value(prompt, "prompt")
    enforce_prompt_size(prompt, MAX_MEDIA_PROMPT_CHARS, MAX_MEDIA_PROMPT_BYTES, "Audio prompt")
    await check_audio_rate_limit(request, authorization, x_client_id)

    url = f"https://gen.pollinations.ai/audio/{prompt}?key={PKEY3}"

    async with httpx.AsyncClient(timeout=None) as client:
        resp = await client.get(url)

    if resp.status_code != 200:
        return JSONResponse(
            status_code=resp.status_code,
            content={"success": False, "error": "Upstream audio generation failed"},
        )

    return Response(resp.content, media_type="audio/mpeg")


# ──────────────────────────────────────────────
# VIDEO GENERATION (Pollinations)
# ──────────────────────────────────────────────

@router.get("/video/{prompt}")
@router.post("/video")
@router.head("/video")
async def genvideo(
    request: Request,
    prompt: str = None,
    authorization: str = Header(None),
    x_client_id: str = Header(None),
):
    if request.method == "HEAD":
        return Response(
            status_code=200,
            headers={
                "Y-prompt": "string — required. The text prompt used to generate the video.",
                "Y-ratio": "string — optional. Aspect ratio of the output video.",
                "Y-ratio-values": "3:2,2:3,1:1",
                "Y-ratio-default": "3:2",
                "Y-mode": "string — optional. Controls generation style.",
                "Y-mode-values": "normal,fun",
                "Y-mode-default": "normal",
                "Y-duration": "integer — optional. Duration in seconds (1–10).",
                "Y-duration-default": "5",
                "Y-image_urls": "array<string> — optional. Up to 2 image URLs for conditioning.",
                "Y-image_urls-max": "2",
                "Y-response_format": "video/mp4",
                "Y-model": "grok-video",
            },
        )

    aspectRatio = "3:2"
    inputMode = "normal"
    duration = 5
    image_urls = None

    if prompt is None:
        user_body = await request.json()
        prompt = user_body.get("prompt")
        ratio = user_body.get("ratio")
        mode = user_body.get("mode")
        image_urls = user_body.get("image_urls")
        duration = user_body.get("duration", 5)

        if ratio not in valid_ratios:
            raise HTTPException(400, f"Invalid aspect ratio '{ratio}'. Must be one of 3:2, 2:3, or 1:1.")
        if ratio in ratios:
            aspectRatio = ratio

        if mode not in valid_modes:
            raise HTTPException(400, f"Invalid mode '{mode}'. Must be 'normal' or 'fun'.")
        if mode in modes:
            inputMode = mode

        if image_urls:
            if not isinstance(image_urls, list):
                raise HTTPException(400, "image_urls must be a list")
            if len(image_urls) > 2:
                raise HTTPException(400, "You may provide at most two image URLs")

        try:
            duration = max(1, min(10, int(duration)))
        except (TypeError, ValueError):
            duration = 5

    prompt = normalize_prompt_value(prompt, "prompt")
    enforce_prompt_size(prompt, MAX_MEDIA_PROMPT_CHARS, MAX_MEDIA_PROMPT_BYTES, "Video prompt")
    await check_video_rate_limit(request, authorization, x_client_id)

    RATIO_MAP = {"3:2": "16:9", "2:3": "9:16", "1:1": "9:16"}
    pollinations_ratio = RATIO_MAP.get(aspectRatio, "16:9")

    encoded_prompt = quote(prompt, safe="")
    params = {
        "model": "ltx-2",
        "duration": duration,
        "aspectRatio": pollinations_ratio,
        "seed": -1,
    }

    temp_assets = []

    if image_urls:
        processed_urls = []
        for img in image_urls[:2]:
            if is_base64_image(img):
                image_id = save_base64_image(img)
                temp_assets.append(image_id)
                served_url = f"{request.base_url}asset-cdn/assets/{image_id}"
                processed_urls.append(served_url)
            else:
                processed_urls.append(img)
        params["image"] = "|".join(processed_urls)

    if inputMode == "fun":
        params["enhance"] = "true"

    query_string = "&".join(f"{k}={quote(str(v), safe='')}" for k, v in params.items())
    url = f"https://gen.pollinations.ai/image/{encoded_prompt}?{query_string}&key={PKEY}"
    print(f"[VIDEO GEN] Pollinations URL: {url}")

    resp = None
    try:
        async with httpx.AsyncClient(timeout=600) as client:
            resp = await client.get(url)
    finally:
        for aid in temp_assets:
            cleanup_image(aid)

    if resp is None:
        raise HTTPException(502, "Video generation request failed")

    if resp.status_code != 200:
        body_text = ""
        try:
            body_text = resp.text
        except Exception:
            pass
        return JSONResponse(
            status_code=resp.status_code,
            content={
                "success": False,
                "error": "Upstream video generation failed",
                "status_code": resp.status_code,
                "message": body_text[:1000],
            },
        )

    if not resp.content:
        raise HTTPException(502, "Pollinations returned empty response")

    return Response(
        content=resp.content,
        media_type="video/mp4",
        headers={
            "Content-Length": str(len(resp.content)),
            "Accept-Ranges": "bytes",
        },
    )


# ──────────────────────────────────────────────
# VIDEO GENERATION (Airforce)
# ──────────────────────────────────────────────

@router.get("/video/airforce/{prompt}")
@router.post("/video/airforce")
async def genvideo_airforce(
    request: Request,
    prompt: str = None,
    authorization: str = Header(None),
    x_client_id: str = Header(None),
):
    if request.method == "HEAD":
        return Response(
            status_code=200,
            headers={
                "Y-prompt": "string — required. The text prompt used to generate the video.",
                "Y-ratio": "string — optional. Aspect ratio of the output video.",
                "Y-ratio-values": "3:2,2:3,1:1",
                "Y-ratio-default": "3:2",
                "Y-mode": "string — optional. Controls generation style.",
                "Y-mode-values": "normal,fun",
                "Y-mode-default": "normal",
                "Y-duration": "integer — optional. Duration in seconds.",
                "Y-duration-default": "5",
                "Y-image_urls": "array<string> — optional. Up to 2 image URLs for conditioning.",
                "Y-image_urls-max": "2",
                "Y-response_format": "video/mp4",
                "Y-model": "grok-imagine-video",
            },
        )

    aspectRatio = "3:2"
    inputMode = "normal"
    image_urls = None

    if prompt is None:
        user_body = await request.json()
        prompt = user_body.get("prompt")
        ratio = user_body.get("ratio")
        mode = user_body.get("mode")
        image_urls = user_body.get("image_urls")

        if ratio not in valid_ratios:
            raise HTTPException(400, f"Invalid aspect ratio {ratio}. Must be one of 3:2, 2:3, or 1:1. Default is 3:2")
        if ratio in ratios:
            aspectRatio = ratio

        if mode not in valid_modes:
            raise HTTPException(400, f"Invalid mode {mode}. Must be 'normal' or 'fun'. Default is normal")
        if mode in modes:
            inputMode = mode

        if image_urls:
            if not isinstance(image_urls, list):
                raise HTTPException(400, "image_urls must be a list")
            if len(image_urls) > 2:
                raise HTTPException(400, "You may provide at most two image URLs")

    prompt = normalize_prompt_value(prompt, "prompt")
    enforce_prompt_size(prompt, MAX_MEDIA_PROMPT_CHARS, MAX_MEDIA_PROMPT_BYTES, "Video prompt")
    await check_video_rate_limit(request, authorization, x_client_id)

    payload = {
        "model": AIRFORCE_VIDEO_MODEL,
        "prompt": prompt,
        "n": 1,
        "size": "1024x1024",
        "response_format": "b64_json",
        "sse": False,
        "mode": inputMode,
        "aspectRatio": aspectRatio,
    }

    if image_urls:
        payload["image_urls"] = image_urls

    async with httpx.AsyncClient(timeout=600) as client:
        resp = await client.post(
            AIRFORCE_API_URL,
            headers={"Authorization": f"Bearer {AIRFORCE_KEY}", "Content-Type": "application/json"},
            json=payload,
        )

    if resp.status_code != 200:
        return JSONResponse(status_code=resp.status_code, content=resp.json())

    if not resp.content:
        raise HTTPException(502, "api.airforce returned empty response")

    try:
        result = resp.json()
        b64_video = result["data"][0]["b64_json"]
    except Exception:
        raise HTTPException(502, f"Invalid api.airforce response: {resp.text[:500]}")

    if not b64_video:
        raise HTTPException(502, "Airforce returned empty b64_json")

    video_bytes = base64.b64decode(b64_video)

    return Response(
        content=video_bytes,
        media_type="video/mp4",
        headers={
            "Content-Length": str(len(video_bytes)),
            "Accept-Ranges": "bytes",
        },
    )


# ──────────────────────────────────────────────
# CHAT COMPLETIONS  (/gen/chat/completions)
# ──────────────────────────────────────────────

async def _check_chat_rate_limit(
    request: Request,
    authorization: Optional[str],
    client_id: Optional[str] = None,
):
    return await enforce_rate_limit(request, authorization, "cloudChatDaily", client_id)


@router.post("/chat/completions")
async def generate_text(
    request: Request,
    authorization: Optional[str] = Header(None),
    x_client_id: Optional[str] = Header(None),
):
    body = await request.json()
    messages = body.get("messages", [])
    if not isinstance(messages, list) or len(messages) == 0:
        raise HTTPException(400, "messages[] is required")

    uses_tools = (
        "tools" in body and isinstance(body["tools"], list) and len(body["tools"]) > 0
    ) or ("tool_choice" in body and body["tool_choice"] not in [None, "none"])

    chosen_model, provider = route_chat(messages, uses_tools=uses_tools)
    _log_routing(chosen_model, provider, messages, uses_tools)

    await _check_chat_rate_limit(request, authorization, x_client_id)

    # Determine whether the caller is an API-key client that should receive
    # raw thinking tokens.
    forward_thinking = _is_api_key_request(request)

    body["model"] = chosen_model
    stream = body.get("stream", False)

    url, api_key = _get_provider_url_and_key(provider)
    headers = {"Authorization": f"Bearer {api_key}"}

    if stream:
        body["stream"] = True

        async def stream_fallback(client: httpx.AsyncClient):
            fallback_body = {
                "model": FALLBACK_MODEL,
                "messages": body["messages"],
                "stream": True,
            }
            fb_url, fb_key = _get_provider_url_and_key(FALLBACK_PROVIDER)
            fb_headers = {"Authorization": f"Bearer {fb_key}"}
            print("[FALLBACK] Starting Groq fallback stream")

            async with client.stream("POST", fb_url, json=fallback_body, headers=fb_headers) as r:
                if r.status_code >= 400:
                    err = (await r.aread()).decode("utf-8", errors="replace")
                    yield f'data: {{"error": "Fallback provider failed: {err[:500]}"}}\n\n'
                    return
                async for line in r.aiter_lines():
                    if not line:
                        yield "\n"
                        continue
                    yield (line if line.startswith("data:") else f"data: {line}\n\n") + "\n"

        async def stream_primary(client: httpx.AsyncClient):
            try:
                async with client.stream("POST", url, json=body, headers=headers) as r:
                    if r.status_code >= 400:
                        print("[STREAM FALLBACK] Primary provider failed → switching to fallback")
                        async for chunk in stream_fallback(client):
                            yield chunk
                        return

                    async for line in r.aiter_lines():
                        if not line:
                            yield "\n"
                            continue
                        if line.startswith("data:"):
                            try:
                                obj = json.loads(line[5:].strip())
                                if isinstance(obj, dict) and isinstance(obj.get("error"), dict):
                                    async for chunk in stream_fallback(client):
                                        yield chunk
                                    return
                            except Exception:
                                pass
                        yield line + "\n"
            except Exception as e:
                print(f"[STREAM ERROR] {e}")
                async for chunk in stream_fallback(client):
                    yield chunk

        async def event_generator():
            sent_metadata = False
            async with httpx.AsyncClient(timeout=None) as client:
                async for chunk in stream_primary(client):
                    # ── emit router metadata once as the very first SSE frame ──
                    if not sent_metadata:
                        meta = {
                            "router_metadata": {
                                "model_name": MODEL_MAP.get(chosen_model, chosen_model)
                            }
                        }
                        yield f"data: {json.dumps(meta)}\n\n"
                        sent_metadata = True

                    # ── pass [DONE] straight through ──────────────────────────
                    if "data: [DONE]" in chunk:
                        yield chunk
                        continue

                    # ── process data: … lines ─────────────────────────────────
                    if chunk.startswith("data:"):
                        raw = chunk[5:].strip()
                        try:
                            obj = json.loads(raw)
                        except Exception:
                            # Not valid JSON — forward verbatim (keeps partial
                            # chunks from blocking the stream).
                            yield chunk
                            continue

                        if not isinstance(obj, dict):
                            yield chunk
                            continue

                        # Normalize usage block whenever it appears.
                        _normalize_usage_block(obj)

                        # ── thinking / reasoning tokens ───────────────────────
                        # Navy models may embed thinking in two ways:
                        #
                        #   1. As delta.reasoning_content (separate field)
                        #   2. Inline inside delta.content wrapped in <think>…</think>
                        #
                        # For API-key callers we always surface both forms.
                        # For browser/session callers we strip reasoning_content
                        # so it doesn't confuse UI clients that don't expect it,
                        # but <think> tags already present in content are left
                        # alone (they arrived that way from upstream).
                        if forward_thinking:
                            # Merge reasoning_content into content as
                            # <think>…</think> and keep the raw field.
                            obj = _inject_reasoning_into_chunk(obj)
                        else:
                            # Strip the non-standard field so browser clients
                            # don't see unexpected keys.
                            try:
                                delta = obj["choices"][0]["delta"]
                                delta.pop("reasoning_content", None)
                                delta.pop("reasoning", None)
                                obj["choices"][0]["delta"] = delta
                            except (KeyError, IndexError, TypeError):
                                pass

                        yield f"data: {json.dumps(obj)}\n\n"
                        continue

                    # ── any other line (comments, keep-alives, …) ─────────────
                    yield chunk

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

    # ── non-streaming ─────────────────────────
    async with httpx.AsyncClient(timeout=None) as client:
        r = await client.post(url, json=body, headers=headers)

        # navy-vision fallback
        if provider == "navy vision" and r.status_code >= 400:
            print("[FALLBACK] Navy vision failed — switching to fallback")
            fb_url, fb_key = _get_provider_url_and_key(FALLBACK_PROVIDER)
            fallback_body = dict(body)
            fallback_body["model"] = FALLBACK_MODEL
            r = await client.post(
                fb_url,
                json=fallback_body,
                headers={"Authorization": f"Bearer {fb_key}"},
            )

    content_type = (r.headers.get("content-type") or "").lower()
    if "application/json" in content_type:
        try:
            payload = r.json()
        except Exception:
            payload = {"error": "Upstream returned invalid JSON"}
        else:
            # Normalize usage fields.
            _normalize_usage_block(payload)

            # ── thinking tokens in non-streaming responses ────────────────────
            # Some navy models put thinking content in
            # message.reasoning_content.  For API-key callers we prepend it to
            # message.content wrapped in <think>…</think>; for others we drop
            # the non-standard field.
            try:
                message = payload["choices"][0]["message"]
                reasoning = (
                    message.pop("reasoning_content", None)
                    or message.pop("reasoning", None)
                    or ""
                )
                if reasoning and isinstance(reasoning, str):
                    if forward_thinking:
                        existing = message.get("content") or ""
                        message["content"] = f"<think>{reasoning}</think>{existing}"
                        # Restore the raw field for clients that want it.
                        message["reasoning_content"] = reasoning
                    # else: already popped — nothing to do.
                    payload["choices"][0]["message"] = message
            except (KeyError, IndexError, TypeError):
                pass

            payload.setdefault("router_metadata", {})["model_name"] = MODEL_MAP.get(
                chosen_model, chosen_model
            )
    else:
        payload = {
            "error": "Upstream returned non-JSON response",
            "status_code": r.status_code,
            "message": r.text[:1000],
        }

    return JSONResponse(status_code=r.status_code, content=payload)


# ──────────────────────────────────────────────
# PROMPT ANALYZE  (/gen/prompt_analyze)
# ──────────────────────────────────────────────

@router.post("/prompt_analyze")
async def analyze_prompt(request: Request):
    body = await request.json()
    messages = body.get("prompt", [])
    if not isinstance(messages, list) or len(messages) == 0:
        raise HTTPException(400, "messages[] is required")

    uses_tools = (
        "tools" in body and isinstance(body["tools"], list) and len(body["tools"]) > 0
    ) or ("tool_choice" in body and body["tool_choice"] not in [None, "none"])

    chosen_model, _ = route_chat(messages, uses_tools=uses_tools)
    return {MODEL_MAP.get(chosen_model, chosen_model)}


# ──────────────────────────────────────────────
# MODELS LIST
# ──────────────────────────────────────────────

@router.get("/models")
def return_models_openai():
    return {
        "object": "list",
        "data": [
            {
                "id": "lightning",
                "object": "model",
                "created": 1767225600,
                "owned_by": "inferenceport-ai",
            }
        ],
    }


# ──────────────────────────────────────────────
# RESPONSES API  (/gen/responses)
# ──────────────────────────────────────────────

def _resp_id(prefix: str) -> str:
    return f"{prefix}_{uuid4().hex}"

def _resp_ts() -> int:
    return int(time())

def _content_to_text(content: Any) -> str:
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = []
        for item in content:
            if isinstance(item, dict) and item.get("type") in ("input_text", "output_text", "text"):
                txt = item.get("text")
                if isinstance(txt, str):
                    parts.append(txt)
        return "".join(parts)
    return ""

def _responses_input_to_messages(
    input_data: Any,
    instructions: Optional[str] = None,
) -> List[Dict[str, Any]]:
    messages: List[Dict[str, Any]] = []
    if instructions:
        messages.append({"role": "developer", "content": instructions})

    if isinstance(input_data, str):
        messages.append({"role": "user", "content": input_data})
        return messages

    if isinstance(input_data, list):
        for item in input_data:
            if isinstance(item, str):
                messages.append({"role": "user", "content": item})
                continue
            if not isinstance(item, dict):
                continue
            role = item.get("role", "user")
            text = _content_to_text(item.get("content", ""))
            if text:
                messages.append({"role": role, "content": text})

    return messages

def _build_responses_payload(
    model: str,
    text: str,
    response_id: str,
    input_tokens: int = 0,
    output_tokens: int = 0,
    tool_calls: Optional[List[Dict[str, Any]]] = None,
) -> Dict[str, Any]:
    # Build content: text part first, then one function_call part per tool call
    content: List[Dict[str, Any]] = []
    if text:
        content.append({"type": "output_text", "text": text, "annotations": []})
    for tc in (tool_calls or []):
        fn = tc.get("function", {})
        content.append({
            "type": "tool_use",
            "id": tc.get("id", _resp_id("tool")),
            "name": fn.get("name", ""),
            "input": json.loads(fn["arguments"]) if fn.get("arguments") else {},
        })

    # Top-level output items: one message item (text) + one per tool call
    output_items: List[Dict[str, Any]] = []

    if text or not tool_calls:
        output_items.append({
            "id": _resp_id("msg"),
            "type": "message",
            "role": "assistant",
            "status": "completed",
            "content": [c for c in content if c["type"] == "output_text"],
        })

    for tc in (tool_calls or []):
        fn = tc.get("function", {})
        output_items.append({
            "id": tc.get("id", _resp_id("tool")),
            "type": "function_call",
            "call_id": tc.get("id", ""),
            "name": fn.get("name", ""),
            "arguments": fn.get("arguments", "{}"),
            "status": "completed",
        })

    return {
        "id": response_id,
        "object": "response",
        "created_at": _resp_ts(),
        "status": "completed",
        "completed_at": _resp_ts(),
        "error": None,
        "incomplete_details": None,
        "instructions": None,
        "max_output_tokens": None,
        "model": model,
        "output": output_items,
        "output_text": text,
        "usage": {
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "total_tokens": input_tokens + output_tokens,
        },
    }


@router.post("/responses")
async def create_responses(
    request: Request,
    authorization: Optional[str] = Header(None),
    x_client_id: Optional[str] = Header(None),
):
    body = await request.json()
    model = body.get("model")
    input_data = body.get("input")
    instructions = body.get("instructions")
    stream = body.get("stream", False)
    tools = body.get("tools")
    tool_choice = body.get("tool_choice")

    if not model:
        raise HTTPException(400, "model is required")
    if input_data is None:
        raise HTTPException(400, "input is required")

    messages = _responses_input_to_messages(input_data, instructions=instructions)
    if not messages:
        raise HTTPException(400, "input could not be parsed")

    uses_tools = bool(tools) or (tool_choice not in [None, "none"])

    # Build extra fields to forward upstream
    extra_body: Dict[str, Any] = {}
    if tools:
        extra_body["tools"] = tools
    if tool_choice is not None:
        extra_body["tool_choice"] = tool_choice

    chosen_model, provider = route_chat(messages, uses_tools=uses_tools)
    _log_routing(chosen_model, provider, messages, uses_tools=uses_tools)
    await _check_chat_rate_limit(request, authorization, x_client_id)

    async def _generate() -> Tuple[str, List[Dict[str, Any]], int, int]:
        data = await call_chat_completions(
            messages, chosen_model, provider, extra_body=extra_body or None
        )
        input_tokens, output_tokens = _extract_usage(data)
        message = data.get("choices", [{}])[0].get("message", {})
        text = message.get("content") or ""
        tool_calls = message.get("tool_calls") or []
        return text, tool_calls, input_tokens, output_tokens

    # ── non-streaming ─────────────────────────
    if stream is False:
        text, tool_calls, input_tokens, output_tokens = await _generate()
        response_id = _resp_id("resp")
        return JSONResponse(
            content=_build_responses_payload(
                chosen_model, text, response_id, input_tokens, output_tokens, tool_calls
            )
        )

    # ── streaming ─────────────────────────────
    async def event_stream():
        response_id = _resp_id("resp")
        item_id     = _resp_id("item")
        ts          = _resp_ts()

        def sse(event_type: str, data: dict) -> str:
            """Emit a properly-formed SSE frame with both event: and data: lines.
            The OpenAI SDK dispatches on the `event:` field — without it most
            events are silently dropped."""
            return f"event: {event_type}\ndata: {json.dumps(data)}\n\n"

        # 1. response.created
        yield sse("response.created", {
            "type": "response.created",
            "response": {
                "id": response_id, "object": "response",
                "created_at": ts, "status": "in_progress", "model": model,
                "output": [], "usage": None,
            },
        })

        # 2. response.in_progress
        yield sse("response.in_progress", {
            "type": "response.in_progress",
            "response": {
                "id": response_id, "object": "response",
                "created_at": ts, "status": "in_progress", "model": model,
            },
        })

        # ── Run _generate() in the background, pinging every 15 s ──────────────
        # Without keepalive bytes, Cloudflare (524) and Codex both drop the
        # connection while the model is thinking or accumulating tool arguments.
        # SSE comment lines (": ping") are invisible to application code but
        # reset every proxy's idle-timeout counter.
        PING_INTERVAL = 15  # seconds
        gen_task: asyncio.Task = asyncio.ensure_future(_generate())

        while not gen_task.done():
            try:
                await asyncio.wait_for(asyncio.shield(gen_task), timeout=PING_INTERVAL)
            except asyncio.TimeoutError:
                yield ": ping\n\n"
            except Exception:
                break  # real error — handled below

        try:
            text, tool_calls, input_tokens, output_tokens = gen_task.result()
        except HTTPException as exc:
            yield sse("response.failed", {
                "type": "response.failed",
                "response": {
                    "id": response_id, "object": "response",
                    "created_at": ts, "status": "failed", "model": chosen_model,
                    "error": {"code": "upstream_error", "message": exc.detail},
                },
            })
            yield "data: [DONE]\n\n"
            return
        except Exception as exc:
            yield sse("response.failed", {
                "type": "response.failed",
                "response": {
                    "id": response_id, "object": "response",
                    "created_at": ts, "status": "failed", "model": chosen_model,
                    "error": {"code": "upstream_error", "message": str(exc)},
                },
            })
            yield "data: [DONE]\n\n"
            return

        output_index = 0

        # ── text output (only emitted if there is text content) ──────────────
        if text:
            yield sse("response.output_item.added", {
                "type": "response.output_item.added",
                "response_id": response_id,
                "output_index": output_index,
                "item": {"id": item_id, "type": "message", "role": "assistant",
                         "status": "in_progress", "content": []},
            })
            yield sse("response.content_part.added", {
                "type": "response.content_part.added",
                "response_id": response_id, "item_id": item_id,
                "output_index": output_index, "content_index": 0,
                "part": {"type": "output_text", "text": "", "annotations": []},
            })
            chunk_size = 64
            for i in range(0, len(text), chunk_size):
                yield sse("response.output_text.delta", {
                    "type": "response.output_text.delta",
                    "response_id": response_id, "item_id": item_id,
                    "output_index": output_index, "content_index": 0,
                    "delta": text[i : i + chunk_size],
                })
            yield sse("response.output_text.done", {
                "type": "response.output_text.done",
                "response_id": response_id, "item_id": item_id,
                "output_index": output_index, "content_index": 0,
                "text": text,
            })
            yield sse("response.content_part.done", {
                "type": "response.content_part.done",
                "response_id": response_id, "item_id": item_id,
                "output_index": output_index, "content_index": 0,
                "part": {"type": "output_text", "text": text, "annotations": []},
            })
            yield sse("response.output_item.done", {
                "type": "response.output_item.done",
                "response_id": response_id, "output_index": output_index,
                "item": {"id": item_id, "type": "message", "role": "assistant",
                         "status": "completed",
                         "content": [{"type": "output_text", "text": text, "annotations": []}]},
            })
            output_index += 1

        # ── tool call outputs (one item per call) ─────────────────────────────
        for tc in (tool_calls or []):
            fn = tc.get("function", {})
            tc_id = tc.get("id", _resp_id("tool"))
            tc_item = {
                "id": tc_id,
                "type": "function_call",
                "call_id": tc_id,
                "name": fn.get("name", ""),
                "arguments": fn.get("arguments", "{}"),
                "status": "completed",
            }
            yield sse("response.output_item.added", {
                "type": "response.output_item.added",
                "response_id": response_id,
                "output_index": output_index,
                "item": {**tc_item, "status": "in_progress"},
            })
            yield sse("response.function_call_arguments.delta", {
                "type": "response.function_call_arguments.delta",
                "response_id": response_id, "item_id": tc_id,
                "output_index": output_index, "call_id": tc_id,
                "delta": fn.get("arguments", "{}"),
            })
            yield sse("response.function_call_arguments.done", {
                "type": "response.function_call_arguments.done",
                "response_id": response_id, "item_id": tc_id,
                "output_index": output_index, "call_id": tc_id,
                "arguments": fn.get("arguments", "{}"),
            })
            yield sse("response.output_item.done", {
                "type": "response.output_item.done",
                "response_id": response_id, "output_index": output_index,
                "item": tc_item,
            })
            output_index += 1

        # ── response.completed ────────────────────────────────────────────────
        yield sse("response.completed", {
            "type": "response.completed",
            "response": _build_responses_payload(
                chosen_model, text, response_id, input_tokens, output_tokens, tool_calls
            ),
        })

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

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