app.py

import os
import time
import hashlib
from fastapi import FastAPI, Request, HTTPException, status, Header
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import Response, JSONResponse, StreamingResponse, RedirectResponse
import httpx
from bs4 import BeautifulSoup
from typing import List, Dict, Any
import asyncio
import re
from random import randint
from urllib.parse import quote
import uuid
import base64
from subscriptions import fetch_subscription
from typing import Optional

WAN_SPACE = "https://huggingface.co/spaces/Wan-AI/Wan2.1"
WAN_API_BASE = f"{WAN_SPACE}/gradio_api"

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["GET", "POST", "HEAD"],
    allow_headers=["*"],
)

@app.get("/")
async def reroute_to_status():
    return RedirectResponse(url="https://inference.js.org", status_code=status.HTTP_308_PERMANENT_REDIRECT)

OLLAMA_LIBRARY_URL = "https://ollama.com/library"

PLAN_ORDER = ["free", "light", "pro", "creator", "professional"]
TIER_CONFIG = {
    "free": {
        "name": "Free Tier",
        "url": "",
        "price": "0.00",
        "limits": {
            "cloudChatDaily": 50,
            "imagesDaily": 10,
            "videosDaily": 3,
            "audioWeekly": 1,
        },
    },
    "light": {
        "name": "InferencePort AI Light",
        "url": "https://buy.stripe.com/test_6oUcN5g665rp7nLgaq8bS00",
        "price": "9.99",
        "limits": {
            "cloudChatDaily": None,
            "imagesDaily": 50,
            "videosDaily": 10,
            "audioWeekly": 5,
        },
    },
    "pro": {
        "name": "InferencePort AI Pro",
        "url": "https://buy.stripe.com/test_bJe9AT2fg6vt23rgaq8bS01",
        "price": "15.99",
        "limits": {
            "cloudChatDaily": None,
            "imagesDaily": 150,
            "videosDaily": 30,
            "audioWeekly": 25,
        },
    },
    "creator": {
        "name": "InferencePort AI Creator",
        "url": "https://buy.stripe.com/test_14AaEX9HIdXV8rPf6m8bS02",
        "price": "29.99",
        "limits": {
            "cloudChatDaily": None,
            "imagesDaily": 300,
            "videosDaily": 50,
            "audioWeekly": 45,
        },
    },
    "professional": {
        "name": "InferencePort AI Professional",
        "url": "https://buy.stripe.com/test_5kQ00jf22cTR0ZncYe8bS03",
        "price": "99.99",
        "limits": {
            "cloudChatDaily": None,
            "imagesDaily": None,
            "videosDaily": None,
            "audioWeekly": 75,
        },
    },
}
USAGE_PERIODS = {
    "cloudChatDaily": "daily",
    "imagesDaily": "daily",
    "videosDaily": "daily",
    "audioWeekly": "weekly",
}
usage_store = {
    "cloudChatDaily": {},
    "imagesDaily": {},
    "videosDaily": {},
    "audioWeekly": {},
}
usage_locks = {
    "cloudChatDaily": {},
    "imagesDaily": {},
    "videosDaily": {},
    "audioWeekly": {},
}
IDENTITY_CACHE_TTL_SECONDS = 60
identity_cache = {}
CLIENT_BIND_TTL_SECONDS = int(os.getenv("CLIENT_BIND_TTL_SECONDS", str(8 * 24 * 60 * 60)))
MAX_CLIENT_ID_LENGTH = 128
client_subject_bindings = {}

MAX_CHAT_PROMPT_CHARS = int(os.getenv("MAX_CHAT_PROMPT_CHARS", "120000"))
MAX_CHAT_PROMPT_BYTES = int(os.getenv("MAX_CHAT_PROMPT_BYTES", "500000"))
MAX_GROQ_PROMPT_CHARS = int(os.getenv("MAX_GROQ_PROMPT_CHARS", "90000"))
MAX_GROQ_PROMPT_BYTES = int(os.getenv("MAX_GROQ_PROMPT_BYTES", "350000"))
MAX_MEDIA_PROMPT_CHARS = int(os.getenv("MAX_MEDIA_PROMPT_CHARS", "4000"))
MAX_MEDIA_PROMPT_BYTES = int(os.getenv("MAX_MEDIA_PROMPT_BYTES", "16000"))

REASONING_KEYWORDS = [
    # explicit reasoning requests
    "prove", "demonstrate", "derive", "justify", "verify",
    "show that", "walk through", "step by step", "reason through",
    "chain of reasoning", "rigorous", "formal proof",

    # analysis/comparison
    "analyze", "analysis of", "compare and contrast",
    "evaluate", "critically assess", "explain why",
    "explain how", "what causes", "implications of",

    # problem solving
    "solve", "solution to", "how would you approach",
    "strategy for", "optimize", "algorithm for",

    # technical domains
    "theorem", "lemma", "corollary",
    "complexity analysis", "big o", "time complexity",
    "mathematical", "statistical", "probabilistic",
    "model the", "simulate",
]

CODE_KEYWORDS = [
    "await", "async", "print(", "console.log(",
    "code", ".ts", ".js", ".py", ".repy", ".rb",
    "gnu", "gcc", "clang", "clang++", "program",
    "coding"
]

CREATIVE_KEYWORDS = [
    # cinematic cues
    "cinematic", "film still", "movie scene",
    "epic", "dramatic lighting", "moody lighting",
    "volumetric lighting", "depth of field",
    "anamorphic lens", "8k", "4k",

    # art styles
    "concept art", "digital painting",
    "fantasy art", "sci-fi", "mythical",
    "cyberpunk", "steampunk",
    "baroque", "surreal", "abstract",
    "oil painting", "watercolor",

    # rendering engines
    "octane render", "unreal engine",
    "ray tracing", "global illumination",

    # emotional narrative framing
    "emotional portrait", "story scene",
    "hero shot", "dramatic pose",
]

STRUCTURED_KEYWORDS = [
    "return as json",
    "output json",
    "json schema",
    "format as json",
    "structured output",
    "extract entities",
    "extract fields",
    "parse this",
    "convert to table",
    "create a table",
    "categorize into",
    "classify",
    "label the following",
    "taxonomy",
    "generate schema",
]

MATH_PATTERNS = [
    r"\b∫\b", r"\b∑\b", r"\b∂\b",
    r"\bmatrix\b",
    r"\blimit\b",
    r"\bintegral\b",
    r"\bderivative\b",
    r"\bdifferential equation\b",
    r"\blinear algebra\b",
    r"\boptimi[sz]e\b",
    r"\bgradient\b",
    r"\bbackprop\b",
    r"\bproof\b",
    r"\btheorem\b",
]

LIGHTWEIGHT_KEYWORDS = [
    "hello", "hi", "hey",
    "thanks", "thank you",
    "define", "definition of",
    "what is", "who is",
    "quick question",
    "short answer",
    "brief explanation",
    "summarize",
    "paraphrase",
    "rewrite this",
]

def is_long_context(messages: list) -> bool:
    total_chars = sum(len(m.get("content", "")) for m in messages)
    return total_chars > 4000


def contains_code(prompt: str) -> bool:
    if "```" in prompt:
        return True
    for kw in CODE_KEYWORDS:
        if kw in prompt:
            return True
    return False

def is_code_heavy(prompt: str, code_present: bool, long_context: bool) -> bool:
    """
    Determines whether the coding task is substantial enough
    to require a code-optimized or larger model.
    """

    if not code_present:
        return False

    heavy_patterns = [
        r"\brefactor\b",
        r"\boptimi[sz]e\b",
        r"\bdebug\b",
        r"\bfix this\b",
        r"\barchitecture\b",
        r"\bdesign pattern\b",
        r"\bscalable\b",
        r"\bmicroservice\b",
        r"\bmultiple files\b",
        r"\bentire project\b",
        r"\bcodebase\b",
        r"\bperformance\b",
    ]

    for pattern in heavy_patterns:
        if re.search(pattern, prompt):
            return True

    if prompt.count("```") >= 2:
        return True

    if long_context:
        return True

    return False

def is_math_heavy(prompt: str) -> bool:
    for pattern in MATH_PATTERNS:
        if re.search(pattern, prompt):
            return True
    return False


def is_structured_task(prompt: str) -> bool:
    for kw in STRUCTURED_KEYWORDS:
        if kw in prompt:
            return True
    return False


def multiple_questions(prompt: str) -> bool:
    return prompt.count("?") >= 3

def extract_user_text(messages: list) -> str:
    return " ".join(
        message_content_to_text(m.get("content"))
        for m in messages
        if isinstance(m, dict) and m.get("role") == "user"
    ).lower()

def normalize_plan_key(plan_name: Optional[str]) -> str:
    if not plan_name:
        return "free"
    normalized = "".join(ch for ch in str(plan_name).lower() if ch.isalpha())
    if "professional" in normalized:
        return "professional"
    if "creator" in normalized:
        return "creator"
    if "pro" in normalized:
        return "pro"
    if "light" in normalized:
        return "light"
    return "free"


def get_usage_period_key(metric: str) -> str:
    now = time.gmtime()
    period = USAGE_PERIODS.get(metric, "daily")
    if period == "weekly":
        iso_year, iso_week, _ = time.strftime("%G %V %u", now).split(" ")
        return f"{iso_year}-W{iso_week}"
    return time.strftime("%Y-%m-%d", now)


def sanitize_client_id(raw_client_id: Optional[str]) -> Optional[str]:
    if not isinstance(raw_client_id, str):
        return None
    trimmed = raw_client_id.strip()
    if not trimmed or len(trimmed) > MAX_CLIENT_ID_LENGTH:
        return None
    if not re.match(r"^[A-Za-z0-9._:-]+$", trimmed):
        return None
    return trimmed


def get_usage_lock(metric: str, subject: str) -> asyncio.Lock:
    metric_locks = usage_locks.get(metric)
    if metric_locks is None:
        metric_locks = {}
        usage_locks[metric] = metric_locks
    lock = metric_locks.get(subject)
    if lock is None:
        lock = asyncio.Lock()
        metric_locks[subject] = lock
    return lock


def build_default_subject(request: Request, client_id: Optional[str]) -> str:
    if client_id:
        client_hash = hashlib.sha256(client_id.encode("utf-8")).hexdigest()[:24]
        return f"client:{client_hash}"
    host = request.client.host if request.client else "unknown"
    user_agent = request.headers.get("user-agent", "")
    ua_hash = (
        hashlib.sha256(user_agent.encode("utf-8")).hexdigest()[:12]
        if user_agent
        else "noua"
    )
    return f"anon:{host}:{ua_hash}"


def bind_client_subject(client_id: Optional[str], subject: str, plan_key: str):
    if not client_id:
        return
    client_subject_bindings[client_id] = {
        "subject": subject,
        "plan_key": plan_key,
        "expires_at": time.time() + CLIENT_BIND_TTL_SECONDS,
    }


def resolve_bound_subject(client_id: Optional[str], fallback_subject: str) -> str:
    if not client_id:
        return fallback_subject
    bound = client_subject_bindings.get(client_id)
    if not bound:
        return fallback_subject
    if bound.get("expires_at", 0) <= time.time():
        client_subject_bindings.pop(client_id, None)
        return fallback_subject
    return bound.get("subject", fallback_subject)


def normalize_prompt_value(prompt: Optional[str], field_name: str = "prompt") -> str:
    if not isinstance(prompt, str):
        raise HTTPException(status_code=400, detail=f"{field_name} is required")
    normalized = prompt.strip()
    if not normalized:
        raise HTTPException(status_code=400, detail=f"{field_name} is required")
    return normalized


def enforce_prompt_size(prompt: str, max_chars: int, max_bytes: int, context: str):
    char_len = len(prompt)
    byte_len = len(prompt.encode("utf-8"))
    if char_len > max_chars or byte_len > max_bytes:
        raise HTTPException(
            status_code=413,
            detail=(
                f"{context} is too large ({char_len} chars, {byte_len} bytes). "
                f"Max allowed is {max_chars} chars or {max_bytes} bytes."
            ),
        )


def message_content_to_text(content: Any) -> str:
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts: List[str] = []
        for item in content:
            if isinstance(item, str):
                parts.append(item)
                continue
            if isinstance(item, dict):
                text = item.get("text")
                if isinstance(text, str):
                    parts.append(text)
        return " ".join(parts)
    return ""


def calculate_messages_size(messages: list) -> tuple[int, int]:
    total_chars = 0
    total_bytes = 0
    for message in messages:
        if not isinstance(message, dict):
            continue
        text = message_content_to_text(message.get("content"))
        if not text:
            continue
        total_chars += len(text)
        total_bytes += len(text.encode("utf-8"))
    return total_chars, total_bytes


def get_usage_snapshot_for_subject(plan_key: str, subject: str) -> Dict[str, Dict[str, Any]]:
    plan = TIER_CONFIG.get(plan_key) or TIER_CONFIG["free"]
    plan_limits = plan.get("limits", {})
    snapshot: Dict[str, Dict[str, Any]] = {}
    for metric in usage_store.keys():
        limit = plan_limits.get(metric)
        window_key = get_usage_period_key(metric)
        entry = usage_store[metric].get(subject)
        used = 0
        if entry and entry.get("window") == window_key:
            used = max(0, int(entry.get("count", 0)))
        remaining = None if limit is None else max(0, int(limit) - used)
        snapshot[metric] = {
            "limit": limit,
            "used": used,
            "remaining": remaining,
            "window": window_key,
            "period": USAGE_PERIODS.get(metric, "daily"),
        }
    return snapshot


async def resolve_rate_limit_identity(
    request: Request,
    authorization: Optional[str],
    client_id: Optional[str] = None,
) -> tuple[str, str]:
    now = time.time()
    normalized_client_id = sanitize_client_id(client_id)
    default_subject = build_default_subject(request, normalized_client_id)
    if not authorization or not authorization.startswith("Bearer "):
        return "free", resolve_bound_subject(normalized_client_id, default_subject)

    token = authorization.split(" ", 1)[1].strip()
    if not token:
        return "free", resolve_bound_subject(normalized_client_id, default_subject)

    cached = identity_cache.get(token)
    if cached and cached.get("expires_at", 0) > now:
        plan_key = cached.get("plan_key", "free")
        subject = cached.get("subject", default_subject)
        bind_client_subject(normalized_client_id, subject, plan_key)
        return plan_key, subject

    try:
        sub = await fetch_subscription(token)
    except Exception:
        return "free", resolve_bound_subject(normalized_client_id, default_subject)

    if not isinstance(sub, dict) or sub.get("error"):
        return "free", resolve_bound_subject(normalized_client_id, default_subject)

    email = sub.get("email")
    if isinstance(email, str) and email.strip():
        subject = f"user:{email.strip().lower()}"
    else:
        subject = default_subject

    plan_key = normalize_plan_key(sub.get("plan_key"))
    identity_cache[token] = {
        "plan_key": plan_key,
        "subject": subject,
        "expires_at": now + IDENTITY_CACHE_TTL_SECONDS,
    }
    bind_client_subject(normalized_client_id, subject, plan_key)
    return plan_key, subject


async def enforce_rate_limit(
    request: Request,
    authorization: Optional[str],
    metric: str,
    client_id: Optional[str] = None,
) -> Dict[str, Optional[int | str]]:
    if metric not in usage_store:
        raise HTTPException(status_code=500, detail=f"Unknown limit metric: {metric}")

    plan_key, subject = await resolve_rate_limit_identity(request, authorization, client_id)
    plan = TIER_CONFIG.get(plan_key) or TIER_CONFIG["free"]
    plan_limits = plan.get("limits", {})
    limit = plan_limits.get(metric)

    window_key = get_usage_period_key(metric)
    lock = get_usage_lock(metric, subject)
    async with lock:
        bucket = usage_store[metric]
        entry = bucket.get(subject)
        if not entry or entry.get("window") != window_key:
            entry = {"window": window_key, "count": 0}
            bucket[subject] = entry

        if limit is not None and entry["count"] >= int(limit):
            raise HTTPException(
                status_code=429,
                detail=f"{metric} limit reached for {plan.get('name', 'current plan')}",
            )

        entry["count"] += 1
        remaining = None if limit is None else max(0, int(limit) - entry["count"])
        return {
            "plan_key": plan_key,
            "remaining": remaining,
            "used": entry["count"],
            "window": window_key,
        }


async def check_audio_rate_limit(
    request: Request,
    authorization: Optional[str],
    client_id: Optional[str] = None,
):
    await enforce_rate_limit(request, authorization, "audioWeekly", client_id)


def is_complex_reasoning(prompt: str) -> bool:
    if len(prompt) > 800:
        return True

    for kw in REASONING_KEYWORDS:
        if kw in prompt:
            return True

    if re.search(r"\b(if|therefore|assume|let x|given that)\b", prompt):
        return True

    return False


def is_lightweight(prompt: str) -> bool:
    if len(prompt) < 100:
        for kw in LIGHTWEIGHT_KEYWORDS:
            if kw in prompt:
                return True
    return False


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

async def check_image_rate_limit(
    request: Request,
    authorization: Optional[str],
    client_id: Optional[str] = None,
):
    await enforce_rate_limit(request, authorization, "imagesDaily", client_id)


async def check_video_rate_limit(
    request: Request,
    authorization: Optional[str],
    client_id: Optional[str] = None,
):
    await enforce_rate_limit(request, authorization, "videosDaily", client_id)

PKEY  = os.getenv("POLLINATIONS_KEY", "")
PKEY2 = os.getenv("POLLINATIONS2_KEY", "")
PKEY3 = os.getenv("POLLINATIONS3_KEY", "")

GROQ_TOOL_MODELS = [
    "openai/gpt-oss-120b",
    "openai/gpt-oss-20b",
    "meta-llama/llama-4-scout-17b-16e-instruct",
    "qwen/qwen3-32b",
    "moonshotai/kimi-k2-instruct",
]

GROQ_NORMAL_MODELS = [
    "llama-3.1-8b-instant",
    "llama-3.3-70b-versatile",
    "meta-llama/llama-4-maverick-17b-128e-instruct",
    "meta-llama/llama-guard-4-12b",
    "openai/gpt-oss-safeguard-20b",
    "qwen/qwen3-32b",
]

CEREBRAS_MODELS = [
    "gpt-oss-120b",
    "llama3.1-8b",
    "qwen-3-235b-a22b-instruct-2507",
    "zai-glm-4.7",
]

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)

@app.head("/status/sfx")
async def head_sfx():
    return Response(
        status_code=200,
        headers={
            "Content-Type": "audio/mpeg",
            "Accept-Ranges": "bytes",
        }
    )

@app.head("/status/image")
async def head_image():
    return Response(
        status_code=200,
        headers={
            "Content-Type": "image/jpeg",
            "Accept-Ranges": "bytes",
        }
    )

@app.head("/status/video")
async def head_video():
    return Response(
        status_code=200,
        headers={
            "Content-Type": "video/mp4",
            "Accept-Ranges": "bytes",
        }
    )

@app.head("/status/text")
async def head_text():
    return Response(
        status_code=200,
        headers={
            "Content-Type": "application/json",
            "Accept-Ranges": "bytes",
        },
    )

@app.get("/status")
async def get_status():
    notify = ""
    services = {
        "Video Generation": {
            "code": 200,
            "state": "ok",
            "message": "Running normally"
        },
        "Image Generation": {
            "code": 200,
            "state": "ok",
            "message": "Running normally"
        },
        "Lightning-Text v2": {
            "code": 200,
            "state": "ok",
            "message": "Running normally"
        },
        "Music/SFX Generation": {
            "code": 200,
            "state": "ok",
            "message": "Running normally"
        }
    }

    overall_state = (
        "ok" if all(s["state"] == "ok" for s in services.values())
        else "degraded"
    )

    return JSONResponse(
        status_code=200,
        content={
            "state": overall_state,
            "services": services,
            "notifications": notify,
            "latest": "2.4.0"
        }
    )

@app.post("/gen/image")
@app.get("/genimg/{prompt}")
async def generate_image(
    request: Request,
    prompt: str = None,
    authorization: Optional[str] = Header(None),
    x_client_id: Optional[str] = Header(None),
):
    timeout = httpx.Timeout(300.0, read=300.0)
    payload: Dict[str, Any] = {}
    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, "Image prompt")

    await check_image_rate_limit(request, authorization, x_client_id)

    mode = payload.get("mode") if isinstance(payload, dict) else None
    if is_cinematic_image_prompt(prompt):
        chosen_model = "flux"
    else:
        chosen_model = "zimage"

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

    print(f"[IMAGE GEN] Routing to model: {chosen_model}")

    url = f"https://gen.pollinations.ai/image/{quote(prompt, safe='')}?model={chosen_model}&key={PKEY2}"
    async with httpx.AsyncClient(timeout = timeout) as client:
        response = await client.get(url)

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

    return Response(
        content=response.content,
        media_type="image/jpeg"
    )
@app.head("/models")
@app.get("/models")
async def get_models() -> List[Dict]:
    async with httpx.AsyncClient() as client:
        response = await client.get(OLLAMA_LIBRARY_URL)
        html = response.text

    soup = BeautifulSoup(html, "html.parser")
    items = soup.select("li[x-test-model]")

    models = []
    for item in items:
        name = item.select_one("[x-test-model-title] span")
        description = item.select_one("p.max-w-lg")
        sizes = [el.get_text(strip=True) for el in item.select("[x-test-size]")]
        pulls = item.select_one("[x-test-pull-count]")
        tags = [t.get_text(strip=True) for t in item.select('span[class*="text-blue-600"]')]
        updated = item.select_one("[x-test-updated]")
        link = item.select_one("a")

        models.append({
            "name": name.get_text(strip=True) if name else "",
            "description": description.get_text(strip=True) if description else "No description",
            "sizes": sizes,
            "pulls": pulls.get_text(strip=True) if pulls else "Unknown",
            "tags": tags,
            "updated": updated.get_text(strip=True) if updated else "Unknown",
            "link": link.get("href") if link else None,
        })

    return models

@app.post("/gen/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")

    total_chars, total_bytes = calculate_messages_size(messages)
    if total_chars > MAX_CHAT_PROMPT_CHARS or total_bytes > MAX_CHAT_PROMPT_BYTES:
        raise HTTPException(
            status_code=413,
            detail=(
                f"Prompt context too large ({total_chars} chars, {total_bytes} bytes). "
                f"Max allowed is {MAX_CHAT_PROMPT_CHARS} chars or {MAX_CHAT_PROMPT_BYTES} bytes."
            ),
        )

    prompt_text = extract_user_text(messages)
    
    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"])
    
    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)

    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
    
    chosen_model = "llama-3.3-70b-versatile"
    provider = "groq"
    if score > 10:
        score = 10
    if uses_tools:
        if score >= 4:
            chosen_model = "openai/gpt-oss-120b"
        else:
            chosen_model = "openai/gpt-oss-20b"
        provider = "groq"

    elif code_present:
    
        if code_heavy and score >= 6:
            chosen_model = "gpt-oss-120b"
            provider = "cerebras"
    
        elif score >= 4:
            chosen_model = "llama-3.3-70b-versatile"
            provider = "groq"

    elif score >= 4:
        chosen_model = "meta-llama/llama-4-scout-17b-16e-instruct"
        provider = "groq"

    if provider == "groq" and (
        total_chars > MAX_GROQ_PROMPT_CHARS or total_bytes > MAX_GROQ_PROMPT_BYTES
    ):
        raise HTTPException(
            status_code=413,
            detail=(
                f"Prompt exceeds Groq-safe size ({total_chars} chars, {total_bytes} bytes). "
                f"Max Groq-safe size is {MAX_GROQ_PROMPT_CHARS} chars or {MAX_GROQ_PROMPT_BYTES} bytes."
            ),
        )

    await check_chat_rate_limit(request, authorization, x_client_id)

    body["model"] = chosen_model
    print(f"""
    [ADVANCED ROUTER]
      Score: {score}
      Uses tools: {uses_tools}
      Long context: {long_context}
      Code present: {code_present}
      Math heavy: {math_heavy}
      Structured: {structured_task}
      Multi-question: {multi_q}
      → Selected: {chosen_model} ({provider})
    """)


    stream = body.get("stream", False)

    if provider == "groq":
        num = randint(1, 2)
        if num == 1:
            API_KEY = os.getenv("GROQ_KEY", "")
        elif num == 2:
            API_KEY = os.getenv("GROQ2_KEY", "")
        if not API_KEY:
            raise HTTPException(500, "Missing GROQ_KEY")

        url = "https://api.groq.com/openai/v1/chat/completions"

    elif provider == "cerebras":
        API_KEY = os.getenv("CER_KEY", "")
        if not API_KEY:
            raise HTTPException(500, "Missing CER_KEY")

        url = "https://api.cerebras.ai/v1/chat/completions"
    else:
        raise HTTPException(500, "Unknown provider routing error")

    headers = {"Authorization": f"Bearer {API_KEY}"}

    if stream:
        body["stream"] = True
    
        async def event_generator():
            try:
                async with httpx.AsyncClient(timeout=None) as client:
                    async with client.stream(
                        "POST",
                        url,
                        json=body,
                        headers=headers,
                    ) as r:
                        if r.status_code >= 400:
                            error_payload = ""
                            try:
                                error_payload = (
                                    (await r.aread()).decode("utf-8", errors="replace")
                                )[:800]
                            except Exception:
                                error_payload = ""
                            safe_error_payload = (
                                error_payload.replace("\\", "\\\\")
                                .replace('"', '\\"')
                                .replace("\n", " ")
                                .replace("\r", " ")
                            )
                            yield (
                                "data: {\"error\": "
                                f"\"Upstream provider error ({r.status_code}): {safe_error_payload}\""
                                "}\n\n"
                            )
                            return
    
                        async for line in r.aiter_lines():
                            if line == "":
                                yield "\n"
                                continue
    
                            yield line + "\n"
    
            except asyncio.CancelledError:
                return
            except Exception as e:
                yield f"data: {{\"error\": \"{str(e)}\"}}\n\n"
    
        return StreamingResponse(
            event_generator(),
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "X-Accel-Buffering": "no",  # critical for nginx
            },
        )
    else:
        async with httpx.AsyncClient(timeout=None) as client:
            r = await client.post(url, json=body, headers=headers)
        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:
            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
        )

    raise HTTPException(500, "Unknown provider routing error")

@app.get("/gen/sfx/{prompt}")
@app.post("/gen/sfx")
async def gensfx(
    request: Request,
    prompt: str = None,
    authorization: Optional[str] = Header(None),
    x_client_id: Optional[str] = Header(None),
):
    payload: Dict[str, Any] = {}
    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=elevenmusic&key={PKEY}"
    async with httpx.AsyncClient(timeout=None) as client:
        response = await client.get(url)
    body_text = ""
    try:
        body_text = response.text
    except Exception:
        pass
    if response.status_code != 200:
        return JSONResponse(
            status_code=response.status_code,
            content={
                "success": False,
                "error": "Upstream music/sfx generation failed",
                "status_code": response.status_code,
                "message": body_text[:1000]
            }
        )
    return Response(
        response.content,
        media_type="audio/mpeg"
    )

@app.get("/gen/tts/{prompt}")
@app.post("/gen/tts")
async def gensfx(
    request: Request,
    prompt: str = None,
    authorization: Optional[str] = Header(None),
    x_client_id: Optional[str] = Header(None),
):
    payload: Dict[str, Any] = {}
    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:
        response = await client.get(url)
    body_text = ""
    try:
        body_text = response.text
    except Exception:
        pass
    if response.status_code != 200:
        return JSONResponse(
            status_code=response.status_code,
            content={
                "success": False,
                "error": "Upstream audio generation failed",
                "status_code": response.status_code,
                "message": body_text[:1000]
            }
        )
    return Response(
        response.content,
        media_type="audio/mpeg"
    )
@app.get("/gen/video/{prompt}")
@app.post("/gen/video")
@app.head("/gen/video")
async def genvideo_airforce(
    request: Request,
    prompt: str = None,
    authorization: Optional[str] = Header(None),
    x_client_id: Optional[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
    ratio = None
    mode = 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(
                status_code=400,
                detail=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(
                status_code=400,
                detail=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")

        # Clamp duration
        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":  "1:1",
    }
    pollinations_ratio = RATIO_MAP.get(aspectRatio, "16:9")

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

    if image_urls:
        params["image"] = "|".join(image_urls[:2])

    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}"

    print(f"[VIDEO GEN] Pollinations URL: {url}")
    url = url + f"&key={PKEY}"

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

    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",
        },
    )

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"}

MAX_VIDEO_RETRIES = 6


@app.get("/gen/video/airforce/{prompt}")
@app.post("/gen/video/airforce")
@app.head("/gen/video/airforce")
async def genvideo_airforce(
    request: Request,
    prompt: str = None,
    authorization: Optional[str] = Header(None),
    x_client_id: Optional[str] = Header(None),
):
    if request.method == "HEAD":
        return Response(
            status_code=200,
            headers={
    
                # Required field
                "Y-prompt": "string — required. The text prompt used to generate the video.",
    
                # Optional fields
                "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",
    
                # Response format
                "Y-response_format": "video/mp4",
    
                # Model info
                "Y-model": "grok-imagine-video"
            }
        )

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

    user_body = {}
    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(
                status_code=400,
                detail=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(
                status_code=400,
                detail=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",
        },
    )

@app.get("/subscription")
async def get_subscription(authorization: Optional[str] = Header(None)):
    if not authorization or not authorization.startswith("Bearer "):
        raise HTTPException(401, "Missing or invalid Authorization header")

    jwt = authorization.split(" ", 1)[1]

    result = await fetch_subscription(jwt)
    if "error" in result:
        raise HTTPException(401, result["error"])

    plan_key = normalize_plan_key(result.get("plan_key"))
    result["plan_key"] = plan_key
    result["plan_name"] = (TIER_CONFIG.get(plan_key) or TIER_CONFIG["free"])["name"]
    return result


@app.get("/usage")
async def get_usage(
    request: Request,
    authorization: Optional[str] = Header(None),
    x_client_id: Optional[str] = Header(None),
):
    plan_key, subject = await resolve_rate_limit_identity(request, authorization, x_client_id)
    plan = TIER_CONFIG.get(plan_key) or TIER_CONFIG["free"]
    usage = get_usage_snapshot_for_subject(plan_key, subject)
    return JSONResponse(
        status_code=200,
        content={
            "plan_key": plan_key,
            "plan_name": plan.get("name", "Free Tier"),
            "usage": usage,
            "generated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        },
    )

@app.get("/tier-config")
async def tier_config():
    plans = []
    for idx, key in enumerate(PLAN_ORDER):
        plan = TIER_CONFIG.get(key)
        if not plan:
            continue
        plans.append(
            {
                "key": key,
                "name": plan["name"],
                "url": plan["url"],
                "price": plan["price"],
                "limits": plan["limits"],
                "order": idx,
            }
        )

    return JSONResponse(
        status_code=200,
        content={
            "defaultPlanKey": "free",
            "plans": plans,
        },
    )

@app.get("/tiers")
async def tiers():
    paid_plans = []
    for key in PLAN_ORDER:
        if key == "free":
            continue
        plan = TIER_CONFIG.get(key)
        if not plan:
            continue
        paid_plans.append(
            {
                "key": key,
                "name": plan["name"],
                "url": plan["url"],
                "price": plan["price"],
                "limits": plan["limits"],
            }
        )

    return JSONResponse(
        status_code=200,
        content=paid_plans,
    )

@app.get("/portal")
def a():
    return RedirectResponse(url="https://billing.stripe.com/p/login/test_6oUcN5g665rp7nLgaq8bS00", status_code=status.HTTP_302_FOUND)