import os
import uuid
import asyncio
from typing import Dict, Any

import replicate
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
import httpx

REPLICATE_API_TOKEN = os.getenv("REPLICATE_API_TOKEN")
if not REPLICATE_API_TOKEN:
    raise RuntimeError("Missing REPLICATE_API_TOKEN env var")

os.environ["REPLICATE_API_TOKEN"] = REPLICATE_API_TOKEN

VIDEO_OUTPUT_DIR = "video_outputs"
os.makedirs(VIDEO_OUTPUT_DIR, exist_ok=True)

app = FastAPI(title="nexus-video-api", version="2.0.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

VIDEO_MODELS = {
    "tiktok":       "minimax/video-01",
    "social":       "minimax/video-01",
    "cinematic":    "minimax/video-01",
    "ads":          "minimax/video-01",
    "music":        "minimax/video-01",
    "music_long":   "minimax/video-01",
    "experimental": "minimax/video-01",
    "avatar":       "minimax/video-01",
}

TASKS: Dict[str, Dict[str, Any]] = {}
TASKS_LOCK = asyncio.Lock()


def create_task(mode: str) -> str:
    task_id = str(uuid.uuid4())
    TASKS[task_id] = {"mode": mode, "status": "queued", "error": None, "video_path": None}
    return task_id


async def set_task(task_id: str, **updates: Any) -> None:
    async with TASKS_LOCK:
        if task_id in TASKS:
            TASKS[task_id].update(updates)


async def run_model(task_id: str, model_id: str, payload: Dict[str, Any]) -> None:
    await set_task(task_id, status="processing", error=None)
    try:
        output = await asyncio.to_thread(
            replicate.run,
            model_id,
            input={"prompt": payload.get("prompt", "")},
        )
        video_url = output[0] if isinstance(output, list) else str(output)

        async with httpx.AsyncClient(timeout=120) as client:
            resp = await client.get(video_url)
            resp.raise_for_status()
            video_bytes = resp.content

        if not video_bytes:
            raise RuntimeError("Empty video output")

        out_path = os.path.join(VIDEO_OUTPUT_DIR, f"{task_id}.mp4")
        with open(out_path, "wb") as f:
            f.write(video_bytes)

        await set_task(task_id, status="done", video_path=out_path)
    except Exception as e:
        await set_task(task_id, status="error", error=str(e), video_path=None)


@app.get("/")
def root():
    return {"status": "ok", "service": "nexus-video-api", "version": "2.0.0"}

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

@app.get("/video/status/{task_id}")
async def video_status(task_id: str):
    task = TASKS.get(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="task_not_found")
    return {"task_id": task_id, "mode": task["mode"], "status": task["status"], "error": task["error"], "ready": bool(task["video_path"]) and task["status"] == "done"}

@app.get("/video/{task_id}")
async def get_video(task_id: str):
    task = TASKS.get(task_id)
    if not task:
        raise HTTPException(status_code=404, detail="task_not_found")
    if not task.get("video_path") or task.get("status") != "done":
        return {"error": "video_not_ready"}
    return FileResponse(task["video_path"], media_type="video/mp4")

@app.post("/video/tiktok")
async def video_tiktok(data: dict):
    task_id = create_task("tiktok")
    asyncio.create_task(run_model(task_id, VIDEO_MODELS["tiktok"], data))
    return {"task_id": task_id}

@app.post("/video/social")
async def video_social(data: dict):
    task_id = create_task("social")
    asyncio.create_task(run_model(task_id, VIDEO_MODELS["social"], data))
    return {"task_id": task_id}

@app.post("/video/cinematic")
async def video_cinematic(data: dict):
    task_id = create_task("cinematic")
    asyncio.create_task(run_model(task_id, VIDEO_MODELS["cinematic"], data))
    return {"task_id": task_id}

@app.post("/video/ads")
async def video_ads(data: dict):
    task_id = create_task("ads")
    asyncio.create_task(run_model(task_id, VIDEO_MODELS["ads"], data))
    return {"task_id": task_id}

@app.post("/video/music")
async def video_music(data: dict):
    task_id = create_task("music")
    asyncio.create_task(run_model(task_id, VIDEO_MODELS["music"], data))
    return {"task_id": task_id}

@app.post("/video/music_long")
async def video_music_long(data: dict):
    task_id = create_task("music_long")
    asyncio.create_task(run_model(task_id, VIDEO_MODELS["music_long"], data))
    return {"task_id": task_id}

@app.post("/video/experimental")
async def video_experimental(data: dict):
    task_id = create_task("experimental")
    asyncio.create_task(run_model(task_id, VIDEO_MODELS["experimental"], data))
    return {"task_id": task_id}