import cv2
import numpy as np
from tqdm import tqdm
import subprocess
import os

class VideoProcessor:
    def __init__(self, config):
        self.config = config.get('video_normalization', {})
    
    def normalize_video(self, input_path, output_path):
        """Normaliza o vídeo usando OpenCV"""
        cap = cv2.VideoCapture(input_path)
        
        # Obter propriedades do vídeo original
        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))
        
        # Configurações de saída
        target_resolution = self.config.get('target_resolution', (width, height))
        target_fps = self.config.get('target_fps', fps)
        
        # Definir codec e writer
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(
            output_path, fourcc, target_fps, target_resolution
        )
        
        print(f"Processando {total_frames} frames...")
        
        for _ in tqdm(range(total_frames), desc="Normalizando vídeo"):
            ret, frame = cap.read()
            if not ret:
                break
            
            # Redimensionar
            frame = cv2.resize(frame, target_resolution)
            
            # Normalização de cor e brilho
            if self.config.get('normalize_brightness', True):
                frame = self.normalize_brightness(frame)
            
            if self.config.get('enhance_contrast', True):
                frame = self.enhance_contrast(frame)
            
            out.write(frame)
        
        cap.release()
        out.release()
        
        # Usar FFmpeg para melhor compressão (opcional)
        self.optimize_with_ffmpeg(output_path)
        
        return output_path
    
    def normalize_brightness(self, frame):
        """Normaliza o brilho do frame"""
        # Converter para YUV e normalizar canal Y (luminância)
        yuv = cv2.cvtColor(frame, cv2.COLOR_BGR2YUV)
        yuv[:,:,0] = cv2.equalizeHist(yuv[:,:,0])
        return cv2.cvtColor(yuv, cv2.COLOR_YUV2BGR)
    
    def enhance_contrast(self, frame):
        """Melhora o contraste usando CLAHE"""
        lab = cv2.cvtColor(frame, cv2.COLOR_BGR2LAB)
        l, a, b = cv2.split(lab)
        clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
        l = clahe.apply(l)
        lab = cv2.merge((l, a, b))
        return cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)
    
    def optimize_with_ffmpeg(self, input_path):
        """Otimiza o vídeo com FFmpeg"""
        temp_path = input_path.replace('.mp4', '_temp.mp4')
        
        cmd = [
            'ffmpeg', '-i', input_path,
            '-c:v', 'libx264', '-preset', 'medium', '-crf', '23',
            '-c:a', 'aac', '-b:a', '128k',
            '-movflags', '+faststart',
            '-y', temp_path
        ]
        
        try:
            subprocess.run(cmd, check=True, capture_output=True)
            os.replace(temp_path, input_path)
        except (subprocess.CalledProcessError, FileNotFoundError):
            # Fallback se FFmpeg não estiver disponível
            if os.path.exists(temp_path):
                os.remove(temp_path)