import os
import torch
import numpy as np
import folder_paths
import av
from server import PromptServer
from aiohttp import web
import comfy.utils

# Custom API route to serve video files from anywhere on the user's system for the frontend preview
@PromptServer.instance.routes.get("/video_ui_custom_view")
async def custom_view(request):
    file_path = request.query.get("filename", "")
    if os.path.exists(file_path) and os.path.isfile(file_path):
        return web.FileResponse(file_path)
    return web.Response(status=404, text="File not found")

# Custom API route for Chunked Uploads to bypass the 413 Payload Too Large error
@PromptServer.instance.routes.post("/video_ui_upload_chunk")
async def upload_chunk(request):
    post = await request.post()
    file = post.get("file")
    filename = post.get("filename")
    chunk_index = int(post.get("chunk_index"))
    total_chunks = int(post.get("total_chunks"))

    upload_dir = folder_paths.get_input_directory()
    file_path = os.path.join(upload_dir, filename)

    # Append to file if it's not the first chunk, otherwise write new
    mode = "ab" if chunk_index > 0 else "wb"
    with open(file_path, mode) as f:
        f.write(file.file.read())

    if chunk_index == total_chunks - 1:
        return web.json_response({"name": filename})
    return web.json_response({"status": "ok"})


class LoadVideoUI:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "video": ("STRING", {"default": ""}),
                "start_time": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100000.0, "step": 0.01}),
                "end_time": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100000.0, "step": 0.01}),
                "duration": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100000.0, "step": 0.01}),
                "start_frame": ("INT", {"default": 0, "min": 0, "max": 10000000, "step": 1}),
                "end_frame": ("INT", {"default": 0, "min": 0, "max": 10000000, "step": 1}),
                "duration_frames": ("INT", {"default": 0, "min": 0, "max": 10000000, "step": 1}),  
                "resize_method": (["maintain aspect ratio", "stretch to fit", "pad", "crop"], {"default": "maintain aspect ratio"}),                
                "custom_width": ("INT", {"default": 0, "min": 0, "max": 100000, "step": 8, "tooltip": "Custom width. 0 means original width."}),
                "custom_height": ("INT", {"default": 0, "min": 0, "max": 100000, "step": 8, "tooltip": "Custom height. 0 means original height."}),
                "frame_rate": ("INT", {"default": 24, "min": 1, "max": 120, "step": 1, "tooltip": "Force the video to a specific frame rate for extraction."}),
                "display_mode": (["seconds", "frames"], {"default": "seconds"}),
                "crop_x": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
                "crop_y": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
                "crop_w": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}),
                "crop_h": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}),
            }
        }

    RETURN_TYPES = ("IMAGE", "AUDIO", "FLOAT", "INT")
    RETURN_NAMES = ("images", "audio", "duration", "frame_count")
    FUNCTION = "load_video"
    CATEGORY = "WhatDreamsCost"

    def load_video(self, video, frame_rate, display_mode, start_time, end_time, duration, start_frame, end_frame, duration_frames, custom_width=0, custom_height=0, resize_method="maintain aspect ratio", crop_x=0.0, crop_y=0.0, crop_w=1.0, crop_h=1.0, **kwargs):
        if not video:
            # Return blank defaults if no video is loaded
            empty_image = torch.zeros((1, 512, 512, 3), dtype=torch.float32)
            empty_audio = {"waveform": torch.zeros((1, 1, 44100)), "sample_rate": 44100}
            return (empty_image, empty_audio, 0.0, 0)

        # 1. Resolve path using ComfyUI standard paths or Absolute Path
        video_path = video  # Try exact/absolute path first
        if not os.path.exists(video_path):
            video_path_annotated = folder_paths.get_annotated_filepath(video)
            if os.path.exists(video_path_annotated):
                video_path = video_path_annotated
            else:
                video_path_input = os.path.join(folder_paths.get_input_directory(), video)
                if os.path.exists(video_path_input):
                    video_path = video_path_input
                else:
                    raise FileNotFoundError(f"Video file not found: {video}")

        # Open container to read streams and metadata
        container = av.open(video_path)
        
        # Determine video stream and duration
        video_stream = container.streams.video[0] if len(container.streams.video) > 0 else None
        video_duration = 0
        if video_stream and video_stream.duration and video_stream.time_base:
            video_duration = float(video_stream.duration * video_stream.time_base)

        orig_w = video_stream.codec_context.width if video_stream else 512
        orig_h = video_stream.codec_context.height if video_stream else 512

        # Determine correct colorspace and color range for PyAV conversion to prevent color shift
        try:
            from av.video.reformatter import Colorspace, ColorRange
            # Improve fallback heuristic to check both dimensions (e.g. 720x1280 vertical video is HD)
            fallback_cs = Colorspace.ITU709 if max(orig_w, orig_h) >= 720 else Colorspace.ITU601
            fallback_cr = ColorRange.MPEG
            dst_range = ColorRange.JPEG # RGB should always be full range
        except ImportError:
            fallback_cs = "itu709" if max(orig_w, orig_h) >= 720 else "itu601"
            fallback_cr = "mpeg"
            dst_range = "jpeg"
            
        src_colorspace = fallback_cs
        src_color_range = fallback_cr
        
        if video_stream and video_stream.codec_context:
            cc = video_stream.codec_context
            
            c_space = getattr(cc, 'colorspace', getattr(cc, 'color_space', None))
            if c_space and hasattr(c_space, 'name') and c_space.name != "UNSPECIFIED":
                src_colorspace = c_space
            elif c_space and isinstance(c_space, str) and "unspecified" not in c_space.lower():
                src_colorspace = c_space
                
            c_range = getattr(cc, 'color_range', None)
            if c_range and hasattr(c_range, 'name') and c_range.name != "UNSPECIFIED":
                src_color_range = c_range
            elif c_range and isinstance(c_range, str) and "unspecified" not in c_range.lower():
                src_color_range = c_range

        target_w = custom_width if custom_width > 0 else orig_w
        target_h = custom_height if custom_height > 0 else orig_h
        
        target_w = target_w - (target_w % 2)
        target_h = target_h - (target_h % 2)
        
        # Calculate manual crop from interactive UI first
        manual_crop_left = int(orig_w * crop_x)
        manual_crop_top = int(orig_h * crop_y)
        manual_crop_right = orig_w - int(orig_w * (crop_x + crop_w))
        manual_crop_bottom = orig_h - int(orig_h * (crop_y + crop_h))
        
        # Ensure we don't crop more than the image
        manual_crop_left = max(0, min(manual_crop_left, orig_w - 1))
        manual_crop_top = max(0, min(manual_crop_top, orig_h - 1))
        manual_crop_right = max(0, min(manual_crop_right, orig_w - manual_crop_left - 1))
        manual_crop_bottom = max(0, min(manual_crop_bottom, orig_h - manual_crop_top - 1))
        
        # After manual crop, the new original dimensions are:
        cropped_orig_w = orig_w - manual_crop_left - manual_crop_right
        cropped_orig_h = orig_h - manual_crop_top - manual_crop_bottom
        
        # If no custom width/height is provided, use the cropped original dimensions
        if custom_width == 0:
            target_w = cropped_orig_w
            target_w = target_w - (target_w % 2)
        if custom_height == 0:
            target_h = cropped_orig_h
            target_h = target_h - (target_h % 2)

        scale_w, scale_h = target_w, target_h
        pad_left = pad_right = pad_top = pad_bottom = 0
        crop_left = crop_right = crop_top = crop_bottom = 0

        if custom_width > 0 or custom_height > 0:
            if resize_method == "maintain aspect ratio" or resize_method == "pad":
                ratio = min(target_w / cropped_orig_w, target_h / cropped_orig_h)
                scale_w = int(cropped_orig_w * ratio)
                scale_h = int(cropped_orig_h * ratio)
                scale_w = scale_w - (scale_w % 2)
                scale_h = scale_h - (scale_h % 2)
                
                if resize_method == "pad":
                    pad_x = target_w - scale_w
                    pad_y = target_h - scale_h
                    pad_left = pad_x // 2
                    pad_right = pad_x - pad_left
                    pad_top = pad_y // 2
                    pad_bottom = pad_y - pad_top
                else:
                    target_w, target_h = scale_w, scale_h

            elif resize_method == "crop":
                ratio = max(target_w / cropped_orig_w, target_h / cropped_orig_h)
                scale_w = int(cropped_orig_w * ratio)
                scale_h = int(cropped_orig_h * ratio)
                scale_w = scale_w - (scale_w % 2)
                scale_h = scale_h - (scale_h % 2)
                
                crop_x = scale_w - target_w
                crop_y = scale_h - target_h
                crop_left = crop_x // 2
                crop_right = crop_x - crop_left
                crop_top = crop_y // 2
                crop_bottom = crop_y - crop_top

            elif resize_method == "stretch to fit":
                scale_w, scale_h = target_w, target_h

        # Determine exact bounds based on frontend mode
        if display_mode == "frames":
            fr = float(frame_rate) if frame_rate > 0 else 24.0
            actual_start_time = float(start_frame) / fr
            actual_end_time = float(end_frame) / fr if (end_frame > 0 and end_frame > start_frame) else video_duration
        else:
            actual_start_time = start_time
            actual_end_time = end_time if (end_time > 0 and end_time > start_time) else video_duration

        if actual_end_time <= 0:
            actual_end_time = float('inf') # Fallback if duration is unknown

        # 2. Extract Video Frames (PyAV)
        frames = []
        image_tensor = None
        frames_loaded = 0
        
        if video_stream:
            video_stream.thread_type = "AUTO" # Enable multithreaded decoding
            
            # Efficiently seek backwards to the nearest keyframe
            if video_stream.time_base:
                seek_pts = int(actual_start_time / float(video_stream.time_base))
            else:
                seek_pts = int(actual_start_time * av.time_base)
            
            container.seek(seek_pts, stream=video_stream, backward=True)
            
            # Custom sampling to force specific framerate 
            frame_interval = 1.0 / float(frame_rate) if frame_rate > 0 else 1.0/24.0
            expected_target_time = actual_start_time
            
            # Pre-calculate expected frames
            alloc_end_time = actual_end_time if actual_end_time != float('inf') else video_duration
            expected_frames = 0
            if alloc_end_time > 0:
                duration_to_extract = alloc_end_time - actual_start_time
                if duration_to_extract > 0:
                    expected_frames = int(np.ceil(duration_to_extract / frame_interval)) + 2
                    
            pbar = comfy.utils.ProgressBar(expected_frames) if expected_frames > 0 else None

            for frame in container.decode(video_stream):
                frame_time = frame.time
                if frame_time is None:
                    frame_time = float(frame.pts * float(video_stream.time_base)) if frame.pts and video_stream.time_base else 0.0

                if frame_time < actual_start_time:
                    continue
                    
                # Add a slight buffer (interval) to ensure we evaluate the boundary correctly
                if frame_time > actual_end_time + frame_interval: 
                    break
                    
                # Fix PyAV color shift by forcing proper colorspace and range conversion.
                # Omit dst_colorspace so swscale defaults naturally for RGB output
                # (passing it can cause the YUV matrix to be applied incorrectly).
                try:
                    frame = frame.reformat(
                        format="rgb24",
                        src_colorspace=src_colorspace,
                        src_color_range=src_color_range,
                        dst_color_range=dst_range
                    )
                    frame_rgb = frame.to_ndarray(format='rgb24')
                except Exception as e:
                    # Fallback: if explicit color reformat fails, use PyAV's default conversion
                    print(f"[LoadVideoUI] Color reformat failed, using default: {e}")
                    frame_rgb = frame.to_ndarray(format='rgb24')
                
                # Apply interactive crop first
                if manual_crop_left > 0 or manual_crop_top > 0 or manual_crop_right > 0 or manual_crop_bottom > 0:
                    frame_rgb = frame_rgb[manual_crop_top:orig_h-manual_crop_bottom, manual_crop_left:orig_w-manual_crop_right, :]
                    
                # Now resize to the scaled dimensions
                if scale_w != cropped_orig_w or scale_h != cropped_orig_h:
                    import cv2
                    frame_rgb = cv2.resize(frame_rgb, (scale_w, scale_h), interpolation=cv2.INTER_AREA)
                
                if crop_left > 0 or crop_top > 0 or crop_right > 0 or crop_bottom > 0:
                    frame_rgb = frame_rgb[crop_top:scale_h-crop_bottom, crop_left:scale_w-crop_right, :]
                if pad_left > 0 or pad_top > 0 or pad_right > 0 or pad_bottom > 0:
                    frame_rgb = np.pad(frame_rgb, ((pad_top, pad_bottom), (pad_left, pad_right), (0, 0)), mode='constant', constant_values=0)
                
                # Duplicate or skip frames perfectly based on timestamps to meet forced framerate.
                # FIX: Use strictly less than (<) for actual_end_time to prevent the loop from fetching an extra +1 frame
                # at the exact boundary of the duration slice!
                while expected_target_time <= frame_time and expected_target_time < actual_end_time - 1e-5:
                    if image_tensor is None and expected_frames > 0:
                        # First frame: allocate the tensor
                        height, width = frame_rgb.shape[:2]
                        alloc_frames = expected_frames + 50 # Add generous buffer to prevent reallocation
                        try:
                            image_tensor = torch.zeros((alloc_frames, height, width, 3), dtype=torch.float32)
                        except Exception as e:
                            print(f"[LoadVideoUI] Pre-allocation failed, falling back to list: {e}")
                            expected_frames = 0 # Disable pre-allocation
                            
                    if image_tensor is not None:
                        # Check bounds (just in case)
                        if frames_loaded >= image_tensor.shape[0]:
                            # Extend tensor if we underestimated
                            extension = torch.zeros((50, image_tensor.shape[1], image_tensor.shape[2], 3), dtype=torch.float32)
                            image_tensor = torch.cat((image_tensor, extension), dim=0)
                            
                        # Insert frame with minimal memory copy directly to tensor
                        image_tensor[frames_loaded] = torch.from_numpy(frame_rgb).float().div_(255.0)
                        frames_loaded += 1
                    else:
                        # Fallback list append if pre-allocation failed
                        frames.append(frame_rgb)
                        
                    if pbar:
                        pbar.update(1)
                        
                    expected_target_time += frame_interval

        # Convert frames to ComfyUI Image standard format [N, H, W, C], float32, range 0.0-1.0
        if image_tensor is not None:
            if frames_loaded > 0:
                image_tensor = image_tensor[:frames_loaded]
            else:
                image_tensor = torch.zeros((1, 512, 512, 3), dtype=torch.float32)
        elif len(frames) > 0:
            frames_np = np.array(frames, dtype=np.float32) / 255.0
            image_tensor = torch.from_numpy(frames_np)
        else:
            # Fallback for an empty slice
            image_tensor = torch.zeros((1, 512, 512, 3), dtype=torch.float32)

        # 3. Extract Audio (PyAV)
        audio_dict = {"waveform": torch.zeros((1, 1, 44100)), "sample_rate": 44100} # Default empty audio
        
        if len(container.streams.audio) > 0:
            try:
                audio_stream = container.streams.audio[0]
                audio_stream.thread_type = "AUTO"
                sample_rate = getattr(audio_stream, 'rate', 44100) or 44100
                
                # We must seek again on the container specifically for the audio stream
                if audio_stream.time_base:
                    seek_pts = int(actual_start_time / float(audio_stream.time_base))
                else:
                    seek_pts = int(actual_start_time * av.time_base)
                    
                container.seek(seek_pts, stream=audio_stream, backward=True)
                
                # Resample to standard float planar format (fltp)
                resampler = av.AudioResampler(format='fltp')
                
                audio_data = []
                first_frame_time = None
                
                for frame in container.decode(audio_stream):
                    frame_time = frame.time
                    if frame_time is None:
                        frame_time = float(frame.pts * float(audio_stream.time_base)) if frame.pts and audio_stream.time_base else 0.0
                        
                    # Give a small 1-second buffer to ensure we catch end frames
                    if frame_time > actual_end_time + 1.0: 
                        break
                        
                    if first_frame_time is None:
                        first_frame_time = frame_time
                        
                    resampled_frames = resampler.resample(frame)
                    for r_frame in resampled_frames:
                        audio_data.append(r_frame.to_ndarray())
                        
                if audio_data:
                    # Concatenate all frames horizontally along the sample axis
                    waveform_np = np.concatenate(audio_data, axis=1)
                    waveform = torch.from_numpy(waveform_np).float()
                    
                    if first_frame_time is None:
                        first_frame_time = 0.0
                        
                    # Calculate exact slice points to trim precisely
                    offset_sec = max(0.0, actual_start_time - first_frame_time)
                    start_sample = int(offset_sec * sample_rate)
                    
                    duration_sec_audio = actual_end_time - actual_start_time
                    end_sample = start_sample + int(duration_sec_audio * sample_rate)
                    
                    # Trim array bounds properly
                    if end_sample > start_sample:
                        waveform = waveform[:, start_sample:end_sample]
                    else:
                        waveform = waveform[:, start_sample:]
                        
                    # Expand to ComfyUI Audio standard [batch_size, channels, samples]
                    waveform = waveform.unsqueeze(0)
                    audio_dict = {"waveform": waveform, "sample_rate": sample_rate}
                    
            except Exception as e:
                # Catch gracefully without breaking the pipeline execution
                print(f"[LoadVideoUI] Audio track extraction skipped or failed: {e}")

        # Always close container to free up system memory lock
        container.close()
        
        # Output accurate final duration in seconds
        final_duration_sec = float(max(0.0, actual_end_time - actual_start_time))
        
        # Accurately output the true number of extracted frames 
        # (Using the shape of the array provides exact 1:1 parity with the timeline's math)
        frame_count = image_tensor.shape[0] if (frames_loaded > 0 or len(frames) > 0) else 0
        if frame_count == 0 and final_duration_sec > 0:
             # Fallback estimation only if PyAV completely failed to decode a valid chunk
             calc_fr = float(frame_rate) if frame_rate > 0 else 24.0
             frame_count = int(np.floor(final_duration_sec * calc_fr))

        return (image_tensor, audio_dict, final_duration_sec, frame_count)