from fastapi import FastAPI, UploadFile, File, Form
from fastapi.responses import FileResponse, JSONResponse
from gradio_client import Client, handle_file
from gradio_client.exceptions import AppError
import tempfile
import os
import shutil

HF_TOKEN = os.getenv("HF_TOKEN")

app = FastAPI(title="ZHSAPI")


@app.get("/health")
def health():
    return {"status": "ok"}


@app.post("/timelapse")
async def generate_timelapse(
    start_frame: UploadFile = File(...),
    end_frame: UploadFile = File(...),
    prompt: str = Form(default="a construction timelapse"),
):
    with tempfile.TemporaryDirectory() as tmpdir:
        start_path = os.path.join(tmpdir, start_frame.filename)
        end_path = os.path.join(tmpdir, end_frame.filename)

        with open(start_path, "wb") as f:
            shutil.copyfileobj(start_frame.file, f)
        with open(end_path, "wb") as f:
            shutil.copyfileobj(end_frame.file, f)

        # Try diffusion model first, fall back to CPU FILM interpolation on quota error
        try:
            client = Client("linoyts/LTX-2-3-First-Last-Frame", token=HF_TOKEN)
            result = client.predict(
                first_image=handle_file(start_path),
                last_image=handle_file(end_path),
                input_audio=None,
                prompt=prompt,
                api_name="/generate_video",
            )
            output_path = result[0] if isinstance(result, (list, tuple)) else result
        except AppError as e:
            if "GPU quota" not in str(e):
                return JSONResponse(status_code=502, content={"error": str(e)})

            # Fallback: CPU-based FILM interpolation (no GPU quota)
            film = Client("freealise/video_frame_interpolation")
            loaded = film.predict(
                f=[handle_file(start_path), handle_file(end_path)],
                r_bg=False,
                api_name="/loadf",
            )
            frames = loaded[0]
            result = film.predict(
                f_in=frames,
                interpolation=4,
                fps_output=0,
                api_name="/infer",
            )
            output_path = result[0]["video"] if isinstance(result[0], dict) else result[0]

        return FileResponse(output_path, media_type="video/mp4", filename="timelapse.mp4")