app.py

import os
import time
from fastapi import FastAPI, Request, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import Response, JSONResponse, StreamingResponse
import httpx
from bs4 import BeautifulSoup
from typing import List, Dict
import asyncio
import re
from random import randint
from urllib.parse import quote
import uuid

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=["*"],
)

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

RATE_LIMIT = 60
WINDOW_SECONDS = 60 * 60 * 24
ip_store = {}  # { ip: { "count": int, "reset": timestamp } }
AUDIO_RATE_LIMIT = 10
AUDIO_WINDOW_SECONDS = 60 * 60 * 24
audio_ip_store = {}

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(
        m.get("content", "")
        for m in messages
        if m.get("role") == "user"
    ).lower()

def check_audio_rate_limit(ip: str):
    now = time.time()

    if ip not in audio_ip_store:
        audio_ip_store[ip] = {
            "count": 0,
            "reset": now + AUDIO_WINDOW_SECONDS
        }

    entry = audio_ip_store[ip]

    if now > entry["reset"]:
        entry["count"] = 0
        entry["reset"] = now + AUDIO_WINDOW_SECONDS

    if entry["count"] >= AUDIO_RATE_LIMIT:
        raise HTTPException(
            status_code=429,
            detail="Daily audio limit reached: 10 per IP"
        )

    entry["count"] += 1


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

def check_rate_limit(ip: str):
    now = time.time()

    if ip not in ip_store:
        ip_store[ip] = {"count": 0, "reset": now + WINDOW_SECONDS}

    entry = ip_store[ip]

    if now > entry["reset"]:
        entry["count"] = 0
        entry["reset"] = now + WINDOW_SECONDS

    if entry["count"] >= RATE_LIMIT:
        raise HTTPException(
            status_code=429,
            detail="Daily limit reached: 25 images per IP"
        )

    entry["count"] += 1

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

CHAT_RATE_LIMIT = 50
CHAT_WINDOW_SECONDS = 60 * 60

chat_ip_store = {}

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",
]

def check_chat_rate_limit(ip: str):
    now = time.time()

    if ip not in chat_ip_store:
        chat_ip_store[ip] = {
            "count": 0,
            "reset": now + CHAT_WINDOW_SECONDS
        }

    entry = chat_ip_store[ip]

    if now > entry["reset"]:
        entry["count"] = 0
        entry["reset"] = now + CHAT_WINDOW_SECONDS

    if entry["count"] >= CHAT_RATE_LIMIT:
        raise HTTPException(
            status_code=429,
            detail="Chat rate limit exceeded"
        )

    entry["count"] += 1
    return entry["count"]

@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=207,
        headers={
            "Content-Type": "video/mp4",
            "Accept-Ranges": "bytes",
        }
    )

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

@app.post("/gen/image")
@app.get("/genimg/{prompt}")
async def generate_image(request: Request, prompt: str = None):
    client_ip = request.client.host
    check_rate_limit(client_ip)
    timeout = httpx.Timeout(300.0, read=300.0)
    if prompt is None:
        prompt = (await request.json()).get("prompt")
    if is_cinematic_image_prompt(prompt):
        chosen_model = "flux"
    else:
        chosen_model = "zimage"

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

    url = f"https://gen.pollinations.ai/image/{prompt}?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):
    body = await request.json()

    messages = body.get("messages", [])
    if not isinstance(messages, list) or len(messages) == 0:
        raise HTTPException(400, "messages[] is required")

    ip = request.client.host
    msg_count = check_chat_rate_limit(ip)
    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.1-8b-instant"
    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 >= 6:
        chosen_model = "gpt-oss-120b"
        provider = "cerebras"
    
    elif score >= 4:
        chosen_model = "llama-3.3-70b-versatile"
        provider = "groq"
    
    elif score >= 3 and structured_task:
        chosen_model = "qwen-3-235b-a22b-instruct-2507"
        provider = "cerebras"
        
    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:
        async def event_generator():
            async with httpx.AsyncClient(timeout=None) as client:
                async with client.stream("POST", url, json=body, headers=headers) as r:
                    async for chunk in r.aiter_raw():
                        yield chunk
    
        return StreamingResponse(
            event_generator(),
            media_type="text/event-stream",
        )
    else:
        async with httpx.AsyncClient(timeout=None) as client:
            r = await client.post(url, json=body, headers=headers)
        
        return JSONResponse(
            status_code=r.status_code,
            content=r.json()
        )

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

@app.get("/gen/sfx/{prompt}")
@app.post("/gen/sfx")
async def gensfx(request: Request, prompt: str = None):
    client_ip = request.client.host
    check_audio_rate_limit(client_ip)
    if prompt is None:
        prompt = (await request.json()).get("prompt")
    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):
    client_ip = request.client.host
    check_rate_limit(client_ip)
    if prompt is None:
        prompt = (await request.json()).get("prompt")
    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")
async def genvideo(request: Request, prompt: str = None):
    client_ip = request.client.host
    check_rate_limit(client_ip)

    if prompt is None:
        return RedirectResponse(url="/gen/video/airforce", status_code=status.HTTP_307_TEMPORARY_REDIRECT)

        body = await request.json()
        prompt = body.get("prompt")
    else:
        return RedirectResponse(url=f"/gen/video/airforce/{prompt}", status_code=status.HTTP_307_TEMPORARY_REDIRECT)
    if not prompt:
        raise HTTPException(400, "Prompt is required")

    encoded_prompt = quote(prompt)
    url = f"https://gen.pollinations.ai/image/{encoded_prompt}?model=grok-video&key={PKEY}"

    async def fetch_with_retry(url, max_retries=6):
        async with httpx.AsyncClient(timeout=300.0) as client:
            for attempt in range(max_retries):
                response = await client.get(url)

                if response.status_code == 200:
                    return response

                body_text = ""
                try:
                    body_text = response.text
                except Exception:
                    pass

                if response.status_code == 429 and "api.airforce" in body_text:
                    wait = 0.5 * (2 ** attempt)
                    print(f"[VIDEO RETRY] 429 from api.airforce, retrying in {wait:.2f}s...")
                    await asyncio.sleep(wait)
                    continue

                content_type = response.headers.get("content-type", "")

                if "application/json" in content_type:
                    return JSONResponse(
                        status_code=response.status_code,
                        content=response.json()
                    )
                else:
                    return JSONResponse(
                        status_code=response.status_code,
                        content={
                            "success": False,
                            "error": "Upstream video generation failed",
                            "status_code": response.status_code,
                            "message": body_text[:1000]
                        }
                    )

            return JSONResponse(
                status_code=503,
                content={
                    "success": False,
                    "error": "Pollinations video generation overloaded (api.airforce). Try again later."
                }
            )

    result = await fetch_with_retry(url)
    
    if isinstance(result, JSONResponse):
        return result
    
    response = result

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

@app.get("/gen/video/exp/{prompt}")
@app.post("/gen/video/exp")
async def genvideo_wan(request: Request, prompt: str = None):
    client_ip = request.client.host
    check_rate_limit(client_ip)

    if prompt is None:
        body = await request.json()
        prompt = body.get("prompt")

    if not prompt:
        raise HTTPException(400, "Prompt is required")

    session_hash = uuid.uuid4().hex

    join_payload = {
        "fn_index": 0,
        "session_hash": session_hash,
        "data": [prompt],
    }

    async with httpx.AsyncClient(timeout=60.0) as client:
        # 1️⃣ Join queue
        join = await client.post(
            f"{WAN_API_BASE}/queue/join",
            json=join_payload,
        )

        if join.status_code != 200:
            raise HTTPException(
                join.status_code,
                f"Failed to join WAN queue: {join.text}"
            )

        for _ in range(120):
            await asyncio.sleep(1)

            poll = await client.get(
                f"{WAN_API_BASE}/queue/data",
                params={"session_hash": session_hash},
            )

            if poll.status_code != 200:
                continue

            data = poll.json()

            if data.get("status") == "complete":
                try:
                    video_url = data["data"][0]
                except Exception:
                    raise HTTPException(
                        500,
                        "WAN completed but returned unexpected format"
                    )

                # 3️⃣ Fetch video
                video = await client.get(video_url)

                if video.status_code != 200:
                    raise HTTPException(
                        502,
                        "Failed to fetch generated video"
                    )

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

            if data.get("status") == "error":
                raise HTTPException(
                    500,
                    f"WAN generation error: {data}"
                )

    raise HTTPException(
        504,
        "WAN video generation timed out"
    )

AIRFORCE_KEY = os.getenv("AIRFORCE")
AIRFORCE_VIDEO_MODEL = "wan-2.6"
AIRFORCE_API_URL = "https://api.airforce/v1/images/generations"

MAX_VIDEO_RETRIES = 6

@app.get("/gen/video/airforce/{prompt}")
@app.post("/gen/video/airforce")
async def genvideo_airforce(request: Request, prompt: str = None):
    if prompt is None:
        body = await request.json()
        prompt = body.get("prompt")
        duration = body.get("duration", 5)

    body = {
        "model": "grok-imagine-video",
        "prompt": prompt,
        "n": 1,
        "size": "1024x1024",
        "response_format": "b64_json",
        "sse": False,
        "mode": "normal",
        "aspectRatio": "3:2"
    }

    async with httpx.AsyncClient(timeout=600) as client:
        resp = await client.post(
            "https://api.airforce/v1/images/generations",
            headers={"Authorization": f"Bearer {AIRFORCE_KEY}", "Content-Type": "application/json"},
            json=body
        )

    if resp.status_code != 200:
        return JSONResponse(status_code=resp.status_code, content=resp.json())
    if not resp.content:
        raise HTTPException(
            status_code=502,
            detail="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",
        },
    )