import os
import sys
import spaces
import gradio as gr
import torch
import argparse
from PIL import Image
import numpy as np
import torchvision.transforms as transforms
from moviepy.editor import VideoFileClip
from diffusers.utils import load_image, load_video
from tqdm import tqdm
from image_gen_aux import DepthPreprocessor

project_root = os.path.dirname(os.path.abspath(__file__))
os.environ["GRADIO_TEMP_DIR"] = os.path.join(project_root, "tmp", "gradio")
sys.path.append(project_root)

try:
    sys.path.append(os.path.join(project_root, "submodules/MoGe"))
    sys.path.append(os.path.join(project_root, "submodules/vggt"))
    os.environ["TOKENIZERS_PARALLELISM"] = "false"
except:
    print("Warning: MoGe not found, motion transfer will not be applied")

HERE_PATH = os.path.normpath(os.path.dirname(__file__))
sys.path.insert(0, HERE_PATH)
from huggingface_hub import hf_hub_download
hf_hub_download(repo_id="EXCAI/Diffusion-As-Shader", filename='spatracker/spaT_final.pth', local_dir=f'{HERE_PATH}/checkpoints/')

from models.pipelines import DiffusionAsShaderPipeline, FirstFrameRepainter, CameraMotionGenerator, ObjectMotionGenerator
from submodules.MoGe.moge.model import MoGeModel
from submodules.vggt.vggt.utils.pose_enc import pose_encoding_to_extri_intri
from submodules.vggt.vggt.models.vggt import VGGT

import torch._dynamo
torch._dynamo.config.suppress_errors = True

# Parse command line arguments
parser = argparse.ArgumentParser(description="Diffusion as Shader Web UI")
parser.add_argument("--port", type=int, default=7860, help="Port to run the web UI on")
parser.add_argument("--share", action="store_true", help="Share the web UI")
parser.add_argument("--gpu", type=int, default=0, help="GPU device ID")
parser.add_argument("--model_path", type=str, default="EXCAI/Diffusion-As-Shader", help="Path to model checkpoint")
parser.add_argument("--output_dir", type=str, default="tmp", help="Output directory")
args = parser.parse_args()

# Use the original GPU ID throughout the entire code for consistency
GPU_ID = args.gpu
DEFAULT_MODEL_PATH = args.model_path
OUTPUT_DIR = args.output_dir

# Create necessary directories
os.makedirs("outputs", exist_ok=True)
# Create project tmp directory instead of using system temp
os.makedirs(os.path.join(project_root, "tmp"), exist_ok=True)
os.makedirs(os.path.join(project_root, "tmp", "gradio"), exist_ok=True)

def load_media(media_path, max_frames=49, transform=None):
    """Load video or image frames and convert to tensor
    
    Args:
        media_path (str): Path to video or image file
        max_frames (int): Maximum number of frames to load
        transform (callable): Transform to apply to frames
        
    Returns:
        Tuple[torch.Tensor, float, bool]: Video tensor [T,C,H,W], FPS, and is_video flag
    """
    if transform is None:
        transform = transforms.Compose([
            transforms.Resize((480, 720)),
            transforms.ToTensor()
        ])
    
    # Determine if input is video or image based on extension
    ext = os.path.splitext(media_path)[1].lower()
    is_video = ext in ['.mp4', '.avi', '.mov']
    
    if is_video:
        # Load video file info
        video_clip = VideoFileClip(media_path)
        duration = video_clip.duration
        original_fps = video_clip.fps
        
        # Case 1: Video longer than 6 seconds, sample first 6 seconds + 1 frame
        if duration > 6.0:
            # 使用 max_frames 参数而不是 sampling_fps
            frames = load_video(media_path, max_frames=max_frames)
            fps = max_frames / 6.0  # 计算等效的 fps
        # Cases 2 and 3: Video shorter than 6 seconds
        else:
            # Load all frames
            frames = load_video(media_path)
            
            # Case 2: Total frames less than max_frames, need interpolation
            if len(frames) < max_frames:
                fps = len(frames) / duration  # Keep original fps
                
                # Evenly interpolate to max_frames
                indices = np.linspace(0, len(frames) - 1, max_frames)
                new_frames = []
                for i in indices:
                    idx = int(i)
                    new_frames.append(frames[idx])
                frames = new_frames
            # Case 3: Total frames more than max_frames but video less than 6 seconds
            else:
                # Evenly sample to max_frames
                indices = np.linspace(0, len(frames) - 1, max_frames)
                new_frames = []
                for i in indices:
                    idx = int(i)
                    new_frames.append(frames[idx])
                frames = new_frames
                fps = max_frames / duration  # New fps to maintain duration
    else:
        # Handle image as single frame
        image = load_image(media_path)
        frames = [image]
        fps = 8  # Default fps for images
        
        # Duplicate frame to max_frames
        while len(frames) < max_frames:
            frames.append(frames[0].copy())
    
    # Convert frames to tensor
    video_tensor = torch.stack([transform(frame) for frame in frames])
    
    return video_tensor, fps, is_video

def save_uploaded_file(file):
    if file is None:
        return None
        
    # Use project tmp directory instead of system temp
    temp_dir = os.path.join(project_root, "tmp")
    
    if hasattr(file, 'name'):
        filename = file.name
    else:
        # Generate a unique filename if name attribute is missing
        import uuid
        ext = ".tmp"
        if hasattr(file, 'content_type'):
            if "image" in file.content_type:
                ext = ".png"
            elif "video" in file.content_type:
                ext = ".mp4"
        filename = f"{uuid.uuid4()}{ext}"
    
    temp_path = os.path.join(temp_dir, filename)
    
    try:
        # Check if file is a FileStorage object or already a path
        if hasattr(file, 'save'):
            file.save(temp_path)
        elif isinstance(file, str):
            # It's already a path
            return file
        else:
            # Try to read and save the file
            with open(temp_path, 'wb') as f:
                f.write(file.read() if hasattr(file, 'read') else file)
    except Exception as e:
        print(f"Error saving file: {e}")
        return None
        
    return temp_path

das_pipeline = None
moge_model = None
vggt_model = None

@spaces.GPU
def get_das_pipeline():
    global das_pipeline
    if das_pipeline is None:
        das_pipeline = DiffusionAsShaderPipeline(gpu_id=GPU_ID, output_dir=OUTPUT_DIR)
    return das_pipeline

@spaces.GPU
def get_moge_model():
    global moge_model
    if moge_model is None:
        das = get_das_pipeline()
        moge_model = MoGeModel.from_pretrained("Ruicheng/moge-vitl").to(das.device)
    return moge_model

@spaces.GPU
def get_vggt_model():
    global vggt_model
    if vggt_model is None:
        das = get_das_pipeline()
        vggt_model = VGGT.from_pretrained("facebook/VGGT-1B").to(das.device)
    return vggt_model

def process_motion_transfer(source, prompt, mt_repaint_option, mt_repaint_image):
    """Process video motion transfer task"""
    try:
        # 保存上传的文件
        input_video_path = save_uploaded_file(source)
        if input_video_path is None:
            return None, None, None, None, None
        
        print(f"DEBUG: Repaint option: {mt_repaint_option}")
        print(f"DEBUG: Repaint image: {mt_repaint_image}")
        
        das = get_das_pipeline()
        video_tensor, fps, is_video = load_media(input_video_path)
        das.fps = fps  # 设置 das.fps 为 load_media 返回的 fps
        
        if not is_video:
            tracking_method = "moge"
            print("Image input detected, using MoGe for tracking video generation.")
        else:
            tracking_method = "cotracker"
        
        repaint_img_tensor = None
        if mt_repaint_image is not None:
            repaint_path = save_uploaded_file(mt_repaint_image)
            repaint_img_tensor, _, _ = load_media(repaint_path)
            repaint_img_tensor = repaint_img_tensor[0]
        elif mt_repaint_option == "Yes":
            repainter = FirstFrameRepainter(gpu_id=GPU_ID, output_dir=OUTPUT_DIR)
            repaint_img_tensor = repainter.repaint(
                video_tensor[0], 
                prompt=prompt,
                depth_path=None
            )
        
        tracking_tensor = None
        tracking_path = None
        if tracking_method == "moge":
            moge = get_moge_model()
            infer_result = moge.infer(video_tensor[0].to(das.device))  # [C, H, W] in range [0,1]
            H, W = infer_result["points"].shape[0:2]
            pred_tracks = infer_result["points"].unsqueeze(0).repeat(49, 1, 1, 1) #[T, H, W, 3]
            poses = torch.eye(4).unsqueeze(0).repeat(49, 1, 1)

            pred_tracks_flatten = pred_tracks.reshape(video_tensor.shape[0], H*W, 3)
  
            cam_motion = CameraMotionGenerator(None)
            cam_motion.set_intr(infer_result["intrinsics"])
  
            pred_tracks = cam_motion.w2s(pred_tracks_flatten, poses).reshape([video_tensor.shape[0], H, W, 3]) # [T, H, W, 3]

            tracking_path, tracking_tensor = das.visualize_tracking_moge(
                pred_tracks.cpu().numpy(), 
                infer_result["mask"].cpu().numpy()
            )
            print('Export tracking video via MoGe')
        else:
            # 使用 cotracker
            pred_tracks, pred_visibility = generate_tracking_cotracker(video_tensor)
            tracking_path, tracking_tensor = das.visualize_tracking_cotracker(pred_tracks, pred_visibility)
            print('Export tracking video via cotracker')
        
        return tracking_path, video_tensor, tracking_tensor, repaint_img_tensor, fps
    except Exception as e:
        import traceback
        print(f"Processing failed: {str(e)}\n{traceback.format_exc()}")
        return None, None, None, None, None
        
def generate_tracking_cotracker(video_tensor, density=30):
    """在CPU上生成跟踪视频，只使用第一帧的深度信息，使用矩阵运算提高效率
    
    参数:
        video_tensor (torch.Tensor): 输入视频张量
        density (int): 跟踪点的密度
        
    返回:
        tuple: (pred_tracks, pred_visibility)
    """
    cotracker = torch.hub.load("facebookresearch/co-tracker", "cotracker3_offline").to("cpu")
    depth_preprocessor = DepthPreprocessor.from_pretrained("Intel/zoedepth-nyu-kitti").to("cpu")
    
    video = video_tensor.unsqueeze(0).to("cpu")
    
    # 只处理第一帧以获取深度图
    print("estimating depth for first frame...")
    frame = (video_tensor[0].permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
    depth = depth_preprocessor(Image.fromarray(frame))[0]
    depth_tensor = transforms.ToTensor()(depth)  # [1, H, W]
    
    # 获取跟踪点和可见性
    print("tracking on CPU...")
    pred_tracks, pred_visibility = cotracker(video, grid_size=density)  # B T N 2,  B T N 1
    
    # 提取维度
    B, T, N, _ = pred_tracks.shape
    H, W = depth_tensor.shape[1], depth_tensor.shape[2]
    
    # 创建带深度的输出张量
    pred_tracks_with_depth = torch.zeros((B, T, N, 3), device="cpu")
    pred_tracks_with_depth[:, :, :, :2] = pred_tracks  # 复制x,y坐标
    
    # 使用矩阵运算一次性处理所有帧和点
    # 重塑pred_tracks为[B*T*N, 2]以便于处理
    flat_tracks = pred_tracks.reshape(-1, 2)
    
    # 将坐标限制在有效图像边界内
    x_coords = flat_tracks[:, 0].clamp(0, W-1).long()
    y_coords = flat_tracks[:, 1].clamp(0, H-1).long()
    
    # 从第一帧的深度图获取所有点的深度值
    depths = depth_tensor[0, y_coords, x_coords]
    
    # 重塑回原始形状并分配给输出张量
    pred_tracks_with_depth[:, :, :, 2] = depths.reshape(B, T, N)

    del cotracker,depth_preprocessor
    
    # 将结果返回
    return pred_tracks_with_depth.squeeze(0), pred_visibility.squeeze(0)

@spaces.GPU(duration=350)
def apply_tracking_unified(video_tensor, tracking_tensor, repaint_img_tensor, prompt, fps):
    """统一的应用跟踪函数"""
    print("--- Entering apply_tracking_unified ---")
    print(f"Prompt received: {prompt}")
    print(f"FPS received: {fps}")
    print(f"Video tensor shape: {video_tensor.shape if video_tensor is not None else None}")
    print(f"Tracking tensor shape: {tracking_tensor.shape if tracking_tensor is not None else None}")
    print(f"Repaint tensor shape: {repaint_img_tensor.shape if repaint_img_tensor is not None else None}")
    try:
        if video_tensor is None or tracking_tensor is None:
            print("Error: Video tensor or tracking tensor is None.")
            return None

        das = get_das_pipeline()
        output_path = das.apply_tracking(
            video_tensor=video_tensor,
            fps=fps,
            tracking_tensor=tracking_tensor,
            img_cond_tensor=repaint_img_tensor,
            prompt=prompt,
            checkpoint_path=DEFAULT_MODEL_PATH,
            num_inference_steps=15
        )

        print(f"das.apply_tracking returned: {output_path}")

        # --- 临时解决方案开始 ---
        # 检查 das.apply_tracking 是否返回 None，并尝试使用日志中看到的固定路径
        potential_fixed_path = os.path.join(project_root, OUTPUT_DIR, "result.mp4") # 构建预期的固定路径
        print(f"Checking potential fixed path: {potential_fixed_path}")

        if output_path is None and os.path.exists(potential_fixed_path):
            print(f"Warning: das.apply_tracking returned None, but found file at {potential_fixed_path}. Using this path.")
            output_path = potential_fixed_path
        # --- 临时解决方案结束 ---

        print(f"最终使用的视频路径: {output_path}")

        # 确保返回的是绝对路径
        if output_path and not os.path.isabs(output_path):
            output_path = os.path.abspath(output_path)

        # 检查文件是否存在
        if output_path and os.path.exists(output_path):
            print(f"文件存在，大小: {os.path.getsize(output_path)} 字节")
            return output_path
        else:
            print(f"警告: 输出文件不存在或路径无效: {output_path}")
            return None
    except Exception as e:
        import traceback
        print(f"Apply tracking failed: {str(e)}\n{traceback.format_exc()}")
        return None

# 添加在 apply_tracking_unified 函数之后，Gradio 界面定义之前

def enable_apply_button(tracking_result):
    """当跟踪视频生成后启用应用按钮"""
    if tracking_result is not None:
        return gr.update(interactive=True)
    return gr.update(interactive=False)

@spaces.GPU
def process_vggt(video_tensor):
    vggt_model = get_vggt_model()
    
    t, c, h, w = video_tensor.shape
    new_width = 518
    new_height = round(h * (new_width / w) / 14) * 14
    resize_transform = transforms.Resize((new_height, new_width), interpolation=Image.BICUBIC)
    video_vggt = resize_transform(video_tensor)  # [T, C, H, W]
    
    if new_height > 518:
        start_y = (new_height - 518) // 2
        video_vggt = video_vggt[:, :, start_y:start_y + 518, :]

    with torch.no_grad():
        with torch.cuda.amp.autocast(dtype=torch.float16):
            video_vggt = video_vggt.unsqueeze(0)  # [1, T, C, H, W]
            aggregated_tokens_list, ps_idx = vggt_model.aggregator(video_vggt.to("cuda"))
        
            extr, intr = pose_encoding_to_extri_intri(vggt_model.camera_head(aggregated_tokens_list)[-1], video_vggt.shape[-2:])
    
    return extr, intr 

def load_examples():
    """加载示例文件路径"""
    samples_dir = os.path.join(project_root, "samples")
    if not os.path.exists(samples_dir):
        print(f"Warning: Samples directory not found at {samples_dir}")
        return []
    
    examples_list = []
    
    # 为每个示例集创建一个示例项
    # 示例1
    example1 = [None] * 5  # [source, repaint_image, prompt, tracking_video, result_video]
    for filename in os.listdir(samples_dir):
        if filename.startswith("sample1_"):
            if filename.endswith("_raw.mp4"):
                example1[0] = os.path.join(samples_dir, filename)
            elif filename.endswith("_repaint.png"):
                example1[1] = os.path.join(samples_dir, filename)
            elif filename.endswith("_tracking.mp4"):
                example1[3] = os.path.join(samples_dir, filename)
            elif filename.endswith("_result.mp4"):
                example1[4] = os.path.join(samples_dir, filename)
    
    # 设置示例1的提示文本
    example1[2] = "A wonderful bright old-fasion red car is riding from left to right sun light is shining on the car, its reflection glittering. In the background is a deserted city in the noon, the roads and buildings are covered with green vegetation."
    
    
    # 示例2
    example2 = [None] * 5  # [source, repaint_image, prompt, tracking_video, result_video]
    for filename in os.listdir(samples_dir):
        if filename.startswith("sample2_"):
            if filename.endswith("_raw.mp4"):
                example2[0] = os.path.join(samples_dir, filename)
            elif filename.endswith("_repaint.png"):
                example2[1] = os.path.join(samples_dir, filename)
            elif filename.endswith("_tracking.mp4"):
                example2[3] = os.path.join(samples_dir, filename)
            elif filename.endswith("_result.mp4"):
                example2[4] = os.path.join(samples_dir, filename)
    
    # 设置示例2的提示文本
    example2[2] = "a rocket lifts off from the table and smoke erupt from its bottom."
    
    # 添加示例到列表
    if example1[0] is not None and example1[3] is not None:
        examples_list.append(example1)
    
    if example2[0] is not None and example2[3] is not None:
        examples_list.append(example2)
    
    # 添加其他示例（如果有）
    sample_prefixes = set()
    for filename in os.listdir(samples_dir):
        if filename.endswith(('.mp4', '.png')):
            prefix = filename.split('_')[0]
            if prefix not in ["sample1", "sample2"]:
                sample_prefixes.add(prefix)
    
    for prefix in sorted(sample_prefixes):
        example = [None] * 5  # [source, repaint_image, prompt, tracking_video, result_video]
        for filename in os.listdir(samples_dir):
            if filename.startswith(f"{prefix}_"):
                if filename.endswith("_raw.mp4"):
                    example[0] = os.path.join(samples_dir, filename)
                elif filename.endswith("_repaint.png"):
                    example[1] = os.path.join(samples_dir, filename)
                elif filename.endswith("_tracking.mp4"):
                    example[3] = os.path.join(samples_dir, filename)
                elif filename.endswith("_result.mp4"):
                    example[4] = os.path.join(samples_dir, filename)
        
        # 添加默认提示文本
        example[2] = "A beautiful scene"
        
        # 只有当至少有源文件和跟踪视频时才添加示例
        if example[0] is not None and example[3] is not None:
            examples_list.append(example)
    
    return examples_list

# Create Gradio interface with updated layout
with gr.Blocks(title="Diffusion as Shader") as demo:
    gr.Markdown("# Diffusion as Shader Web UI")
    gr.Markdown("### [Project Page](https://igl-hkust.github.io/das/) | [GitHub](https://github.com/IGL-HKUST/DiffusionAsShader)")

    # 创建隐藏状态变量来存储中间结果
    video_tensor_state = gr.State(None)
    tracking_tensor_state = gr.State(None)
    repaint_img_tensor_state = gr.State(None)
    fps_state = gr.State(None)

    with gr.Row():
        left_column = gr.Column(scale=1)
        right_column = gr.Column(scale=1)

    with left_column:
        gr.Markdown("### 1. Upload Source")
        gr.Markdown("Upload a video, We will extract the motion from it")
        source_preview = gr.Video(label="Source Preview")
        source_upload = gr.UploadButton("Upload Source", file_types=["video"])

        def update_source_preview(file):
            if file is None:
                return None
            path = save_uploaded_file(file)
            return path
        
        source_upload.upload(
            fn=update_source_preview,
            inputs=[source_upload],
            outputs=[source_preview]
        )
        
        gr.Markdown("### 2. Enter the prompt")
        common_prompt = gr.Textbox(label="Describe the scene and the motion you want to create: ", lines=2)
        
        gr.Markdown("### 3. Select a task")
        with gr.Tabs() as task_tabs:
            # Motion Transfer tab
            with gr.TabItem("Motion Transfer"):
                
                gr.Markdown("#### 3.1 Process the first frame of Source")
                gr.Markdown("DaS can produce novel videos while maintaining the features of the first frame and all the motion of the Source. You can use FLUX.1 to repaint the first frame of the Source")
                # Simplified controls - Radio buttons for Yes/No and separate file upload
                with gr.Row():
                    mt_repaint_option = gr.Radio(
                        label="Repaint First Frame (Optional)",
                        choices=["No", "Yes"],
                        value="No"
                    )
                gr.Markdown("Or if you want to use your own image as repainted first frame, please upload the image in below.")
                
                mt_repaint_upload = gr.UploadButton("Upload Repaint Image (Optional)", file_types=["image"])
                mt_repaint_preview = gr.Image(label="Repaint Image Preview")
                
                mt_repaint_upload.upload(
                    fn=update_source_preview,
                    inputs=[mt_repaint_upload],
                    outputs=[mt_repaint_preview]
                )

            with gr.TabItem("Camera Control"):
                 gr.Markdown("Camera Control is not available in Huggingface Space, please deploy our [GitHub project](https://github.com/IGL-HKUST/DiffusionAsShader) on your own machine")

            with gr.TabItem("Object Manipulation"):
                 gr.Markdown("Object Manipulation is not available in Huggingface Space, please deploy our [GitHub project](https://github.com/IGL-HKUST/DiffusionAsShader) on your own machine")

    with right_column:
        
        gr.Markdown("### 4. Generate Tracking Video")
        gr.Markdown("'Generate Tracking Video' is used to preserve all motion from the Source. You need to generate tracking video before producing the final result.")
        mt_run_btn = gr.Button("Generate Tracking", variant="primary", size="lg")
        tracking_video = gr.Video(label="Tracking Video")
        
        apply_tracking_btn = gr.Button("5. Generate Video", variant="primary", size="lg", interactive=False)
        output_video = gr.Video(label="Generated Video")

        # mt_run_btn 的 click 事件定义
        mt_run_btn.click(
            fn=process_motion_transfer,
            inputs=[
                source_upload, common_prompt,
                mt_repaint_option, mt_repaint_upload
            ],
            outputs=[tracking_video, video_tensor_state, tracking_tensor_state, repaint_img_tensor_state, fps_state]
        ).then(
            fn=enable_apply_button,
            inputs=[tracking_video],
            outputs=[apply_tracking_btn]
        )

        # apply_tracking_btn 的 click 事件定义
        apply_tracking_btn.click(
            fn=apply_tracking_unified,
            inputs=[
                video_tensor_state,
                tracking_tensor_state,
                repaint_img_tensor_state,
                common_prompt,  # common_prompt 现在可用
                fps_state
            ],
            outputs=[output_video]
        )

    examples_list = load_examples()
    gr.Markdown("### Examples (For Workflow Demo Only)")
    gr.Markdown("The following examples are only for demonstrating DaS's workflow and output quality. If you want to actually generate tracking or videos, the program will not run unless you manually upload files from your devices.")
    if examples_list:
        with gr.Blocks() as examples_block:
            gr.Examples(
                examples=examples_list,
                inputs=[source_preview, mt_repaint_preview, common_prompt, tracking_video, output_video],
                outputs=[source_preview, mt_repaint_preview, common_prompt, tracking_video, output_video],
                fn=lambda *args: args,
                cache_examples=True,
                label="Examples"
            )
    

# Launch interface
if __name__ == "__main__":
    print(f"Using GPU: {GPU_ID}")
    print(f"Web UI will start on port {args.port}")
    if args.share:
        print("Creating public link for remote access")
    
    # Launch interface
    demo.launch(share=args.share, server_port=args.port)