import cv2
import numpy as np
from PIL import Image, ImageEnhance
import os
from pathlib import Path
from typing import Tuple, Optional, List
import argparse

class SubEffectEnhancer:
    """Subtitle Chase Effect Enhancer"""
    
    def __init__(self):
        self.effect_presets = {
            'vibrant': {'brightness': 1.2, 'contrast': 1.3, 'saturation': 1.4},
            'soft': {'brightness': 1.1, 'contrast': 1.1, 'saturation': 1.2},
            'dramatic': {'brightness': 1.3, 'contrast': 1.5, 'saturation': 1.6},
            'cinematic': {'brightness': 0.9, 'contrast': 1.4, 'saturation': 1.3}
        }
    
    def enhance_image(self, image_path: str, output_path: str, 
                     brightness: float = 1.2, 
                     contrast: float = 1.3, 
                     saturation: float = 1.4,
                     sharpness: float = 1.2,
                     glow_intensity: float = 0.1) -> bool:
        """
        Enhance subtitle chase effect for single image
        
        Args:
            image_path: Input image path
            output_path: Output image path
            brightness: Brightness enhancement factor
            contrast: Contrast enhancement factor
            saturation: Saturation enhancement factor
            sharpness: Sharpness enhancement factor
            glow_intensity: Glow effect intensity
            
        Returns:
            bool: Whether processing was successful
        """
        try:
            # Open image using PIL
            with Image.open(image_path) as img:
                # Convert to RGB mode (ensure processing color images)
                if img.mode != 'RGB':
                    img = img.convert('RGB')
                
                # Apply various enhancement effects
                enhanced_img = self._apply_enhancements(img, brightness, contrast, 
                                                      saturation, sharpness)
                
                # Add glow effect
                if glow_intensity > 0:
                    enhanced_img = self._add_glow_effect(enhanced_img, glow_intensity)
                
                # Save result
                enhanced_img.save(output_path, quality=95)
                print(f"Image enhancement completed: {output_path}")
                return True
                
        except Exception as e:
            print(f"Error processing image: {e}")
            return False
    
    def enhance_video(self, video_path: str, output_path: str,
                     brightness: float = 1.2,
                     contrast: float = 1.3,
                     saturation: float = 1.4,
                     frame_skip: int = 0) -> bool:
        """
        Enhance subtitle chase effect for video
        
        Args:
            video_path: Input video path
            output_path: Output video path
            brightness: Brightness enhancement factor
            contrast: Contrast enhancement factor
            saturation: Saturation enhancement factor
            frame_skip: Frame skip processing (0 means process all frames)
            
        Returns:
            bool: Whether processing was successful
        """
        try:
            cap = cv2.VideoCapture(video_path)
            
            # Get video information
            fps = cap.get(cv2.CAP_PROP_FPS)
            width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
            height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            
            # Create video writer
            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
            out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
            
            frame_count = 0
            processed_frames = 0
            
            print(f"Starting video processing: {video_path}")
            print(f"Total frames: {total_frames}, FPS: {fps}")
            
            while True:
                ret, frame = cap.read()
                if not ret:
                    break
                
                # Frame skip processing
                if frame_skip > 0 and frame_count % (frame_skip + 1) != 0:
                    frame_count += 1
                    continue
                
                # Convert BGR to RGB
                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                
                # Convert to PIL image for processing
                pil_img = Image.fromarray(frame_rgb)
                enhanced_pil = self._apply_enhancements(pil_img, brightness, contrast, saturation)
                
                # Convert back to OpenCV format
                enhanced_frame = cv2.cvtColor(np.array(enhanced_pil), cv2.COLOR_RGB2BGR)
                
                # Write to output video
                out.write(enhanced_frame)
                
                processed_frames += 1
                frame_count += 1
                
                # Display progress
                if frame_count % 30 == 0:
                    progress = (frame_count / total_frames) * 100
                    print(f"Processing progress: {progress:.1f}%")
            
            # Release resources
            cap.release()
            out.release()
            
            print(f"Video processing completed: {output_path}")
            print(f"Processed frames: {processed_frames}")
            return True
            
        except Exception as e:
            print(f"Error processing video: {e}")
            return False
    
    def _apply_enhancements(self, img: Image.Image, brightness: float, 
                           contrast: float, saturation: float, 
                           sharpness: float = 1.0) -> Image.Image:
        """Apply various image enhancement effects"""
        # Brightness enhancement
        enhancer = ImageEnhance.Brightness(img)
        img = enhancer.enhance(brightness)
        
        # Contrast enhancement
        enhancer = ImageEnhance.Contrast(img)
        img = enhancer.enhance(contrast)
        
        # Saturation enhancement
        enhancer = ImageEnhance.Color(img)
        img = enhancer.enhance(saturation)
        
        # Sharpness enhancement
        if sharpness != 1.0:
            enhancer = ImageEnhance.Sharpness(img)
            img = enhancer.enhance(sharpness)
        
        return img
    
    def _add_glow_effect(self, img: Image.Image, intensity: float) -> Image.Image:
        """Add glow effect"""
        # Convert image to array
        img_array = np.array(img)
        
        # Create Gaussian blur version
        blurred = cv2.GaussianBlur(img_array, (0, 0), 10)
        
        # Blend original and blurred images
        result = cv2.addWeighted(img_array, 1.0, blurred, intensity, 0)
        
        return Image.fromarray(result)
    
    def batch_process(self, input_dir: str, output_dir: str, 
                     preset: str = 'vibrant') -> None:
        """
        Batch process all images and videos in directory
        
        Args:
            input_dir: Input directory
            output_dir: Output directory
            preset: Effect preset
        """
        input_path = Path(input_dir)
        output_path = Path(output_dir)
        output_path.mkdir(parents=True, exist_ok=True)
        
        # Get preset parameters
        if preset in self.effect_presets:
            params = self.effect_presets[preset]
        else:
            params = self.effect_presets['vibrant']
        
        # Process image files
        image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff'}
        for img_file in input_path.iterdir():
            if img_file.suffix.lower() in image_extensions:
                output_file = output_path / f"enhanced_{img_file.name}"
                self.enhance_image(str(img_file), str(output_file), **params)
        
        # Process video files
        video_extensions = {'.mp4', '.avi', '.mov', '.mkv'}
        for vid_file in input_path.iterdir():
            if vid_file.suffix.lower() in video_extensions:
                output_file = output_path / f"enhanced_{vid_file.stem}.mp4"
                self.enhance_video(str(vid_file), str(output_file), **params)

def main():
    """Main function"""
    parser = argparse.ArgumentParser(description='Subtitle Chase Effect Enhancement Tool')
    parser.add_argument('--input', '-i', required=True, help='Input file or directory path')
    parser.add_argument('--output', '-o', required=True, help='Output file or directory path')
    parser.add_argument('--preset', '-p', choices=['vibrant', 'soft', 'dramatic', 'cinematic'], 
                       default='vibrant', help='Effect preset')
    parser.add_argument('--brightness', '-b', type=float, default=1.2, help='Brightness enhancement factor')
    parser.add_argument('--contrast', '-c', type=float, default=1.3, help='Contrast enhancement factor')
    parser.add_argument('--saturation', '-s', type=float, default=1.4, help='Saturation enhancement factor')
    parser.add_argument('--glow', '-g', type=float, default=0.1, help='Glow effect intensity')
    
    args = parser.parse_args()
    
    enhancer = SubEffectEnhancer()
    
    # Check if input path exists
    input_path = Path(args.input)
    if not input_path.exists():
        print(f"Error: Input path does not exist: {args.input}")
        return
    
    # Process single file
    if input_path.is_file():
        output_path = Path(args.output)
        
        if input_path.suffix.lower() in {'.jpg', '.jpeg', '.png', '.bmp', '.tiff'}:
            # Process image
            enhancer.enhance_image(
                str(input_path), str(output_path),
                brightness=args.brightness,
                contrast=args.contrast,
                saturation=args.saturation,
                glow_intensity=args.glow
            )
        elif input_path.suffix.lower() in {'.mp4', '.avi', '.mov', '.mkv'}:
            # Process video
            enhancer.enhance_video(
                str(input_path), str(output_path),
                brightness=args.brightness,
                contrast=args.contrast,
                saturation=args.saturation
            )
        else:
            print(f"Unsupported file format: {input_path.suffix}")
    
    # Process directory
    elif input_path.is_dir():
        enhancer.batch_process(args.input, args.output, args.preset)

if __name__ == "__main__":
    mai