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, Query, 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.security import verify_download_token
from renderer.core.utils import safe_filename
from renderer.jobs import JobManager
from renderer.scenes import Timeline
from renderer import RenderEngine
from renderer.subtitles import SubtitleEvent, SubtitleGenerator
from renderer.templates import apply_preset, list_presets, list_templates

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"
    layout: str = "fill"
    effect: str | None = None


class RenderPayload(BaseModel):
    scenes: list[ScenePayload]
    template: str = "tiktok_classic"
    preset: str | None = None
    output_name: str = "render.mp4"
    voiceover: str | None = None
    background_music: str | None = None
    subtitle_format: str = "ass"
    auto_subtitles: bool = False
    subtitle_language: str | None = None
    whisper_model_size: str | None = None
    preview: bool = False
    audio_normalize: bool = False
    watermark: str | None = None
    watermark_position: str = "bottom-right"
    intro: str | None = None
    outro: str | None = None
    callback_url: str | None = None
    export_target: str | None = None
    priority: int = 0
    scheduled_at: float | None = None
    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


class TranscribePayload(BaseModel):
    audio: str
    model_size: str | None = None
    language: str | None = None
    task: str = "transcribe"
    beam_size: int = Field(default=5, ge=1, le=10)
    vad_filter: bool = True
    word_timestamps: bool = True


class SubtitlePayload(BaseModel):
    events: list[dict[str, Any]]
    format: str = "srt"
    template: str = "tiktok_classic"


class SceneBuildPayload(BaseModel):
    script: str
    assets: list[str]
    duration: float | None = None
    transition: str = "fade"


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


@api.get("/presets")
def presets() -> dict[str, list[str]]:
    return {"presets": list_presets(), "caption_templates": list_templates()}


@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))
    record = job_manager.get(job_id)
    return {
        "job_id": job_id,
        "status_url": f"/status/{job_id}",
        "download_url": f"/download/{job_id}?token={record.download_token}",
    }


@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()))
    record = job_manager.get(job_id)
    return {"job_id": job_id, "status_url": f"/status/{job_id}", "download_url": f"/download/{job_id}?token={record.download_token}"}


@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))
    record = job_manager.get(job_id)
    return {"job_id": job_id, "status_url": f"/status/{job_id}", "download_url": f"/download/{job_id}?token={record.download_token}"}


@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()))
    record = job_manager.get(job_id)
    return {"job_id": job_id, "status_url": f"/status/{job_id}", "download_url": f"/download/{job_id}?token={record.download_token}"}


@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.post("/transcribe")
def transcribe(payload: TranscribePayload) -> dict:
    from renderer.core.ingest import AssetIngestor
    from renderer.core.utils import temp_workdir

    try:
        with temp_workdir(settings.temp_dir, "transcribe") as work:
            audio = AssetIngestor(settings).resolve(payload.audio, Path(work) / "inputs", "audio")
            return RenderEngine(settings).transcribe(
                audio,
                model_size=payload.model_size,
                language=payload.language,
                task=payload.task,
                beam_size=payload.beam_size,
                vad_filter=payload.vad_filter,
                word_timestamps=payload.word_timestamps,
            )
    except Exception as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc


@api.post("/subtitles")
def subtitles(payload: SubtitlePayload) -> FileResponse:
    try:
        events = [SubtitleEvent(float(event["start"]), float(event["end"]), str(event["text"])) for event in payload.events]
        path = settings.temp_dir / f"subtitles_{uuid.uuid4().hex}.{payload.format}"
        generator = SubtitleGenerator()
        if payload.format == "ass":
            generator.write_ass(events, path, payload.template)
            media_type = "text/x-ssa"
        else:
            generator.write_srt(events, path)
            media_type = "application/x-subrip"
        return FileResponse(path, media_type=media_type, filename=path.name)
    except Exception as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc


@api.post("/scene-builder")
def scene_builder(payload: SceneBuildPayload) -> dict[str, list[dict[str, Any]]]:
    if not payload.assets:
        raise HTTPException(status_code=400, detail="At least one asset is required")
    words = payload.script.split()
    total_duration = payload.duration or max(3.0, len(words) * 0.35)
    per_scene = total_duration / len(payload.assets)
    chunk_size = max(1, round(len(words) / len(payload.assets)))
    captions = [" ".join(words[index : index + chunk_size]) for index in range(0, len(words), chunk_size)]
    scenes = [
        {
            "start": round(index * per_scene, 3),
            "duration": round(per_scene, 3),
            "media": asset,
            "caption": captions[index] if index < len(captions) else "",
            "transition": payload.transition,
        }
        for index, asset in enumerate(payload.assets)
    ]
    return {"scenes": scenes}


@api.post("/transcribe/upload")
async def transcribe_upload(
    file: UploadFile = File(...),
    model_size: str | None = Form(default=None),
    language: str | None = Form(default=None),
    task: str = Form(default="transcribe"),
    beam_size: int = Form(default=5),
    vad_filter: bool = Form(default=True),
    word_timestamps: bool = Form(default=True),
) -> dict:
    try:
        uploads = await _stage_uploads([file])
        audio = next(iter(uploads.values()))
        return RenderEngine(settings).transcribe(
            audio,
            model_size=model_size,
            language=language,
            task=task,
            beam_size=beam_size,
            vad_filter=vad_filter,
            word_timestamps=word_timestamps,
        )
    except Exception as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc


@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.post("/cancel/{job_id}")
def cancel(job_id: str) -> dict:
    try:
        return job_manager.cancel(job_id).__dict__
    except KeyError as exc:
        raise HTTPException(status_code=404, detail="Job not found") from exc


@api.post("/admin/cleanup")
def cleanup(older_than_seconds: int | None = None) -> dict[str, int]:
    return job_manager.cleanup(older_than_seconds)


@api.get("/download/{job_id}")
def download(job_id: str, token: str | None = Query(default=None)) -> 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}")
    if record.download_token and not verify_download_token(settings.signing_secret, job_id, token):
        raise HTTPException(status_code=403, detail="Invalid or missing download token")
    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:
    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:
    payload = RenderPayload.model_validate(apply_preset(payload.model_dump()))
    request = RenderRequest(
        scenes=[Scene(**scene.model_dump()) for scene in payload.scenes],
        template=payload.template,
        preset=payload.preset,
        output_name=payload.output_name,
        voiceover=payload.voiceover,
        background_music=payload.background_music,
        subtitle_format=payload.subtitle_format,  # type: ignore[arg-type]
        auto_subtitles=payload.auto_subtitles,
        subtitle_language=payload.subtitle_language,
        whisper_model_size=payload.whisper_model_size,
        preview=payload.preview,
        audio_normalize=payload.audio_normalize,
        watermark=payload.watermark,
        watermark_position=payload.watermark_position,
        intro=payload.intro,
        outro=payload.outro,
        callback_url=payload.callback_url,
        export_target=payload.export_target,
        priority=payload.priority,
        scheduled_at=payload.scheduled_at,
        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