"""
Video processing operations for YOLO dataset creation.

Provides utilities to extract frames from video files for dataset creation.
"""

import concurrent.futures
import os
import threading
from pathlib import Path

os.environ.setdefault("NO_ALBUMENTATIONS_UPDATE", "1")

import cv2

from ydt.core.logger import get_logger

logger = get_logger(__name__)

# Global stop event for graceful shutdown
_stop_event = threading.Event()


def extract_frames(
    video_path: str | Path,
    frames_output_dir: str | Path,
    step: int = 40,
    supported_formats: list[str] = None,
) -> int:
    """
    Extract frames from video files with specified interval.

    Can process single video file or directory containing multiple videos.
    Creates separate output directory for each video.

    Args:
        video_path: Path to video file or directory containing videos
        frames_output_dir: Root output directory for extracted frames
        step: Frame extraction interval (extract every Nth frame)
        supported_formats: List of supported video extensions

    Returns:
        Total number of frames extracted across all videos

    Raises:
        FileNotFoundError: If video path doesn't exist
        ValueError: If no supported video files found

    Examples:
        >>> # Extract from single video
        >>> extract_frames("video.mp4", "./frames", step=30)
        120

        >>> # Extract from directory of videos
        >>> extract_frames("./videos", "./output", step=40)
        450
    """
    if supported_formats is None:
        supported_formats = [".mp4", ".avi", ".mkv", ".mov", ".flv", ".wmv", ".mpg", ".mpeg"]

    video_path = Path(video_path)
    frames_output_dir = Path(frames_output_dir)

    if not video_path.exists():
        raise FileNotFoundError(f"Video path not found: {video_path}")

    # Ensure output directory exists
    frames_output_dir.mkdir(parents=True, exist_ok=True)

    # Find all video files
    video_files = []

    if video_path.is_file():
        # Single file
        if video_path.suffix in supported_formats:
            video_files.append(video_path)
        else:
            raise ValueError(f"Unsupported video format: {video_path.suffix}")
    else:
        # Directory - find all video files
        # Use a set to avoid duplicates (rglob is case-insensitive on Windows)
        video_files_set = set()
        for ext in supported_formats:
            video_files_set.update(video_path.rglob(f"*{ext}"))
            video_files_set.update(video_path.rglob(f"*{ext.upper()}"))
        video_files = sorted(video_files_set)

    if not video_files:
        raise ValueError(f"No supported video files found in: {video_path}")

    logger.info(f"Found {len(video_files)} video files")

    total_saved_count = 0
    processed_videos = 0

    try:
        # Process each video file
        for video_file in video_files:
            logger.info(f"Processing video: {video_file.name}")

            # Create output directory for this video
            video_name = video_file.stem
            video_output_dir = frames_output_dir / f"{video_name}_frames"
            video_output_dir.mkdir(exist_ok=True)

            # Open video file
            cap = cv2.VideoCapture(str(video_file))
            if not cap.isOpened():
                logger.warning(f"Cannot open video file: {video_file}")
                continue

            # Get video information and print
            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            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))
            duration = total_frames / fps if fps > 0 else 0

            logger.info(f"  Total frames: {total_frames}")
            logger.info(f"  FPS: {fps:.2f}")
            logger.info(f"  Resolution: {width}x{height}")
            logger.info(f"  Duration: {duration:.2f} seconds")
            logger.info(f"  Output directory: {video_output_dir}")

            frame_count = 0
            saved_count = 0

            try:
                # Extract frames
                while True:
                    ret, frame = cap.read()
                    if not ret:
                        break

                    if frame_count % step == 0:
                        # Generate filename
                        output_path = video_output_dir / f"frame_{frame_count:06d}.jpg"

                        # Handle existing files
                        counter = 1
                        while output_path.exists():
                            output_path = (
                                video_output_dir / f"frame_{frame_count:06d}_{counter}.jpg"
                            )
                            counter += 1

                        # Save frame
                        cv2.imwrite(str(output_path), frame)
                        saved_count += 1

                        # Progress update
                        if saved_count % 10 == 0:
                            progress = (frame_count / total_frames) * 100
                            logger.info(
                                f"  Progress: {frame_count}/{total_frames} frames ({progress:.1f}%), "
                                f"saved: {saved_count} images"
                            )

                    frame_count += 1

            finally:
                # Always release video capture
                cap.release()

            logger.info(f"  Completed: {video_file.name}")
            logger.info(f"  Processed frames: {frame_count}")
            logger.info(f"  Saved images: {saved_count}")
            total_saved_count += saved_count
            processed_videos += 1

        logger.info("All videos processed!")
        logger.info(f"Total images saved: {total_saved_count}")
        logger.info(f"Output directory: {frames_output_dir}")

    except KeyboardInterrupt:
        logger.warning("\nKeyboardInterrupt received! Stopping gracefully...")
        logger.info(f"Processed {processed_videos}/{len(video_files)} videos")
        logger.info(f"Total images saved before interruption: {total_saved_count}")
        raise  # Re-raise to exit properly

    return total_saved_count


def _process_single_video(video_file: Path, frames_output_dir: Path, step: int) -> int:
    """
    Process a single video file (internal function for threading).

    Args:
        video_file: Path to video file
        frames_output_dir: Root output directory
        step: Frame extraction interval

    Returns:
        Number of frames extracted from this video
    """
    # Create thread-local logger for thread safety
    thread_logger = get_logger(f"{__name__}_thread_{threading.get_ident()}")

    # Create output directory for this video
    video_name = video_file.stem
    video_output_dir = frames_output_dir / f"{video_name}_frames"
    video_output_dir.mkdir(exist_ok=True)

    # Open video file
    cap = cv2.VideoCapture(str(video_file))
    if not cap.isOpened():
        thread_logger.warning(f"Cannot open video file: {video_file}")
        return 0

    # Get video information
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    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 / fps if fps > 0 else 0

    thread_logger.info(f"Processing: {video_file.name}")
    thread_logger.info(
        f"  Total frames: {total_frames}, FPS: {fps:.2f}, Resolution: {width}x{height}"
    )

    frame_count = 0
    saved_count = 0

    # Extract frames
    while True:
        # Check for stop signal
        if _stop_event.is_set():
            thread_logger.info(f"Stopping: {video_file.name} (interrupted)")
            break

        ret, frame = cap.read()
        if not ret:
            break

        if frame_count % step == 0:
            # Generate filename
            output_path = video_output_dir / f"frame_{frame_count:06d}.jpg"

            # Handle existing files
            counter = 1
            while output_path.exists():
                output_path = video_output_dir / f"frame_{frame_count:06d}_{counter}.jpg"
                counter += 1

            # Save frame
            cv2.imwrite(str(output_path), frame)
            saved_count += 1

            # Progress update (less frequent for multi-threading)
            if saved_count % 50 == 0:
                progress = (frame_count / total_frames) * 100
                thread_logger.info(
                    f"  {video_file.name}: {frame_count}/{total_frames} ({progress:.1f}%), saved: {saved_count}"
                )

        frame_count += 1

    # Release resources
    cap.release()

    thread_logger.info(f"Completed: {video_file.name} - Saved {saved_count} frames")
    return saved_count


def extract_frames_parallel(
    video_path: str | Path,
    frames_output_dir: str | Path,
    step: int = 40,
    max_workers: int = None,
    supported_formats: list[str] = None,
) -> int:
    """
    Extract frames from video files using parallel processing for multiple videos.

    Process multiple video files concurrently using thread pools for faster processing.
    Creates separate output directory for each video.

    Args:
        video_path: Path to video file or directory containing videos
        frames_output_dir: Root output directory for extracted frames
        step: Frame extraction interval (extract every Nth frame)
        max_workers: Maximum number of concurrent workers (default: CPU count)
        supported_formats: List of supported video extensions

    Returns:
        Total number of frames extracted across all videos

    Raises:
        FileNotFoundError: If video path doesn't exist
        ValueError: If no supported video files found

    Examples:
        >>> # Extract from directory with parallel processing
        >>> extract_frames_parallel("./videos", "./output", step=40, max_workers=4)
        450

        >>> # Auto-detect worker count
        >>> extract_frames_parallel("./videos", "./output", step=30)
        1200
    """
    if supported_formats is None:
        supported_formats = [".mp4", ".avi", ".mkv", ".mov", ".flv", ".wmv", ".mpg", ".mpeg"]

    video_path = Path(video_path)
    frames_output_dir = Path(frames_output_dir)

    if not video_path.exists():
        raise FileNotFoundError(f"Video path not found: {video_path}")

    # Ensure output directory exists
    frames_output_dir.mkdir(parents=True, exist_ok=True)

    # Find all video files
    video_files = []

    if video_path.is_file():
        # Single file
        if video_path.suffix in supported_formats:
            video_files.append(video_path)
        else:
            raise ValueError(f"Unsupported video format: {video_path.suffix}")
    else:
        # Directory - find all video files
        # Use a set to avoid duplicates (rglob is case-insensitive on Windows)
        video_files_set = set()
        for ext in supported_formats:
            video_files_set.update(video_path.rglob(f"*{ext}"))
            video_files_set.update(video_path.rglob(f"*{ext.upper()}"))
        video_files = sorted(video_files_set)

    if not video_files:
        raise ValueError(f"No supported video files found in: {video_path}")

    # If only one video file, use regular processing
    if len(video_files) == 1:
        logger.info("Single video file detected, using regular processing")
        return extract_frames(video_path, frames_output_dir, step, supported_formats)

    # Determine worker count
    import os

    if max_workers is None:
        max_workers = min(len(video_files), os.cpu_count() or 4)

    max_workers = min(max_workers, len(video_files))  # Don't exceed video count
    max_workers = max(1, max_workers)  # At least 1 worker

    logger.info(f"Found {len(video_files)} video files")
    logger.info(f"Using {max_workers} parallel workers for processing")

    total_saved_count = 0
    completed_count = 0

    # Clear stop event at start
    _stop_event.clear()

    # Process videos in parallel
    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
        # Submit all video processing tasks
        future_to_video = {
            executor.submit(_process_single_video, video_file, frames_output_dir, step): video_file
            for video_file in video_files
        }

        try:
            # Collect results with timeout to allow interrupt checking
            pending = set(future_to_video.keys())
            while pending:
                # Wait with short timeout to check for interrupts
                done, pending = concurrent.futures.wait(
                    pending, timeout=0.5, return_when=concurrent.futures.FIRST_COMPLETED
                )

                for future in done:
                    video_file = future_to_video[future]
                    completed_count += 1

                    try:
                        saved_count = future.result()
                        total_saved_count += saved_count
                        logger.info(
                            f"Progress: {completed_count}/{len(video_files)} videos completed"
                        )
                    except Exception as e:
                        logger.error(f"Error processing {video_file.name}: {e}")

            logger.info("All videos processed in parallel!")
            logger.info(f"Total images saved: {total_saved_count}")
            logger.info(f"Output directory: {frames_output_dir}")

        except KeyboardInterrupt:
            logger.warning("\nKeyboardInterrupt received! Stopping all tasks...")

            # Signal all running threads to stop
            _stop_event.set()

            # Cancel all pending futures
            for future in pending:
                future.cancel()

            # Wait briefly for threads to stop
            concurrent.futures.wait(pending, timeout=2.0)

            logger.info(f"Stopped. Processed {completed_count}/{len(video_files)} videos")
            logger.info(f"Total images saved before interruption: {total_saved_count}")

            # Clear stop event for next run
            _stop_event.clear()
            raise

    return total_saved_count
