import os
import sys
import argparse
from PIL import Image
project_root = os.path.dirname(os.path.abspath(__file__))
try:
    sys.path.append(os.path.join(project_root, "submodules/MoGe"))
    os.environ["TOKENIZERS_PARALLELISM"] = "false"
except:
    print("Warning: MoGe not found, motion transfer will not be applied")
    
import torch
import numpy as np
from PIL import Image
import torchvision.transforms as transforms
from moviepy.editor import VideoFileClip
from diffusers.utils import load_image, load_video

from models.pipelines import DiffusionAsShaderPipeline, FirstFrameRepainter, CameraMotionGenerator, ObjectMotionGenerator
from submodules.MoGe.moge.model import MoGeModel

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]: Video tensor [T,C,H,W] and FPS
    """
    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:
        frames = load_video(media_path)
        fps = len(frames) / VideoFileClip(media_path).duration
    else:
        # Handle image as single frame
        image = load_image(media_path)
        frames = [image]
        fps = 8  # Default fps for images
    
    # Ensure we have exactly max_frames
    if len(frames) > max_frames:
        frames = frames[:max_frames]
    elif len(frames) < max_frames:
        last_frame = frames[-1]
        while len(frames) < max_frames:
            frames.append(last_frame.copy())
            
    # Convert frames to tensor
    video_tensor = torch.stack([transform(frame) for frame in frames])
    
    return video_tensor, fps, is_video

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_path', type=str, default=None, help='Path to input video/image')
    parser.add_argument('--prompt', type=str, required=True, help='Repaint prompt')
    parser.add_argument('--output_dir', type=str, default='outputs', help='Output directory')
    parser.add_argument('--gpu', type=int, default=0, help='GPU device ID')
    parser.add_argument('--checkpoint_path', type=str, default="EXCAI/Diffusion-As-Shader", help='Path to model checkpoint')
    parser.add_argument('--depth_path', type=str, default=None, help='Path to depth image')
    parser.add_argument('--tracking_path', type=str, default=None, help='Path to tracking video, if provided, camera motion and object manipulation will not be applied')
    parser.add_argument('--repaint', type=str, default=None, 
                       help='Path to repainted image, or "true" to perform repainting, if not provided use original frame')
    parser.add_argument('--camera_motion', type=str, default=None, 
                    help='Camera motion mode: "trans <dx> <dy> <dz>" or "rot <axis> <angle>" or "spiral <radius>"')
    parser.add_argument('--object_motion', type=str, default=None, help='Object motion mode: up/down/left/right')
    parser.add_argument('--object_mask', type=str, default=None, help='Path to object mask image (binary image)')
    parser.add_argument('--tracking_method', type=str, default='spatracker', choices=['spatracker', 'moge'], 
                    help='Tracking method to use (spatracker or moge)')
    args = parser.parse_args()
    
    # Load input video/image
    video_tensor, fps, is_video = load_media(args.input_path)
    if not is_video:
        args.tracking_method = "moge"
        print("Image input detected, using MoGe for tracking video generation.")

    # Initialize pipeline
    das = DiffusionAsShaderPipeline(gpu_id=args.gpu, output_dir=args.output_dir)
    if args.tracking_method == "moge" and args.tracking_path is None:
        moge = MoGeModel.from_pretrained("Ruicheng/moge-vitl").to(das.device)
    
    # Repaint first frame if requested
    repaint_img_tensor = None
    if args.repaint:
        if args.repaint.lower() == "true":
            repainter = FirstFrameRepainter(gpu_id=args.gpu, output_dir=args.output_dir)
            repaint_img_tensor = repainter.repaint(
                video_tensor[0], 
                prompt=args.prompt,
                depth_path=args.depth_path
            )
        else:
            repaint_img_tensor, _, _ = load_media(args.repaint)
            repaint_img_tensor = repaint_img_tensor[0]  # Take first frame

    # Generate tracking if not provided
    tracking_tensor = None
    pred_tracks = None
    cam_motion = CameraMotionGenerator(args.camera_motion)

    if args.tracking_path:
        tracking_tensor, _, _ = load_media(args.tracking_path)
        
    elif args.tracking_method == "moge":
        # Use the first frame from previously loaded video_tensor
        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]
        cam_motion.set_intr(infer_result["intrinsics"])

        # Apply object motion if specified
        if args.object_motion:
            if args.object_mask is None:
                raise ValueError("Object motion specified but no mask provided. Please provide a mask image with --object_mask")
                
            # Load mask image
            mask_image = Image.open(args.object_mask).convert('L')  # Convert to grayscale
            mask_image = transforms.Resize((480, 720))(mask_image)  # Resize to match video size
            # Convert to binary mask
            mask = torch.from_numpy(np.array(mask_image) > 127)  # Threshold at 127
            
            motion_generator = ObjectMotionGenerator(device=das.device)

            pred_tracks = motion_generator.apply_motion(
                pred_tracks=pred_tracks,
                mask=mask,
                motion_type=args.object_motion,
                distance=50,
                num_frames=49,
                tracking_method="moge"
            )
            print("Object motion applied")

        # Apply camera motion if specified
        if args.camera_motion:
            poses = cam_motion.get_default_motion() # shape: [49, 4, 4]
            print("Camera motion applied")
        else:
            # no poses
            poses = torch.eye(4).unsqueeze(0).repeat(49, 1, 1)
        # change pred_tracks into screen coordinate
        pred_tracks_flatten = pred_tracks.reshape(video_tensor.shape[0], H*W, 3)
        pred_tracks = cam_motion.w2s(pred_tracks_flatten, poses).reshape([video_tensor.shape[0], H, W, 3]) # [T, H, W, 3]
        _, tracking_tensor = das.visualize_tracking_moge(
            pred_tracks.cpu().numpy(), 
            infer_result["mask"].cpu().numpy()
        )
        print('export tracking video via MoGe.')

    else:
        # Generate tracking points
        pred_tracks, pred_visibility, T_Firsts = das.generate_tracking_spatracker(video_tensor)

        # Apply camera motion if specified
        if args.camera_motion:
            poses = cam_motion.get_default_motion() # shape: [49, 4, 4]
            pred_tracks = cam_motion.apply_motion_on_pts(pred_tracks, poses)
            print("Camera motion applied")
        
        # Apply object motion if specified
        if args.object_motion:
            if args.object_mask is None:
                raise ValueError("Object motion specified but no mask provided. Please provide a mask image with --object_mask")
                
            # Load mask image
            mask_image = Image.open(args.object_mask).convert('L')  # Convert to grayscale
            mask_image = transforms.Resize((480, 720))(mask_image)  # Resize to match video size
            # Convert to binary mask
            mask = torch.from_numpy(np.array(mask_image) > 127)  # Threshold at 127
            
            motion_generator = ObjectMotionGenerator(device=das.device)
            
            pred_tracks = motion_generator.apply_motion(
                pred_tracks=pred_tracks.squeeze(),
                mask=mask,
                motion_type=args.object_motion,
                distance=50,
                num_frames=49,
                tracking_method="spatracker"
            ).unsqueeze(0)
            print(f"Object motion '{args.object_motion}' applied using mask from {args.object_mask}")
    
        # Generate tracking tensor from modified tracks
        _, tracking_tensor = das.visualize_tracking_spatracker(video_tensor, pred_tracks, pred_visibility, T_Firsts)
    
    das.apply_tracking(
        video_tensor=video_tensor,
        fps=8,
        tracking_tensor=tracking_tensor,
        img_cond_tensor=repaint_img_tensor,
        prompt=args.prompt,
        checkpoint_path=args.checkpoint_path
    )