# main.py

import requests
import uuid
import time
import brotli
import json
import asyncio
from typing import Literal

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field

# --- FastAPI App Initialization ---
app = FastAPI(
    title="Sora Video Generation API",
    description="An unofficial API wrapper to generate videos using the Bylo.ai Sora model.",
    version="1.0.0",
)

# --- Pydantic Models for Request and Response ---
class SoraRequest(BaseModel):
    prompt: str = Field(..., description="The text prompt to generate the video from.")
    ratio: Literal["portrait", "landscape"] = Field(
        "portrait", 
        description="The aspect ratio of the generated video."
    )

# --- Core Logic (Adapted from your script) ---

def decode_response(resp: requests.Response):
    """
    Safely decode response: Brotli → UTF-8 → JSON, with automatic fallback if Brotli fails.
    """
    raw_data = resp.content
    encoding = resp.headers.get("Content-Encoding", "").lower()

    if encoding == "br":
        try:
            raw_data = brotli.decompress(raw_data)
        except brotli.error:
            pass  # Fall back to raw UTF-8 if not actually Brotli compressed

    try:
        return json.loads(raw_data.decode("utf-8"))
    except Exception as je:
        snippet = raw_data[:500]
        raise RuntimeError(f"Failed to decode JSON: {je}\nRaw response snippet:\n{snippet}")

def sora(prompt: str, ratio: str = "portrait"):
    """
    The core function to interact with the Bylo.ai Sora API.
    This is a synchronous function intended to be run in a thread.
    """
    if not prompt:
        raise ValueError("Prompt is required.")
    if ratio not in ["portrait", "landscape"]:
        raise ValueError("Available ratios: portrait, landscape.")

    base_url = "https://api.bylo.ai/aimodels/api/v1/ai"

    headers = {
        "accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7",
        "accept-language": "id-ID,id;q=0.9,en-US;q=0.8,en;q=0.7",
        "accept-encoding": "gzip, deflate, br",
        "cache-control": "max-age=0",
        "connection": "keep-alive",
        "content-type": "application/json; charset=UTF-8",
        "dnt": "1",
        "origin": "https://bylo.ai",
        "pragma": "no-cache",
        "referer": "https://bylo.ai/features/sora-2",
        "sec-ch-prefers-color-scheme": "dark",
        "sec-ch-ua": '"Chromium";v="137", "Not/A)Brand";v="24"',
        "sec-ch-ua-mobile": "?1",
        "user-agent": "Mozilla/5.0 (Linux; Android 15; SM-F958 Build/AP3A.240905.015) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/130.0.6723.86 Mobile Safari/537.36",
        "x-requested-with": "XMLHttpRequest",
        "uniqueId": str(uuid.uuid4()).replace("-", "")
    }

    payload = {
        "prompt": prompt,
        "channel": "SORA2",
        "pageId": 536,
        "source": "bylo.ai",
        "watermarkFlag": True,
        "privateFlag": True,
        "isTemp": True,
        "vipFlag": True,
        "model": "sora_video2",
        "videoType": "text-to-video",
        "aspectRatio": ratio
    }

    try:
        # STEP 1: Create generation task
        create_resp = requests.post(f"{base_url}/video/create", headers=headers, json=payload)
        create_resp.raise_for_status()
        task_data = decode_response(create_resp)

        task_id = task_data.get("data")
        if not task_id:
            raise RuntimeError(f"API did not return task_id. Full response:\n{json.dumps(task_data, indent=2)}")

        # STEP 2: Poll until done
        # Set a reasonable timeout to prevent infinite loops
        polling_start_time = time.time()
        timeout_seconds = 300  # 5 minutes

        while True:
            if time.time() - polling_start_time > timeout_seconds:
                raise RuntimeError("Polling timed out after 5 minutes.")

            poll_resp = requests.get(f"{base_url}/{task_id}?channel=SORA2", headers=headers)
            poll_resp.raise_for_status()
            poll_json = decode_response(poll_resp)

            state = poll_json.get("data", {}).get("state", 0)
            if state > 0:
                complete_data_str = poll_json.get("data", {}).get("completeData")
                if not complete_data_str:
                    raise RuntimeError(f"Generation completed but no 'completeData' found. Full response: \n{poll_json}")
                return json.loads(complete_data_str)

            time.sleep(5) # Increased sleep time for polling

    except requests.exceptions.RequestException as re:
        raise RuntimeError(f"Network error during Sora generation: {re}")
    except Exception as e:
        # Re-raise exceptions to be caught by the endpoint handler
        raise RuntimeError(f"Sora generation failed: {e}")


# --- API Endpoint ---
@app.post("/generate", summary="Create a new Sora video")
async def generate_video(request: SoraRequest):
    """
    Accepts a prompt and ratio, then generates a video by polling the Bylo.ai API.
    """
    try:
        # Run the synchronous, blocking sora function in a separate thread
        result = await asyncio.to_thread(sora, request.prompt, request.ratio)
        return result
    except ValueError as ve:
        # Handle bad input from the user
        raise HTTPException(status_code=400, detail=str(ve))
    except RuntimeError as re:
        # Handle errors from the external API or other processing issues
        raise HTTPException(status_code=500, detail=str(re))

@app.get("/", include_in_schema=False)
def root():
    return {"message": "Sora API is running. See /docs for usage."}