"""NVDEC Worker Manager for StreamingGateway integration.

This module provides a manager for the NVDEC hardware decoding backend.
NVDEC now supports fully dynamic camera configuration at runtime (add, remove, update),
restarting workers as needed. Outputs to CUDA IPC ring buffers (NV12 format) for zero-copy GPU inference pipelines.
"""

import logging
import multiprocessing as mp
import time
from typing import Dict, List, Optional, Any

from .nvdec import (
    nvdec_pool_process,
    StreamConfig,
    DemuxerType,
    CUPY_AVAILABLE,
    PYNVCODEC_AVAILABLE,
    RING_BUFFER_AVAILABLE,
)

logger = logging.getLogger(__name__)


def is_nvdec_available() -> bool:
    """Check if NVDEC backend is available.

    Requires:
        - CuPy with CUDA support
        - PyNvVideoCodec for NVDEC hardware decode
        - cuda_shm_ring_buffer module for CUDA IPC
    """
    return CUPY_AVAILABLE and PYNVCODEC_AVAILABLE and RING_BUFFER_AVAILABLE


def get_available_gpu_count() -> int:
    """Detect the number of available CUDA GPUs.

    Returns:
        Number of available GPUs, or 1 if detection fails.
    """
    if not CUPY_AVAILABLE:
        return 1

    try:
        import cupy as cp
        return cp.cuda.runtime.getDeviceCount()
    except Exception as e:
        logger.warning(f"Failed to detect GPU count: {e}, defaulting to 1")
        return 1


class NVDECWorkerManager:
    """Manager for NVDEC worker processes - fully dynamic camera configuration.

    This manager wraps the existing nvdec_pool_process function to integrate
    with StreamingGateway. Key features:

    - Fully dynamic camera configuration (add, remove, update at runtime)
        by restarting workers on any camera config change
    - Outputs to CUDA IPC ring buffers (not Redis/Kafka)
    - NV12 format output (50% smaller than RGB)
    - One worker process per GPU
    """

    def __init__(
        self,
        camera_configs: List[Dict[str, Any]],
        stream_config: Dict[str, Any],  # Unused but kept for interface consistency
        gpu_id: int = 0,
        num_gpus: int = 0,  # 0 = auto-detect all available GPUs
        nvdec_pool_size: int = 8,
        nvdec_burst_size: int = 4,
        frame_width: int = 640,
        frame_height: int = 640,
        num_slots: int = 32,
        target_fps: int = 0,  # 0 = use per-camera FPS from config
        duration_sec: float = 0,  # 0 = infinite
        demuxer_type: str = "nvc",  # "nvc" or "gstreamer" - gstreamer provides ABSOLUTE RTP timestamps
    ):
        """Initialize NVDEC Worker Manager.

        Args:
            camera_configs: List of camera configuration dicts with keys:
                - camera_id or stream_key: Unique identifier (used for ring buffer naming)
                - source: Video file path or RTSP URL
                - width: Optional frame width (default: frame_width)
                - height: Optional frame height (default: frame_height)
                - fps: FPS limit for this camera (used by default)
            stream_config: Stream configuration (unused, for interface consistency)
            gpu_id: Primary GPU device ID (starting GPU for round-robin assignment)
            num_gpus: Number of GPUs to use (0 = auto-detect all available GPUs)
            nvdec_pool_size: Number of NVDEC decoders per GPU
            nvdec_burst_size: Frames per stream before rotating to next
            frame_width: Default output frame width (used if camera config doesn't specify)
            frame_height: Default output frame height (used if camera config doesn't specify)
            num_slots: Ring buffer slots per camera
            target_fps: Global FPS override (0 = use per-camera FPS from config)
            duration_sec: Duration to run (0 = infinite until stop)
            demuxer_type: Demuxer backend - "nvc" (PyNvVideoCodec, fastest for files) or
                          "gstreamer" (provides ABSOLUTE RTP timestamps for RTSP streams)
        """
        if not is_nvdec_available():
            raise RuntimeError(
                "NVDEC not available. Requires CuPy, PyNvVideoCodec, and cuda_shm_ring_buffer"
            )

        self.camera_configs = camera_configs
        self.stream_config = stream_config
        self.gpu_id = gpu_id

        # Auto-detect GPUs if num_gpus is 0
        if num_gpus <= 0:
            detected_gpus = get_available_gpu_count()
            self.num_gpus = min(detected_gpus, 8)  # Max 8 GPUs
            logger.info(f"Auto-detected {detected_gpus} GPU(s), using {self.num_gpus}")
        else:
            self.num_gpus = min(num_gpus, 8)  # Max 8 GPUs
        self.nvdec_pool_size = nvdec_pool_size
        self.nvdec_burst_size = nvdec_burst_size
        self.frame_width = frame_width
        self.frame_height = frame_height
        self.num_slots = num_slots
        self.target_fps = target_fps
        self.duration_sec = duration_sec if duration_sec > 0 else float('inf')
        self.demuxer_type = demuxer_type  # "nvc" or "gstreamer"

        self._workers: List[mp.Process] = []
        self._stop_event: Optional[mp.Event] = None  # type: ignore[valid-type]
        self._result_queue: Optional[mp.Queue] = None
        self._shared_frame_count: Optional[mp.Value] = None  # type: ignore[valid-type]
        self._gpu_frame_counts: Dict[int, Any] = {}  # Per-GPU counters (mp.Value)
        self._start_time: Optional[float] = None
        self._is_running = False

        # Convert camera configs to StreamConfig objects and assign to GPUs
        self._stream_configs: List[StreamConfig] = []
        self._gpu_camera_assignments: Dict[int, List[StreamConfig]] = {
            i: [] for i in range(self.num_gpus)
        }
        self._camera_to_gpu: Dict[str, int] = {}

        self._prepare_camera_configs()

        logger.info(
            f"NVDECWorkerManager initialized: {len(camera_configs)} cameras, "
            f"{self.num_gpus} GPU(s), pool_size={nvdec_pool_size}, demuxer={demuxer_type}"
        )

    def _prepare_camera_configs(self):
        """Convert dict configs to StreamConfig and distribute across GPUs.

        Ring buffers are named using camera_id for SHM identification.
        Per-camera FPS from config is used by default (target_fps=0 means use config FPS).
        """
        # Clear existing configs to prevent duplicates on re-prepare
        self._stream_configs.clear()
        self._gpu_camera_assignments = {i: [] for i in range(self.num_gpus)}
        self._camera_to_gpu.clear()

        for i, config in enumerate(self.camera_configs):
            # Extract camera ID (support both camera_id and stream_key)
            # This ID is used for naming the CUDA IPC ring buffer
            camera_id = config.get('camera_id') or config.get('stream_key') or f"cam_{i:04d}"

            # Extract video source
            source = config.get('source') or config.get('video_path')
            if not source:
                logger.warning(f"Camera {camera_id} has no source, skipping")
                continue

            # Extract dimensions (use per-camera config or fallback to defaults)
            width = config.get('width') or self.frame_width
            height = config.get('height') or self.frame_height

            # Determine FPS: use global override if set, otherwise per-camera FPS from config
            if self.target_fps > 0:
                # Global FPS override is set
                fps = self.target_fps
            else:
                # Use per-camera FPS from config (default streaming FPS)
                fps = config.get('fps', 10)  # Default to 10 FPS if not specified

            # Assign to GPU (round-robin starting from gpu_id)
            gpu_id = (self.gpu_id + i) % self.num_gpus

            stream_config = StreamConfig(
                camera_id=camera_id,
                video_path=source,
                width=width,
                height=height,
                target_fps=fps,
                gpu_id=gpu_id,
                demuxer_type=self.demuxer_type,
            )

            self._stream_configs.append(stream_config)
            self._gpu_camera_assignments[gpu_id].append(stream_config)
            self._camera_to_gpu[camera_id] = gpu_id

            logger.debug(f"Camera {camera_id}: source={source}, {width}x{height}@{fps}fps, GPU{gpu_id}")

    def start(self) -> None:
        """Start NVDEC worker processes (one per GPU)."""
        if self._is_running:
            logger.warning("NVDECWorkerManager is already running")
            return

        if not self._stream_configs:
            logger.warning("No cameras configured, nothing to start")
            return

        ctx = mp.get_context("spawn")
        self._stop_event = ctx.Event()
        self._result_queue = ctx.Queue()
        self._shared_frame_count = ctx.Value('L', 0)  # Global counter (all GPUs)
        self._start_time = time.perf_counter()

        # Create per-GPU frame counters
        self._gpu_frame_counts = {}
        for gpu_id in range(self.num_gpus):
            if self._gpu_camera_assignments[gpu_id]:  # Only if GPU has cameras
                self._gpu_frame_counts[gpu_id] = ctx.Value('L', 0)

        total_num_streams = len(self._stream_configs)
        total_num_gpus = len([g for g in range(self.num_gpus) if self._gpu_camera_assignments[g]])

        logger.info(f"Starting NVDEC: {total_num_streams} cameras across {total_num_gpus} GPUs")

        # Start one process per GPU that has cameras
        for gpu_id in range(self.num_gpus):
            gpu_cameras = self._gpu_camera_assignments[gpu_id]
            if not gpu_cameras:
                continue

            p = ctx.Process(
                target=nvdec_pool_process,
                args=(
                    gpu_id,                     # process_id
                    gpu_cameras,                # camera_configs (List[StreamConfig])
                    self.nvdec_pool_size,       # pool_size
                    self.duration_sec,          # duration_sec
                    self._result_queue,         # result_queue
                    self._stop_event,           # stop_event
                    self.nvdec_burst_size,      # burst_size
                    self.num_slots,             # num_slots
                    self.target_fps,            # target_fps
                    self._shared_frame_count,   # shared_frame_count (global)
                    self._gpu_frame_counts,     # gpu_frame_counts (per-GPU dict)
                    total_num_streams,          # total_num_streams
                    total_num_gpus,             # total_num_gpus
                    self.demuxer_type,          # demuxer_type
                ),
                name=f"NVDECWorker-GPU{gpu_id}",
                daemon=False,
            )
            p.start()
            self._workers.append(p)  # type: ignore[arg-type]
            logger.info(f"Started NVDEC worker on GPU {gpu_id} with {len(gpu_cameras)} cameras")

        self._is_running = True
        logger.info(f"NVDECWorkerManager started: {len(self._workers)} workers")

    def stop(self, timeout: float = 15.0) -> None:
        """Stop all worker processes.

        Args:
            timeout: Maximum time to wait for each worker to stop gracefully
        """
        if not self._is_running:
            logger.warning("NVDECWorkerManager is not running")
            return

        logger.info("Stopping NVDECWorkerManager...")

        # Signal workers to stop
        if self._stop_event:
            self._stop_event.set()  # type: ignore[attr-defined]

        # Wait for workers to finish
        for p in self._workers:
            p.join(timeout=timeout)
            if p.is_alive():
                logger.warning(f"Worker {p.name} did not stop gracefully, terminating")
                p.terminate()
                p.join(timeout=2.0)

        self._workers.clear()
        self._is_running = False
        logger.info("NVDECWorkerManager stopped")

    def get_worker_statistics(self) -> Dict[str, Any]:
        """Return statistics from workers.

        Returns:
            Dict with keys:
                - num_workers: Number of worker processes
                - running_workers: Number of currently running workers
                - total_cameras: Total cameras across all workers
                - gpu_assignments: Cameras per GPU
                - total_frames: Total frames processed (from shared counter)
                - elapsed_sec: Time since start
                - aggregate_fps: Overall FPS
                - per_stream_fps: Average FPS per camera
                - backend: 'nvdec'
                - gpu_results: Per-GPU results from result queue
        """
        stats = {
            'backend': 'nvdec',
            'num_workers': len(self._workers),
            'running_workers': sum(1 for p in self._workers if p.is_alive()),
            'total_cameras': len(self._stream_configs),
            'gpu_assignments': {
                gpu_id: len(cameras)
                for gpu_id, cameras in self._gpu_camera_assignments.items()
            },
            'nvdec_config': {
                'gpu_id': self.gpu_id,
                'num_gpus': self.num_gpus,
                'pool_size': self.nvdec_pool_size,
                'burst_size': self.nvdec_burst_size,
                'frame_size': f"{self.frame_width}x{self.frame_height}",
                'num_slots': self.num_slots,
                'target_fps': self.target_fps,
                'demuxer_type': self.demuxer_type,
            },
        }

        # Add frame count and FPS
        if self._shared_frame_count:
            total_frames = self._shared_frame_count.value  # type: ignore[attr-defined]
            stats['total_frames'] = total_frames

            if self._start_time:
                elapsed = time.perf_counter() - self._start_time
                stats['elapsed_sec'] = elapsed
                stats['aggregate_fps'] = total_frames / elapsed if elapsed > 0 else 0
                stats['per_stream_fps'] = (
                    stats['aggregate_fps'] / len(self._stream_configs)  # type: ignore[operator]
                    if self._stream_configs else 0
                )

        # Add per-GPU frame counts and FPS
        if self._gpu_frame_counts and self._start_time:
            elapsed = time.perf_counter() - self._start_time
            gpu_stats = {}
            for gpu_id, counter in self._gpu_frame_counts.items():
                gpu_frames = counter.value  # type: ignore[attr-defined]
                num_cams = len(self._gpu_camera_assignments.get(gpu_id, []))
                gpu_fps = gpu_frames / elapsed if elapsed > 0 else 0
                gpu_per_cam = gpu_fps / num_cams if num_cams > 0 else 0
                gpu_stats[f'GPU{gpu_id}'] = {
                    'frames': gpu_frames,
                    'cameras': num_cams,
                    'fps': gpu_fps,
                    'fps_per_cam': gpu_per_cam,
                }
            stats['per_gpu_stats'] = gpu_stats

        # Collect any available results from queue (non-blocking)
        gpu_results = []
        if self._result_queue:
            while True:
                try:
                    result = self._result_queue.get_nowait()
                    gpu_results.append(result)
                except:
                    break
        stats['gpu_results'] = gpu_results

        return stats

    def get_camera_assignments(self) -> Dict[str, int]:
        """Return mapping of camera_id to GPU ID.

        Returns:
            Dict mapping camera_id -> gpu_id
        """
        return self._camera_to_gpu.copy()

    def add_camera(self, camera_config: Dict[str, Any]) -> bool:
        """Add a new camera at runtime by restarting NVDEC workers with updated config."""
        camera_id = camera_config.get('camera_id') or camera_config.get('stream_key')
        if not camera_id:
            logger.error("Camera config missing camera_id/stream_key")
            return False
        # Prevent duplicate
        if camera_id in self._camera_to_gpu:
            logger.warning(f"Camera {camera_id} already exists in NVDEC config")
            return False
        self.camera_configs.append(camera_config)
        self._prepare_camera_configs()
        self.restart_workers()
        logger.info(f"Added camera {camera_id} and restarted NVDEC workers")
        return True

    def remove_camera(self, stream_key: str) -> bool:
        """Remove a camera at runtime by restarting NVDEC workers with updated config."""
        found = False
        for i, config in enumerate(self.camera_configs):
            cid = config.get('camera_id') or config.get('stream_key')
            if cid == stream_key:
                del self.camera_configs[i]
                found = True
                break
        if not found:
            logger.warning(f"Camera {stream_key} not found in NVDEC config")
            return False
        self._prepare_camera_configs()
        self.restart_workers()
        logger.info(f"Removed camera {stream_key} and restarted NVDEC workers")
        return True

    def update_camera(self, camera_config: Dict[str, Any]) -> bool:
        """Update a camera at runtime by restarting NVDEC workers with updated config."""
        camera_id = camera_config.get('camera_id') or camera_config.get('stream_key')
        if not camera_id:
            logger.error("Camera config missing camera_id/stream_key")
            return False
        found = False
        for i, config in enumerate(self.camera_configs):
            cid = config.get('camera_id') or config.get('stream_key')
            if cid == camera_id:
                self.camera_configs[i] = camera_config
                found = True
                break
        if not found:
            logger.warning(f"Camera {camera_id} not found in NVDEC config for update")
            return False
        self._prepare_camera_configs()
        self.restart_workers()
        logger.info(f"Updated camera {camera_id} and restarted NVDEC workers")
        return True

    def restart_workers(self):
        """Stop and restart all NVDEC workers with the current camera config."""
        self.stop()
        self.start()

    @property
    def is_running(self) -> bool:
        """Check if the manager is currently running."""
        return self._is_running

    def __enter__(self):
        """Context manager entry."""
        self.start()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self.stop()
