import os
import sys
import subprocess
import base64
import tempfile
import traceback

# Install latest diffusers (supports SkyReelsV2DiffusionForcingPipeline) + 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-long] torch version: {torch.__version__}", flush=True)
print(f"[video-long] CUDA available: {torch.cuda.is_available()}", flush=True)
if torch.cuda.is_available():
    print(f"[video-long] GPU: {torch.cuda.get_device_name(0)}", flush=True)
    print(f"[video-long] VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB", flush=True)

app = Flask(__name__)
pipe = None

def load_model():
    global pipe
    from diffusers import (
        AutoModel,
        SkyReelsV2DiffusionForcingPipeline,
        UniPCMultistepScheduler,
    )
    print("[video-long] Loading SkyReels-V2-DF-1.3B-540P...", flush=True)

    model_id = "Skywork/SkyReels-V2-DF-1.3B-540P-Diffusers"

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

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

    # Flow shift 8.0 for T2V
    pipe.scheduler = UniPCMultistepScheduler.from_config(
        pipe.scheduler.config, flow_shift=8.0
    )
    pipe.to("cuda")
    print("[video-long] Model loaded on CUDA!", flush=True)

@app.route("/health", methods=["GET"])
def health():
    if pipe is not None:
        return jsonify({"status": "healthy"}), 200
    return jsonify({"status": "loading"}), 503

@app.route("/debug", methods=["GET"])
def debug():
    if pipe is None:
        return jsonify({"status": "loading"}), 503
    import inspect
    sig = inspect.signature(pipe.__call__)
    params = {}
    for name, param in sig.parameters.items():
        params[name] = str(param.default) if param.default != inspect.Parameter.empty else "required"
    return jsonify({
        "pipeline_class": type(pipe).__name__,
        "parameters": params,
        "device": str(pipe.device) if hasattr(pipe, 'device') else 'unknown',
    })

@app.route("/", methods=["POST"])
def generate():
    try:
        from diffusers.utils import export_to_video

        data = request.get_json(force=True)
        # Support both HF format (inputs) and backend format (prompt)
        prompt = data.get("inputs", "") or data.get("prompt", "")
        params = data.get("parameters", {})

        # Diffusion Forcing params - defaults for ~10s video
        num_frames = int(params.get("num_frames", 257))       # 257=10s, 737=30s, 1457=60s
        height = int(params.get("height", 544))                # 544 for 540P
        width = int(params.get("width", 960))                  # 960 for 540P
        steps = int(params.get("num_inference_steps", 30))
        fps = int(params.get("fps", 24))
        base_num_frames = int(params.get("base_num_frames", 97))
        ar_step = int(params.get("ar_step", 5))                # 5 for async (better quality, faster for long)
        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", 5.0))

        print(f"[video-long] Generating {num_frames} frames ({num_frames/fps:.1f}s): {prompt[:100]}", flush=True)
        print(f"[video-long] Params: {width}x{height}, steps={steps}, ar_step={ar_step}", flush=True)

        # Build kwargs - only pass causal_block_size when ar_step > 0
        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,
        )
        if ar_step > 0:
            call_kwargs["causal_block_size"] = causal_block_size

        print(f"[video-long] Pipeline call kwargs: {call_kwargs}", flush=True)
        output = pipe(**call_kwargs).frames[0]

        # Write to temp file
        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
        print(f"[video-long] Done: {len(output)} frames, {duration:.1f}s, {len(video_bytes)} bytes", flush=True)

        return jsonify({
            "video": video_b64,
            "format": "mp4",
            "frames": len(output),
            "duration": duration,
            "success": True,
            "video_base64": video_b64,
            "duration_seconds": duration,
            "model": "skyreels-v2-df-1.3b",
        })
    except Exception as e:
        tb = traceback.format_exc()
        print(f"[video-long] ERROR: {e}", flush=True)
        print(f"[video-long] TRACEBACK:\n{tb}", flush=True)
        return jsonify({"error": str(e), "traceback": tb, "success": False}), 200

# Also add /api/generate-video route for backend compatibility
@app.route("/api/generate-video", methods=["POST"])
def api_generate_video():
    """Handle the backend's /api/generate-video format by forwarding to generate()"""
    return generate()

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