import subprocess
import shutil
import os
import sys
from pathlib import Path

BASE = Path(__file__).resolve().parent
CKPT_VACE = BASE / "Wan2.1-VACE-1.3B"
CKPT_FLF = BASE / "Wan2.1-FLF2V-14B-720P"

def run_cmd(cmd):
    """Run a command and handle errors"""
    try:
        result = subprocess.run(cmd, check=True, capture_output=True, text=True)
        return result.stdout
    except subprocess.CalledProcessError as e:
        print(f"Command failed: {' '.join(cmd)}")
        print(f"Error: {e.stderr}")
        raise

def check_wan_installation():
    """Check if Wan2.1 and VACE are properly installed"""
    try:
        # Check if the repositories exist locally
        wan_path = BASE / "Wan2.1"
        vace_path = BASE / "VACE"
        
        if not wan_path.exists():
            print("Cloning Wan2.1 repository...")
            run_cmd(["git", "clone", "https://github.com/Wan-Video/Wan2.1.git", str(wan_path)])
        
        if not vace_path.exists():
            print("Cloning VACE repository...")
            run_cmd(["git", "clone", "https://github.com/ali-vilab/VACE.git", str(vace_path)])
        
        # Add to Python path
        sys.path.insert(0, str(wan_path))
        sys.path.insert(0, str(vace_path))
        
        return True
    except Exception as e:
        print(f"Failed to setup Wan2.1/VACE: {e}")
        return False

def generate_video(ref_img, first_frame, last_frame, output_path):
    """Generate video from reference image and first/last frames"""
    if not check_wan_installation():
        raise RuntimeError("Wan2.1/VACE setup failed")
    
    # Create necessary directories
    os.makedirs("processed", exist_ok=True)
    os.makedirs("results", exist_ok=True)
    
    # Preprocess frames
    vace_script = BASE / "VACE" / "vace_preprocess.py"
    if vace_script.exists():
        run_cmd([
            sys.executable, str(vace_script),
            "--task", "extension",
            "--mode", "firstlastframe",
            "--image", f"{first_frame},{last_frame}",
            "--expand_num", "240"
        ])
    else:
        # Fallback: create a simple template video
        print("VACE preprocessing not available, using fallback method")
        # This would need actual implementation
    
    # Generate video
    generate_script = BASE / "Wan2.1" / "generate.py"
    if generate_script.exists():
        run_cmd([
            sys.executable, str(generate_script),
            "--task", "flf2v-14B",
            "--size", "1280*720",
            "--ckpt_dir", str(CKPT_FLF),
            "--src_ref_images", ref_img,
            "--src_video", "processed/src_video.mp4",
            "--fps", "24",
            "--offload_model", "True",
            "--t5_cpu"
        ])
        
        # Move result
        if os.path.exists("results/final.mp4"):
            shutil.move("results/final.mp4", output_path)
    else:
        raise FileNotFoundError("Wan2.1 generate.py not found")

def generate_image(ref_img, prompt, output_path):
    """Generate image from reference image and prompt"""
    if not check_wan_installation():
        raise RuntimeError("Wan2.1/VACE setup failed")
    
    # Create necessary directories
    os.makedirs("results", exist_ok=True)
    
    # Generate image
    generate_script = BASE / "Wan2.1" / "generate.py"
    if generate_script.exists():
        run_cmd([
            sys.executable, str(generate_script),
            "--task", "vace-1.3B",
            "--size", "832*480",
            "--ckpt_dir", str(CKPT_VACE),
            "--src_ref_images", ref_img,
            "--frame_num", "1",
            "--prompt", prompt or " "
        ])
        
        # Move result
        if os.path.exists("results/final_frame000.png"):
            shutil.move("results/final_frame000.png", output_path)
    else:
        raise FileNotFoundError("Wan2.1 generate.py not found")