import cv2
import numpy as np
from PIL import Image
import os
import subprocess
import tempfile

class VideoSampler:
    """Extracts frames and audio for forensic analysis with lazy-loading support."""
    
    def get_info(self, video_path: str):
        """Quickly probes video for meta information."""
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened(): return {}
        
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        fps = cap.get(cv2.CAP_PROP_FPS)
        duration = total_frames / (fps if fps > 0 else 1)
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        cap.release()
        
        return {
            "duration": round(duration, 2),
            "fps": round(fps, 1),
            "total_frames": total_frames,
            "dimensions": f"{width}x{height}"
        }

    def extract_audio(self, video_path: str) -> str:
        """Isolated audio extraction for lazy Phase 3 processing."""
        try:
            temp_dir = tempfile.gettempdir()
            audio_path = os.path.join(temp_dir, f"{os.path.basename(video_path)}_audio.wav")
            if os.path.exists(audio_path): return audio_path
            
            print(f"[VideoSampler] [LAZY] Extracting audio: {video_path}")
            subprocess.run([
                "ffmpeg", "-i", video_path, 
                "-vn", "-acodec", "pcm_s16le", "-ar", "16000", "-ac", "1", 
                audio_path, "-y", "-loglevel", "quiet"
            ], check=True)
            return audio_path
        except Exception as e:
            print(f"[VideoSampler] Audio Extraction Failed: {e}")
            return None

    def extract_frames(self, video_path: str, count: int = 8, 특정_indices: list = None):
        """
        V11.1 CPU Optimized: Single-Pass FFmpeg Batch Extraction.
        Extracts all frames in one command, drastically reducing process overhead.
        """
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened(): return [], []
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        cap.release()

        if 특정_indices:
            indices = 특정_indices
        else:
            if total_frames <= count:
                indices = list(range(total_frames))
            else:
                indices = np.linspace(0, total_frames - 1, count, dtype=int)
        
        frames_np = []
        frames_pil = []
        
        # Create a select filter string for the indices
        # Example: select='eq(n\,0)+eq(n\,10)+eq(n\,20)'
        select_filter = "+".join([f"eq(n\,{idx})" for idx in indices])
        
        with tempfile.TemporaryDirectory() as tmpdir:
            out_pattern = os.path.join(tmpdir, "frame_%03d.jpg")
            
            try:
                # Single pass extraction
                cmd = [
                    "ffmpeg", "-i", video_path,
                    "-vf", f"select='{select_filter}'",
                    "-vsync", "0", # vsync 0 is more reliable for select filter
                    "-q:v", "2",
                    out_pattern, "-y", "-loglevel", "quiet"
                ]
                subprocess.run(cmd, check=True)
                
                # Read back the files
                # FFmpeg with vsync 0/vfr might name files frame_001, frame_002...
                # We sort them to ensure chronological order matching our indices
                saved_files = sorted([f for f in os.listdir(tmpdir) if f.startswith("frame_")])
                for f_name in saved_files:
                    out_path = os.path.join(tmpdir, f_name)
                    img = Image.open(out_path).convert("RGB")
                    frames_pil.append(img)
                    frames_np.append(np.array(img))
            except Exception as e:
                print(f"[VideoSampler] Batch FFmpeg failed: {e}. Falling back to iterative extraction.")
                # Fallback: Extract frames one by one if batch fails
                for idx in indices:
                    timestamp = idx / 30.0 # Heuristic if FPS unknown, or use cap.get
                    out_path = os.path.join(tmpdir, f"fb_{idx}.jpg")
                    subprocess.run([
                        "ffmpeg", "-ss", str(max(0, timestamp - 0.1)), "-i", video_path,
                        "-frames:v", "1", out_path, "-y", "-loglevel", "quiet"
                    ])
                    if os.path.exists(out_path):
                        img = Image.open(out_path).convert("RGB")
                        frames_pil.append(img)
                        frames_np.append(np.array(img))

        return frames_np, frames_pil

    def extract_8_frames(self, video_path: str):
        """
        Legacy Method for V10 Pipeline.
        Refactored to NOT extract audio here (Lazy audio is handled by the pipeline).
        """
        info = self.get_info(video_path)
        frames_np, frames_pil = self.extract_frames(video_path, count=8)
        
        return frames_np, frames_pil, {**info, "audio_path": None}