"""
SkyReels-V2-DF-14B - UNLIMITED LENGTH Highest Quality Video Generation
Job Queue Pattern: POST returns job_id, poll /status/{job_id}, GET /result/{job_id}
14B model on A100 80GB - supports 540P and 720P, videos up to 60s+
"""
import os
import sys
import subprocess
import base64
import tempfile
import traceback
import threading
import time
import uuid

# Install dependencies
subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "--upgrade",
    "diffusers", "flask", "accelerate", "sentencepiece", "protobuf", "imageio[ffmpeg]",
    "transformers", "huggingface_hub", "ftfy", "einops"])

import torch
from flask import Flask, request, jsonify

print(f"[video-14b] torch: {torch.__version__}, CUDA: {torch.cuda.is_available()}", flush=True)
if torch.cuda.is_available():
    print(f"[video-14b] GPU: {torch.cuda.get_device_name(0)}", flush=True)
    print(f"[video-14b] VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB", flush=True)

app = Flask(__name__)
pipe = None
jobs = {}

# Model selection - default 540P, can be overridden with MODEL_RESOLUTION env var
RESOLUTION = os.environ.get("MODEL_RESOLUTION", "540P")
if RESOLUTION == "720P":
    MODEL_ID = "Skywork/SkyReels-V2-DF-14B-720P-Diffusers"
    DEFAULT_HEIGHT = 720
    DEFAULT_WIDTH = 1280
    DEFAULT_BASE_FRAMES = 121
else:
    MODEL_ID = "Skywork/SkyReels-V2-DF-14B-540P-Diffusers"
    DEFAULT_HEIGHT = 544
    DEFAULT_WIDTH = 960
    DEFAULT_BASE_FRAMES = 97

print(f"[video-14b] Using model: {MODEL_ID} ({RESOLUTION})", flush=True)


def load_model():
    global pipe
    from diffusers import (
        AutoModel,
        SkyReelsV2DiffusionForcingPipeline,
        UniPCMultistepScheduler,
    )
    print(f"[video-14b] Loading {MODEL_ID}...", flush=True)

    vae = AutoModel.from_pretrained(
        MODEL_ID,
        subfolder="vae",
        torch_dtype=torch.float32,
    )

    pipe = SkyReelsV2DiffusionForcingPipeline.from_pretrained(
        MODEL_ID,
        vae=vae,
        torch_dtype=torch.bfloat16,
    )

    pipe.scheduler = UniPCMultistepScheduler.from_config(
        pipe.scheduler.config, flow_shift=8.0
    )
    pipe.to("cuda")
    print(f"[video-14b] Model loaded on CUDA! ({RESOLUTION})", flush=True)


def generate_video_background(job_id, prompt, params):
    """Run video generation in a background thread."""
    global jobs
    try:
        from diffusers.utils import export_to_video

        jobs[job_id]["status"] = "generating"
        jobs[job_id]["progress"] = "Starting generation..."

        num_frames = int(params.get("num_frames", 257))
        height = int(params.get("height", DEFAULT_HEIGHT))
        width = int(params.get("width", DEFAULT_WIDTH))
        steps = int(params.get("num_inference_steps", 10))
        fps = int(params.get("fps", 24))
        base_num_frames = int(params.get("base_num_frames", DEFAULT_BASE_FRAMES))
        ar_step = int(params.get("ar_step", 5))
        causal_block_size = int(params.get("causal_block_size", 5))
        overlap_history = int(params.get("overlap_history", 17))
        addnoise_condition = int(params.get("addnoise_condition", 20))
        guidance_scale = float(params.get("guidance_scale", 6.0))

        duration_sec = num_frames / fps
        jobs[job_id]["progress"] = f"Generating {num_frames} frames ({duration_sec:.1f}s) at {RESOLUTION}..."
        print(f"[video-14b] Job {job_id}: {num_frames} frames ({duration_sec:.1f}s), {prompt[:80]}", flush=True)

        call_kwargs = dict(
            prompt=prompt,
            num_inference_steps=steps,
            height=height,
            width=width,
            num_frames=num_frames,
            base_num_frames=base_num_frames,
            ar_step=ar_step,
            overlap_history=overlap_history,
            addnoise_condition=addnoise_condition,
            guidance_scale=guidance_scale,
        )
        if ar_step > 0:
            call_kwargs["causal_block_size"] = causal_block_size

        output = pipe(**call_kwargs).frames[0]

        jobs[job_id]["progress"] = "Encoding video..."

        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp:
            tmp_path = tmp.name
        export_to_video(output, tmp_path, fps=fps)

        with open(tmp_path, "rb") as f:
            video_bytes = f.read()
        os.unlink(tmp_path)

        video_b64 = base64.b64encode(video_bytes).decode("utf-8")
        duration = len(output) / fps

        jobs[job_id]["status"] = "completed"
        jobs[job_id]["video_b64"] = video_b64
        jobs[job_id]["frames"] = len(output)
        jobs[job_id]["duration"] = duration
        jobs[job_id]["progress"] = f"Done: {len(output)} frames, {duration:.1f}s"
        print(f"[video-14b] Job {job_id} completed: {len(output)} frames, {duration:.1f}s", flush=True)

    except Exception as e:
        tb = traceback.format_exc()
        jobs[job_id]["status"] = "failed"
        jobs[job_id]["error"] = str(e)
        jobs[job_id]["traceback"] = tb
        print(f"[video-14b] Job {job_id} FAILED: {e}", flush=True)
        print(f"[video-14b] Traceback:\n{tb}", flush=True)


@app.route("/health", methods=["GET"])
def health():
    if pipe is not None:
        return jsonify({"status": "healthy", "model": "skyreels-v2-df-14b", "resolution": RESOLUTION}), 200
    return jsonify({"status": "loading"}), 503


@app.route("/", methods=["POST"])
@app.route("/api/generate-video", methods=["POST"])
def submit_video_job():
    """Submit a video generation job. Returns job_id immediately."""
    try:
        data = request.get_json(force=True)
        prompt = data.get("inputs", "") or data.get("prompt", "")
        params = data.get("parameters", {})

        if not prompt:
            return jsonify({"error": "Prompt is required", "success": False}), 400

        job_id = str(uuid.uuid4())[:8]
        jobs[job_id] = {
            "status": "queued",
            "progress": "Queued for generation",
            "video_b64": None,
            "frames": None,
            "duration": None,
            "error": None,
            "prompt": prompt,
            "params": params,
            "created_at": time.time(),
        }

        thread = threading.Thread(
            target=generate_video_background,
            args=(job_id, prompt, params),
            daemon=True
        )
        thread.start()

        print(f"[video-14b] Job {job_id} submitted: {prompt[:80]}", flush=True)

        return jsonify({
            "success": True,
            "job_id": job_id,
            "status": "queued",
            "message": f"Video generation started. Poll /status/{job_id} for progress.",
            "status_url": f"/status/{job_id}",
            "result_url": f"/result/{job_id}",
            "model": "skyreels-v2-df-14b",
            "resolution": RESOLUTION,
        })

    except Exception as e:
        return jsonify({"error": str(e), "success": False}), 500


@app.route("/status/<job_id>", methods=["GET"])
def job_status(job_id):
    if job_id not in jobs:
        return jsonify({"error": "Job not found", "success": False}), 404

    job = jobs[job_id]
    return jsonify({
        "success": True,
        "job_id": job_id,
        "status": job["status"],
        "progress": job["progress"],
        "frames": job["frames"],
        "duration": job["duration"],
        "error": job["error"],
    })


@app.route("/result/<job_id>", methods=["GET"])
def job_result(job_id):
    if job_id not in jobs:
        return jsonify({"error": "Job not found", "success": False}), 404

    job = jobs[job_id]
    if job["status"] != "completed":
        return jsonify({
            "success": False,
            "status": job["status"],
            "progress": job["progress"],
            "message": "Video not ready yet. Keep polling /status/" + job_id,
        }), 202

    return jsonify({
        "success": True,
        "job_id": job_id,
        "video": job["video_b64"],
        "video_base64": job["video_b64"],
        "format": "mp4",
        "frames": job["frames"],
        "duration": job["duration"],
        "duration_seconds": job["duration"],
        "model": "skyreels-v2-df-14b",
        "resolution": RESOLUTION,
    })


@app.route("/debug", methods=["GET"])
def debug():
    return jsonify({
        "pipe_loaded": pipe is not None,
        "model": "skyreels-v2-df-14b",
        "resolution": RESOLUTION,
        "model_id": MODEL_ID,
        "active_jobs": {k: {"status": v["status"], "progress": v["progress"]} for k, v in jobs.items()},
        "total_jobs": len(jobs),
    })


if __name__ == "__main__":
    load_model()
    app.run(host="0.0.0.0", port=8000, threaded=True)