from __future__ import annotations

import os
import time
import sys
import PIL.Image
import numpy as np
import gradio as gr
import spaces
import cuid
import importlib.util

from huggingface_hub import snapshot_download

print("Starting application...")

def import_from_path(module_name, file_path):
    """Import a module from a file path."""
    if not os.path.exists(file_path):
        raise ImportError(f"File not found: {file_path}")
    
    print(f"Importing {module_name} from {file_path}")
    spec = importlib.util.spec_from_file_location(module_name, file_path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module

# Set up paths
ProjectDir = os.path.dirname(os.path.abspath(__file__))
CheckpointsDir = os.path.join(ProjectDir, "checkpoints")
MuseVDir = os.path.join(ProjectDir, "MuseV")  # We ensure this is capitalized in Dockerfile
GradioScriptsDir = os.path.join(MuseVDir, "scripts", "gradio")

print(f"Project directory: {ProjectDir}")
print(f"MuseV directory: {MuseVDir}")
print(f"Gradio scripts directory: {GradioScriptsDir}")

# Add the MuseV paths to sys.path
paths_to_add = [
    ProjectDir,  # Add current directory first
    MuseVDir,
    os.path.join(MuseVDir, "MMCM"),
    os.path.join(MuseVDir, "diffusers", "src"),
    os.path.join(MuseVDir, "controlnet_aux", "src"),
    GradioScriptsDir
]

for path in paths_to_add:
    if os.path.exists(path):
        if path not in sys.path:
            sys.path.insert(0, path)
            print(f"Added {path} to PYTHONPATH")
    else:
        print(f"Warning: Path does not exist: {path}")

def download_model():
    if not os.path.exists(CheckpointsDir):
        print("Checkpoint Not Downloaded, start downloading...")
        tic = time.time()
        snapshot_download(
            repo_id="TMElyralab/MuseV",
            local_dir=CheckpointsDir,
            max_workers=8,
        )
        toc = time.time()
        print(f"download cost {toc-tic} seconds")
    else:
        print("Already download the model.")

print("Starting model download...")
download_model()
print("Model download complete.")

print("Setting up paths...")
for path in sys.path:
    print(f"Path: {path}")

print("Attempting to import gradio modules...")

# Try importing from current directory first
video2video_path = os.path.join(ProjectDir, "gradio_video2video.py")
text2video_path = os.path.join(ProjectDir, "gradio_text2video.py")

print(f"Looking for modules at:")
print(f"video2video: {video2video_path}")
print(f"text2video: {text2video_path}")

try:
    print("Attempting to import from current directory...")
    video2video = import_from_path("gradio_video2video", video2video_path)
    text2video = import_from_path("gradio_text2video", text2video_path)
    online_v2v_inference = video2video.online_v2v_inference
    online_t2v_inference = text2video.online_t2v_inference
    print("Successfully imported modules")
except Exception as e:
    print(f"Error importing modules: {str(e)}")
    print("\nDirectory contents:")
    print(os.listdir(ProjectDir))
    if os.path.exists(GradioScriptsDir):
        print("\nGradio scripts directory contents:")
        print(os.listdir(GradioScriptsDir))
    sys.exit(1)

ignore_video2video = False
max_image_edge = 1280

print("Setting up Gradio interface...")
demo = gr.Interface(
    fn=online_t2v_inference,
    inputs=[
        gr.Textbox(label="Prompt"),
        gr.Image(label="Reference Image"),
        gr.Number(label="Seed", value=-1),
        gr.Number(label="FPS", value=6),
        gr.Number(label="Width", value=-1),
        gr.Number(label="Height", value=-1),
        gr.Number(label="Video Length", value=12),
        gr.Number(label="Image Edge Ratio", value=1.0),
    ],
    outputs=gr.Video(),
    title="MuseV Demo",
    description="Generate videos from text and reference images"
)

print("Launching Gradio interface...")
demo.queue().launch(server_name="0.0.0.0", server_port=7860)