import os
import time
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
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": None,
            "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": {},
}
IDENTITY_CACHE_TTL_SECONDS = 60
identity_cache = {}

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 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)


async def resolve_rate_limit_identity(
    request: Request,
    authorization: Optional[str],
) -> tuple[str, str]:
    now = time.time()
    default_subject = f"ip:{request.client.host if request.client else 'unknown'}"
    if not authorization or not authorization.startswith("Bearer "):
        return "free", default_subject

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

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

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

    if not isinstance(sub, dict) or sub.get("error"):
        return "free", 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,
    }
    return plan_key, subject


async def enforce_rate_limit(
    request: Request,
    authorization: Optional[str],
    metric: str,
) -> 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)
    plan = TIER_CONFIG.get(plan_key) or TIER_CONFIG["free"]
    plan_limits = plan.get("limits", {})
    limit = plan_limits.get(metric)

    if limit is None:
        return {"plan_key": plan_key, "remaining": None}

    window_key = get_usage_period_key(metric)
    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 entry["count"] >= int(limit):
        raise HTTPException(
            status_code=429,
            detail=f"{metric} limit reached for {plan.get('name', 'current plan')}",
        )

    entry["count"] += 1
    remaining = max(0, int(limit) - entry["count"])
    return {"plan_key": plan_key, "remaining": remaining}


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


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]):
    await enforce_rate_limit(request, authorization, "imagesDaily")


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

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]):
    return await enforce_rate_limit(request, authorization, "cloudChatDaily")

@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),
):
    await check_image_rate_limit(request, authorization)
    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"

    if prompt is None:
        if (await request.json()).get("mode"):
            if mode == "fantasy":
                chosen_model = "flux"
            elif 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)):
    body = await request.json()
    messages = body.get("messages", [])
    if not isinstance(messages, list) or len(messages) == 0:
        raise HTTPException(400, "messages[] is required")

    await check_chat_rate_limit(request, authorization)
    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"

    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:
                        content_type = r.headers.get(
                            "content-type", "text/event-stream"
                        )
    
                        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)
        
        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,
    authorization: Optional[str] = Header(None),
):
    await check_audio_rate_limit(request, authorization)
    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,
    authorization: Optional[str] = Header(None),
):
    await check_audio_rate_limit(request, authorization)
    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")
@app.head("/gen/video")
async def genvideo(
    request: Request,
    prompt: str = None,
    authorization: Optional[str] = Header(None),
):
    if request.method != "HEAD":
        await check_video_rate_limit(request, authorization)
    return RedirectResponse(url="/gen/video/airforce", status_code=status.HTTP_307_TEMPORARY_REDIRECT)

    if prompt is None:

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

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

    await check_video_rate_limit(request, authorization)

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

    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 not prompt:
            raise HTTPException(400, "Prompt is required")

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

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

    if prompt is None:
        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("/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)