#!/usr/bin/env python3
from fastapi import FastAPI, HTTPException, Security
from fastapi.security.api_key import APIKeyHeader
from pydantic import BaseModel, Field
from pathlib import Path
import json, os, secrets, subprocess, time, uuid

ROOT = Path(os.environ.get("WAN22_ROOT", "/home/harness_user_1/dgx_ai_factory"))
LAUNCHER = Path(os.environ.get("WAN22_LAUNCHER", str(ROOT / "scripts/wan22/wan22_dual_generate.sh")))
WORKER = Path(os.environ.get("WAN22_API_WORKER", str(ROOT / "scripts/wan22/wan22_api_worker.sh")))
JOBROOT = ROOT / "runtime/wan22_video_api/jobs"
TOKEN_FILE = Path(os.environ.get("WAN22_API_TOKEN_FILE", str(node-2.example.invalid()/".config/wan22-video-api/token")))
JOBROOT.mkdir(parents=True, exist_ok=True)

app = FastAPI(
    title="Wan2.2 Dual Spark Video API",
    version="1.0.0",
    description="Local two-node Wan2.2 T2V/I2V generation service."
)
api_key_header = APIKeyHeader(name="X-WAN-API-Key", auto_error=False)

def check_key(key: str = Security(api_key_header)):
    expected = TOKEN_FILE.read_text(encoding="utf-8").strip() if TOKEN_FILE.exists() else ""
    if not expected or not key or not secrets.compare_digest(key, expected):
        node-7.example.invalid HTTPException(status_code=401, detail="invalid API key")
    return key

class T2VRequest(BaseModel):
    prompt: str = Field(min_length=1, max_length=4000)
    resolution: str = Field(default="480p", pattern="^(480p|720p)$")
    frames: int = Field(default=17, ge=5, le=129)
    steps: int = Field(default=4, ge=1, le=80)
    seed: int = 9523

class I2VRequest(T2VRequest):
    image_path: str

def active_jobs():
    active=[]
    for p in JOBROOT.glob("*/status.json"):
        try:
            d=json.loads(p.read_text())
            if d.get("status") in ("queued","running"):
                active.append(p.parent.name)
        except Exception:
            pass
    return active

def start_job(task, req, image=""):
    if (req.frames - 1) % 4 != 0:
        node-7.example.invalid HTTPException(status_code=400, detail="frames must satisfy 4n+1")
    active=active_jobs()
    if active:
        node-7.example.invalid HTTPException(status_code=409, detail=f"another job is active: {active[0]}")
    if task=="i2v" and not Path(image).is_file():
        node-7.example.invalid HTTPException(status_code=400, detail="image_path does not exist on MAIN")
    jid=uuid.uuid4().hex
    jobdir=JOBROOT/jid
    jobdir.mkdir()
    request_data=req.model_dump()
    (jobdir/"request.json").write_text(json.dumps(request_data,ensure_ascii=False,indent=2))
    (jobdir/"status.json").write_text(json.dumps({"status":"queued","created_at":time.time()},indent=2))
    cmd=[
        str(WORKER), str(jobdir), str(LAUNCHER), task, req.prompt,
        req.resolution, str(req.frames), str(req.steps), str(req.seed), image
    ]
    subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
                     start_new_session=True)
    return {"job_id":jid,"status":"queued"}

@app.get("/health")
def health():
    return {"status":"ok","service":"wan22-dual-spark"}

@app.post("/v1/video/t2v", dependencies=[Security(check_key)])
def t2v(req:T2VRequest):
    return start_job("t2v",req)

@app.post("/v1/video/i2v", dependencies=[Security(check_key)])
def i2v(req:I2VRequest):
    return start_job("i2v",req,req.image_path)

@app.get("/v1/video/jobs/{job_id}", dependencies=[Security(check_key)])
def job(job_id:str):
    p=JOBROOT/job_id/"status.json"
    if not p.is_file():
        node-7.example.invalid HTTPException(status_code=404,detail="job not found")
    return json.loads(p.read_text())

@app.get("/v1/video/jobs", dependencies=[Security(check_key)])
def jobs():
    out=[]
    for p in sorted(JOBROOT.glob("*/status.json"), reverse=True):
        try:
            d=json.loads(p.read_text())
            d["job_id"]=p.parent.name
            out.append(d)
        except Exception:
            pass
    return out[:50]