import os
import subprocess
from datetime import datetime
from pathlib import Path
import gradio as gr

# -----------------------------
# Setup paths and env
# -----------------------------
HF_HOME = "/app/hf_cache"
os.environ["HF_HOME"] = HF_HOME
os.environ["TRANSFORMERS_CACHE"] = HF_HOME
os.makedirs(HF_HOME, exist_ok=True)

PRETRAINED_DIR = "/app/pretrained"
os.makedirs(PRETRAINED_DIR, exist_ok=True)


# -----------------------------
# Step 1: Optional Model Download
# -----------------------------
def download_models():
    expected_model = os.path.join(PRETRAINED_DIR, "RAFT/raft-things.pth")
    if not Path(expected_model).exists():
        print("⚙️ Downloading pretrained models...")
        try:
            subprocess.check_call(["bash", "download/download_models.sh"])
            print("✅ Models downloaded.")
        except subprocess.CalledProcessError as e:
            print(f"❌ Model download failed: {e}")
    else:
        print("✅ Pretrained models already exist.")


download_models()


# -----------------------------
# Step 2: Inference Logic
# -----------------------------
def run_epic_inference(video_path, caption, motion_type):
    temp_input_path = "/app/temp_input.mp4"
    output_dir = f"/app/output_{motion_type}"
    traj_name = motion_type
    traj_txt = f"/app/inference/v2v_data/test/trajs/{traj_name}.txt"

    # Save uploaded video
    if video_path:
        os.system(f"cp '{video_path}' {temp_input_path}")
    
    command = [
    "python",
    "/app/inference/v2v_data/inference.py",
    "--video_path",
    temp_input_path,
    "--stride",
    "1",
    "--out_dir",
    output_dir,
    "--radius_scale",
    "1",
    "--camera",
    "traj",
    "--mask",
    "--target_pose", "0", "30", "-0.6", "0", "0",
    "--traj_txt",
    traj_txt,
    "--save_name",
    f"amalfi-coast_traj_{traj_name}",
    "--mode",
    "gradual",
    "--out_dir",
    output_dir,
    ]

    # Run inference command
    try:
        result = subprocess.run(command, capture_output=True, text=True, check=True)
        logs = result.stdout
    except subprocess.CalledProcessError as e:
        logs = f"❌ Inference failed:\n{e.stderr}"
        return logs, None

    # Locate the output video
    output_video = Path(output_dir) / f"amalfi-coast_traj_{traj_name}.mp4"
    if output_video.exists():
        return logs, str(output_video)
    else:
        return f"Inference succeeded but no output video found in {output_dir}", None


# -----------------------------
# Step 3: Create Gradio UI
# -----------------------------
demo = gr.Interface(
    fn=run_epic_inference,
    inputs=[
        gr.Video(label="Upload Video (MP4)"),
        gr.Textbox(label="Caption", placeholder="e.g., Amalfi coast with boats"),
        gr.Dropdown(
            choices=["zoom_in", "rotate", "orbit", "pan", "loop1"],
            label="Camera Motion Type",
            value="zoom_in",
        ),
    ],
    outputs=[gr.Textbox(label="Inference Logs"), gr.Video(label="Generated Video")],
    title="🎬 EPiC: Efficient Video Camera Control",
    description="Upload a video, describe the scene, and apply cinematic camera motion using pretrained EPiC models.",
)

# -----------------------------
# Step 4: Launch App
# -----------------------------
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)