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

# GGUF & video libraries
from gguf_runtime import GGUFModel  # Make sure gguf-runtime is installed
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"

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

# -------------------- Download model --------------------
from huggingface_hub import hf_hub_download

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

# -------------------- Load GGUF model --------------------
print("Loading WAN2 GGUF model...")
wan2_model = GGUFModel(model_path)
print("✅ WAN2 model loaded")

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

# -------------------- 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 ----
    # This generates frames from prompt
    frames = wan2_model.generate_video(
        prompt=request.prompt,
        steps=request.steps
    )

    # ---- 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)