import cv2
import torch
import numpy as np
import concurrent.futures
import os

# --- The Parallel Video Loader Node ---
class ParallelSharpnessLoader:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "video_path": ("STRING", {"default": "C:\\path\\to\\video.mp4"}),
                # scan_limit: 0 means ALL frames. Max set high to allow large user inputs.
                "scan_limit": ("INT", {"default": 1440, "min": 0, "max": 10000000, "step": 1, "label": "Max Frames to Scan (0=All)"}),
                "frame_scan_step": ("INT", {"default": 5, "min": 1, "step": 1, "label": "Analyze Every Nth Frame"}),
                "return_count": ("INT", {"default": 4, "min": 1, "max": 1024, "step": 1, "label": "Best Frames Count"}),
                "min_distance": ("INT", {"default": 24, "min": 0, "max": 10000, "step": 1, "label": "Min Distance (Frames)"}),
                # FIXED: Max limit increased to 10 million to prevent slider locking at 2048
                "skip_start": ("INT", {"default": 0, "min": 0, "max": 10000000, "step": 1}),
            },
        }

    RETURN_TYPES = ("IMAGE", "STRING")
    RETURN_NAMES = ("images", "scores_info")
    FUNCTION = "load_video"
    CATEGORY = "BetaHelper/Video"

    def calculate_sharpness(self, frame_data):
        gray = cv2.cvtColor(frame_data, cv2.COLOR_BGR2GRAY)
        return cv2.Laplacian(gray, cv2.CV_64F).var()

    def load_video(self, video_path, scan_limit, frame_scan_step, return_count, min_distance, skip_start):
        # 1. Validation
        if not os.path.exists(video_path):
            video_path = video_path.strip('"')
        if not os.path.exists(video_path):
            raise FileNotFoundError(f"Video not found: {video_path}")

        cap = cv2.VideoCapture(video_path)
        
        # 2. Scanning (Pass 1 - Fast)
        print(f"xx- Parallel Loader: Scanning {video_path}...")
        
        frame_scores = [] 
        current_frame = skip_start
        scanned_count = 0
        
        # Seek to start
        if skip_start > 0:
            cap.set(cv2.CAP_PROP_POS_FRAMES, skip_start)

        # Thread Pool
        with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor:
            futures = []
            
            while True:
                if scan_limit > 0 and scanned_count >= scan_limit:
                    break
                
                ret, frame = cap.read()
                if not ret: break 

                future = executor.submit(self.calculate_sharpness, frame)
                futures.append((current_frame, future))
                scanned_count += 1
                
                # Manual Stepping
                if frame_scan_step > 1:
                    for _ in range(frame_scan_step - 1):
                        if not cap.grab(): break
                        current_frame += 1
                
                current_frame += 1

            # Gather results
            for idx, future in futures:
                frame_scores.append((idx, future.result()))

        cap.release()

        # 3. Selection
        frame_scores.sort(key=lambda x: x[1], reverse=True)
        selected = []
        
        for idx, score in frame_scores:
            if len(selected) >= return_count: break
            
            if all(abs(s[0] - idx) >= min_distance for s in selected):
                selected.append((idx, score))

        # Sort selected by timeline
        selected.sort(key=lambda x: x[0])
        print(f"xx- Selected Frames: {[f[0] for f in selected]}")

        # 4. Extraction (Pass 2)
        cap = cv2.VideoCapture(video_path)
        output_tensors = []
        info_log = []

        for idx, score in selected:
            cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
            ret, frame = cap.read()
            if ret:
                frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                frame = frame.astype(np.float32) / 255.0
                output_tensors.append(torch.from_numpy(frame))
                info_log.append(f"F:{idx} (Score:{int(score)})")
        
        cap.release()

        if not output_tensors:
            return (torch.zeros((1,64,64,3)), "Failed")

        return (torch.stack(output_tensors), ", ".join(info_log))

NODE_CLASS_MAPPINGS = {
    "ParallelSharpnessLoader": ParallelSharpnessLoader
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "ParallelSharpnessLoader": "Parallel Video Loader (Sharpness)"
}