# app.py
from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
import uvicorn
import os
import uuid
import subprocess
import json

from huggingface_hub import hf_hub_download
from moviepy.editor import ImageSequenceClip

app = FastAPI(title="WAN2 GGUF API", version="1.0")

# -------------------- Directories --------------------
MODEL_REPO = "calcuis/wan2-gguf"
MODEL_FILE = "wan2.2-animate-14b-q4_0.gguf"
MODEL_DIR = "models"
OUTPUT_DIR = "outputs"
NODE_CLI = "gguf-node-cli.js"  # Path to your gguf-node CLI

os.makedirs(MODEL_DIR, exist_ok=True)
os.makedirs(OUTPUT_DIR, exist_ok=True)

# -------------------- Download model --------------------
model_path = hf_hub_download(
    repo_id=MODEL_REPO,
    filename=MODEL_FILE,
    local_dir=MODEL_DIR
)
print("✅ Model downloaded to:", model_path)

# -------------------- Request schema --------------------
class PromptRequest(BaseModel):
    prompt: str
    steps: int = 20

# -------------------- Helper: GGUF Node CLI --------------------
def generate_frames_with_node(prompt, steps=20):
    """
    Calls gguf-node CLI to generate frames.
    CLI should output JSON array of frame image paths.
    """
    try:
        result = subprocess.run(
            ["node", NODE_CLI, "--model", model_path, "--prompt", prompt, "--steps", str(steps)],
            capture_output=True,
            text=True,
            check=True
        )
        frames = json.loads(result.stdout)
        return frames
    except subprocess.CalledProcessError as e:
        print("Error calling gguf-node:", e.stderr)
        return []

# -------------------- Routes --------------------
@app.get("/")
def root():
    return {"message": "WAN2 GGUF API is running!"}

@app.get("/generate")
def generate_video_get(q: str, steps: int = 20):
    """Allows GET requests with ?q=... for browser testing"""
    return generate_video(PromptRequest(prompt=q, steps=steps))

@app.post("/generate")
def generate_video(request: PromptRequest):
    """Generates video from prompt using WAN2 GGUF"""
    # Unique filename
    file_id = str(uuid.uuid4())
    file_path = os.path.join(OUTPUT_DIR, f"{file_id}.mp4")

    # ---- WAN2 inference via Node CLI ----
    frames = generate_frames_with_node(request.prompt, request.steps)

    if not frames:
        return {"status": "error", "message": "Failed to generate frames"}

    # ---- Save frames as MP4 ----
    clip = ImageSequenceClip(frames, fps=12)
    clip.write_videofile(file_path, codec="libx264", audio=False, verbose=False, logger=None)

    # Build full URL for Hugging Face Space
    base_url = "https://abrahamdw882-wan2-api.hf.space"
    video_url = f"{base_url}/file/{file_id}.mp4"

    return {
        "status": "success",
        "model_file": MODEL_FILE,
        "prompt": request.prompt,
        "steps": request.steps,
        "video_url": video_url
    }

# -------------------- Serve output videos --------------------
app.mount("/file", StaticFiles(directory=OUTPUT_DIR), name="file")

# -------------------- Run server --------------------
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)