api.py

from __future__ import annotations

import json
import shutil
import uuid
from pathlib import Path
from typing import Any

from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.responses import FileResponse
from pydantic import BaseModel, Field

from renderer.core.config import Settings
from renderer.core.ingest import stage_upload
from renderer.core.models import AIReelsRequest, RenderRequest, Scene
from renderer.core.utils import safe_filename
from renderer.jobs import JobManager
from renderer.scenes import Timeline

settings = Settings()
settings.ensure_dirs()
job_manager = JobManager(settings)
api = FastAPI(title="Basyx FFmpeg Rendering Engine", version="1.0.0")


class ScenePayload(BaseModel):
    start: float = Field(ge=0)
    duration: float = Field(gt=0)
    media: str
    caption: str = ""
    transition: str = "fade"
    background: str = "blur"


class RenderPayload(BaseModel):
    scenes: list[ScenePayload]
    template: str = "tiktok_classic"
    output_name: str = "render.mp4"
    voiceover: str | None = None
    background_music: str | None = None
    subtitle_format: str = "ass"
    normalize: bool = True
    metadata: dict[str, Any] = Field(default_factory=dict)


class AIReelsPayload(BaseModel):
    script: str
    voiceover: str
    assets: list[str]
    template: str = "tiktok_classic"
    output_name: str = "ai_reel.mp4"
    background_music: str | None = None


class BatchPayload(BaseModel):
    jobs: list[RenderPayload]


class UploadedAsset(BaseModel):
    filename: str
    path: str
    reference: str


@api.get("/health")
def health() -> dict[str, str]:
    return {"status": "ok"}


@api.post("/render")
def render(payload: RenderPayload | AIReelsPayload) -> dict[str, str]:
    if isinstance(payload, AIReelsPayload):
        job_id = job_manager.submit_ai_reels(AIReelsRequest(**payload.model_dump()))
    else:
        job_id = job_manager.submit_render(_render_request(payload))
    return {"job_id": job_id, "status_url": f"/status/{job_id}", "download_url": f"/download/{job_id}"}


@api.post("/render/ai-reels")
def render_ai_reels(payload: AIReelsPayload) -> dict[str, str]:
    job_id = job_manager.submit_ai_reels(AIReelsRequest(**payload.model_dump()))
    return {"job_id": job_id, "status_url": f"/status/{job_id}", "download_url": f"/download/{job_id}"}


@api.post("/render/batch")
def render_batch(payload: BatchPayload) -> dict[str, list[str]]:
    job_ids = job_manager.submit_batch([_render_request(job) for job in payload.jobs])
    return {"job_ids": job_ids}


@api.post("/render/upload")
async def render_upload(request_json: str = Form(...), files: list[UploadFile] = File(default=[])) -> dict[str, str]:
    uploads = await _stage_uploads(files)
    payload_data = _replace_upload_refs(json.loads(request_json), uploads)
    payload = RenderPayload.model_validate(payload_data)
    job_id = job_manager.submit_render(_render_request(payload))
    return {"job_id": job_id, "status_url": f"/status/{job_id}", "download_url": f"/download/{job_id}"}


@api.post("/render/ai-reels/upload")
async def render_ai_reels_upload(request_json: str = Form(...), files: list[UploadFile] = File(default=[])) -> dict[str, str]:
    uploads = await _stage_uploads(files)
    payload_data = _replace_upload_refs(json.loads(request_json), uploads)
    payload = AIReelsPayload.model_validate(payload_data)
    job_id = job_manager.submit_ai_reels(AIReelsRequest(**payload.model_dump()))
    return {"job_id": job_id, "status_url": f"/status/{job_id}", "download_url": f"/download/{job_id}"}


@api.post("/assets/upload")
async def upload_assets(files: list[UploadFile] = File(...)) -> dict[str, list[UploadedAsset]]:
    uploads = await _stage_uploads(files)
    assets = [
        UploadedAsset(filename=filename, path=path, reference=f"upload://{filename}") for filename, path in uploads.items()
    ]
    return {"assets": [asset.model_dump() for asset in assets]}


@api.get("/status/{job_id}")
def status(job_id: str) -> dict:
    try:
        return job_manager.get(job_id).__dict__
    except KeyError as exc:
        raise HTTPException(status_code=404, detail="Job not found") from exc


@api.get("/download/{job_id}")
def download(job_id: str) -> FileResponse:
    try:
        record = job_manager.get(job_id)
    except KeyError as exc:
        raise HTTPException(status_code=404, detail="Job not found") from exc
    if record.state != "COMPLETED" or not record.output_path:
        raise HTTPException(status_code=409, detail=f"Job is {record.state}")
    path = Path(record.output_path)
    if not path.exists():
        raise HTTPException(status_code=404, detail="Output file is missing")
    return FileResponse(path, media_type="video/mp4", filename=path.name)


@api.post("/inspect")
def inspect_asset(path: str) -> dict:
    from renderer import RenderEngine

    try:
        return RenderEngine(settings).inspect_asset(path)
    except Exception as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc


def _render_request(payload: RenderPayload) -> RenderRequest:
    request = RenderRequest(
        scenes=[Scene(**scene.model_dump()) for scene in payload.scenes],
        template=payload.template,
        output_name=payload.output_name,
        voiceover=payload.voiceover,
        background_music=payload.background_music,
        subtitle_format=payload.subtitle_format,  # type: ignore[arg-type]
        normalize=payload.normalize,
        metadata=payload.metadata,
    )
    Timeline(request.scenes)
    return request


async def _stage_uploads(files: list[UploadFile]) -> dict[str, str]:
    upload_dir = settings.temp_dir / "uploads" / uuid.uuid4().hex
    staged: dict[str, str] = {}
    total_bytes = 0
    for upload in files:
        filename = safe_filename(upload.filename or f"asset_{len(staged)}")
        target = upload_dir / filename
        target.parent.mkdir(parents=True, exist_ok=True)
        with target.open("wb") as output:
            while True:
                chunk = await upload.read(1024 * 1024)
                if not chunk:
                    break
                total_bytes += len(chunk)
                if total_bytes > settings.max_download_bytes:
                    shutil.rmtree(upload_dir, ignore_errors=True)
                    raise HTTPException(status_code=413, detail="Uploaded assets exceed MAX_DOWNLOAD_BYTES")
                output.write(chunk)
        staged[filename] = str(stage_upload(target, upload_dir, filename))
    return staged


def _replace_upload_refs(value: Any, uploads: dict[str, str]) -> Any:
    if isinstance(value, dict):
        return {key: _replace_upload_refs(item, uploads) for key, item in value.items()}
    if isinstance(value, list):
        return [_replace_upload_refs(item, uploads) for item in value]
    if isinstance(value, str) and value.startswith("upload://"):
        name = safe_filename(value.removeprefix("upload://"))
        if name not in uploads:
            raise HTTPException(status_code=400, detail=f"Missing uploaded file for reference: upload://{name}")
        return uploads[name]
    return value