#!/usr/bin/env python3
"""Orin NVDEC Decoder Pool - GStreamer subprocess decode for Jetson Orin.

On Jetson Orin, the desktop CUVID API (cuvidGetDecoderCaps, cuvidCreateDecoder, etc.)
does NOT exist. PyNvVideoCodec fails with "Could not load function cuvidGetDecoderCaps".

Additionally, GStreamer Python bindings have plugin loading issues in
multiprocessing.spawn() worker processes (avdec_h264 and other libav plugins
fail to load in the spawned worker's GStreamer registry).

This module provides OrinNVDECDecoderPool as a drop-in replacement for
NVDECDecoderPool. Each camera gets its own gst-launch-1.0 subprocess that
writes raw NV12 frames to stdout via fdsink.

Activated by: MATRICE_PLATFORM=orin environment variable.

Pipeline per camera (subprocess):
    gst-launch-1.0 rtspsrc -> rtph264depay -> h264parse -> avdec_h264
    -> videoconvert -> videoscale -> video/x-raw,NV12,640x640 -> fdsink fd=1

On Orin unified memory, CPU decode is efficient since CPU and GPU share DRAM.
"""

from __future__ import annotations

import os
import shutil
import signal
import logging
import subprocess
import threading
import time
import uuid
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple

import numpy as np

logger = logging.getLogger(__name__)

try:
    import cupy as cp
    CUPY_AVAILABLE = True
except ImportError:
    cp = None  # type: ignore[assignment]
    CUPY_AVAILABLE = False

TARGET_WIDTH = 640
TARGET_HEIGHT = 640
NV12_HEIGHT = 960  # 640 * 1.5 for NV12
FRAME_SIZE = TARGET_WIDTH * NV12_HEIGHT  # 614,400 bytes


@dataclass
class OrinStreamState:
    """State for a gst-launch-1.0 subprocess camera stream."""
    stream_id: int
    camera_id: str
    video_path: str
    width: int = TARGET_WIDTH
    height: int = TARGET_HEIGHT
    stream_type: str = "rtsp"
    process: Any = None  # subprocess.Popen
    source_fps: float = 30.0
    frames_decoded: int = 0
    session_id: str = ""
    session_start_ns: int = 0
    errors: int = 0
    _lock: threading.Lock = field(default_factory=threading.Lock)


class OrinNVDECDecoderPool:
    """Orin-compatible NVDEC decoder pool using gst-launch-1.0 subprocesses.

    Drop-in replacement for NVDECDecoderPool on Jetson Orin where CUVID API
    is not available. Uses gst-launch-1.0 subprocess per camera for decode
    and outputs NV12 CuPy tensors identical to the desktop path.

    API matches NVDECDecoderPool exactly:
        - assign_stream(stream_id, camera_id, video_path, ...)
        - decode_round(decoder_idx, frames_per_stream) -> (total, [(cam, tensor, ts, ...)])
        - get_camera_ids_for_decoder(decoder_idx)
        - get_source_fps_for_decoder(decoder_idx)
        - release()
    """

    def __init__(self, pool_size: int, gpu_id: int = 0, demuxer_type: str = "gstreamer"):
        if shutil.which('gst-launch-1.0') is None:
            raise RuntimeError(
                "gst-launch-1.0 not found. Install GStreamer: "
                "apt-get install gstreamer1.0-tools gstreamer1.0-plugins-base "
                "gstreamer1.0-plugins-good gstreamer1.0-plugins-bad gstreamer1.0-libav"
            )

        self.pool_size = pool_size
        self.gpu_id = gpu_id
        self.demuxer_type = demuxer_type
        self.decoders = [None] * pool_size  # Placeholder for API compat
        self.actual_pool_size = pool_size
        self.streams_per_decoder: List[List[OrinStreamState]] = [[] for _ in range(pool_size)]
        self._all_streams: Dict[str, OrinStreamState] = {}
        self._frame_size = FRAME_SIZE

        if CUPY_AVAILABLE and cp is not None:
            cp.cuda.Device(gpu_id).use()

        logger.info(
            f"[Orin] Created subprocess NVDEC pool: {pool_size} slots on GPU {gpu_id}, "
            f"frame_size={self._frame_size} bytes"
        )

    def _start_pipeline(self, stream: OrinStreamState) -> bool:
        """Start gst-launch-1.0 subprocess for a stream."""
        try:
            if stream.stream_type == "rtsp":
                pipeline_str = (
                    f"rtspsrc location={stream.video_path} latency=100 "
                    f"buffer-mode=auto protocols=tcp ! "
                    f"rtph264depay ! h264parse ! avdec_h264 ! "
                    f"videoconvert ! videoscale ! "
                    f"video/x-raw,format=NV12,width={stream.width},height={stream.height} ! "
                    f"fdsink fd=1"
                )
            else:
                pipeline_str = (
                    f"filesrc location={stream.video_path} ! decodebin ! "
                    f"videoconvert ! videoscale ! "
                    f"video/x-raw,format=NV12,width={stream.width},height={stream.height} ! "
                    f"fdsink fd=1"
                )

            cmd = f"gst-launch-1.0 -q -e {pipeline_str}"
            logger.info(f"[Orin] {stream.camera_id}: Starting subprocess pipeline")
            logger.debug(f"[Orin] Command: {cmd}")

            proc = subprocess.Popen(
                cmd,
                shell=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                bufsize=self._frame_size * 4,  # Buffer 4 frames
                preexec_fn=os.setsid,  # Own process group for cleanup
            )

            # Give the pipeline a moment to start and check it didn't exit
            time.sleep(0.5)
            if proc.poll() is not None:
                stderr = proc.stderr.read().decode('utf-8', errors='replace')
                logger.error(
                    f"[Orin] {stream.camera_id}: Pipeline exited immediately "
                    f"(rc={proc.returncode}): {stderr[:500]}"
                )
                return False

            stream.process = proc
            stream.session_start_ns = time.time_ns()
            stream.session_id = str(uuid.uuid4())[:8]
            stream.frames_decoded = 0
            stream.errors = 0

            logger.info(
                f"[Orin] {stream.camera_id}: Subprocess pipeline started "
                f"(pid={proc.pid}, session={stream.session_id})"
            )
            return True

        except Exception as e:
            logger.error(f"[Orin] {stream.camera_id}: Pipeline start failed: {e}")
            return False

    def _stop_pipeline(self, stream: OrinStreamState):
        """Stop gst-launch-1.0 subprocess for a stream."""
        try:
            if stream.process and stream.process.poll() is None:
                try:
                    os.killpg(os.getpgid(stream.process.pid), signal.SIGTERM)
                except (ProcessLookupError, PermissionError):
                    pass
                try:
                    stream.process.wait(timeout=3)
                except subprocess.TimeoutExpired:
                    try:
                        os.killpg(os.getpgid(stream.process.pid), signal.SIGKILL)
                    except (ProcessLookupError, PermissionError):
                        pass
            stream.process = None
        except Exception as e:
            logger.warning(f"[Orin] {stream.camera_id}: Pipeline stop error: {e}")

    def _restart_pipeline(self, stream: OrinStreamState) -> bool:
        """Restart subprocess pipeline on error."""
        self._stop_pipeline(stream)
        time.sleep(0.5)
        return self._start_pipeline(stream)

    def assign_stream(self, stream_id: int, camera_id: str, video_path: str,
                      width: int = TARGET_WIDTH, height: int = TARGET_HEIGHT,
                      stream_type: str = "file", demuxer_type: str = None) -> bool:
        """Assign a camera stream and start its gst-launch-1.0 subprocess."""
        if self.actual_pool_size == 0:
            return False

        decoder_idx = stream_id % self.actual_pool_size

        stream = OrinStreamState(
            stream_id=stream_id,
            camera_id=camera_id,
            video_path=video_path,
            width=width,
            height=height,
            stream_type=stream_type,
        )

        success = self._start_pipeline(stream)
        self.streams_per_decoder[decoder_idx].append(stream)
        self._all_streams[camera_id] = stream

        if success:
            logger.info(
                f"[Orin] {camera_id}: Assigned to decoder slot {decoder_idx} "
                f"({stream_type})"
            )
        else:
            logger.warning(f"[Orin] {camera_id}: Pipeline not started, will retry")

        return success

    def _pull_frame(self, stream: OrinStreamState) -> Optional[Tuple[Any, int]]:
        """Pull one NV12 frame from gst-launch-1.0 stdout pipe.

        Returns:
            (nv12_tensor, timestamp_ns) or None if no frame available.
        """
        if stream.process is None:
            return None

        # Check if process is still alive
        if stream.process.poll() is not None:
            stream.errors += 1
            if stream.errors <= 3:
                stderr = ""
                try:
                    stderr = stream.process.stderr.read().decode('utf-8', errors='replace')[:300]
                except Exception:
                    pass
                logger.warning(
                    f"[Orin] {stream.camera_id}: Subprocess exited "
                    f"(rc={stream.process.returncode}): {stderr}"
                )
            if stream.errors >= 3:
                logger.info(f"[Orin] {stream.camera_id}: Restarting pipeline")
                self._restart_pipeline(stream)
            return None

        try:
            # Read exactly one frame worth of bytes
            data = stream.process.stdout.read(self._frame_size)
            if not data or len(data) < self._frame_size:
                stream.errors += 1
                if stream.errors >= 10:
                    logger.warning(
                        f"[Orin] {stream.camera_id}: Too many read errors, "
                        f"restarting (got {len(data) if data else 0} bytes)"
                    )
                    self._restart_pipeline(stream)
                return None

            # Reshape to NV12 layout: (960, 640, 1)
            nv12_np = np.frombuffer(data, dtype=np.uint8).reshape(
                int(stream.height * 1.5), stream.width, 1
            )

            # Upload to GPU via CuPy (shared DRAM on Orin)
            tensor = cp.asarray(nv12_np)

            # Timestamp: session start + frame-count-based offset
            stream.frames_decoded += 1
            timestamp_ns = stream.session_start_ns + int(
                stream.frames_decoded * 1_000_000_000 / stream.source_fps
            )

            stream.errors = 0

            # Log first frame and periodic status
            if stream.frames_decoded == 1:
                logger.info(
                    f"[Orin] {stream.camera_id}: First frame decoded! "
                    f"shape={tensor.shape}, size={len(data)}"
                )
            elif stream.frames_decoded % 500 == 0:
                elapsed = (time.time_ns() - stream.session_start_ns) / 1e9
                fps = stream.frames_decoded / elapsed if elapsed > 0 else 0
                logger.info(
                    f"[Orin] {stream.camera_id}: {stream.frames_decoded} frames, "
                    f"{fps:.1f} fps"
                )

            return tensor, timestamp_ns

        except Exception as e:
            stream.errors += 1
            if stream.errors <= 3:
                logger.warning(f"[Orin] {stream.camera_id}: Pull frame error: {e}")
            if stream.errors >= 10:
                logger.warning(f"[Orin] {stream.camera_id}: Too many errors, restarting")
                self._restart_pipeline(stream)
            return None

    def decode_round(self, decoder_idx: int, frames_per_stream: int = 4,
                     target_h: int = TARGET_HEIGHT, target_w: int = TARGET_WIDTH
                     ) -> Tuple[int, List[Tuple[str, Any, int, str, str, int, Optional[int]]]]:
        """Decode frames from all streams assigned to this decoder slot.

        Returns same format as NVDECDecoderPool.decode_round:
            (total_frames, [(camera_id, nv12_tensor, timestamp_ns,
                             stream_type, session_id, session_start_ns, rtp_timestamp), ...])
        """
        if decoder_idx >= self.actual_pool_size:
            return 0, []

        streams = self.streams_per_decoder[decoder_idx]
        if not streams:
            return 0, []

        total_frames = 0
        decoded_frames: List[Tuple[str, Any, int, str, str, int, Optional[int]]] = []

        for stream in streams:
            if stream.process is None or (stream.process and stream.process.poll() is not None):
                self._start_pipeline(stream)
                continue

            for _ in range(frames_per_stream):
                result = self._pull_frame(stream)
                if result is None:
                    break

                tensor, timestamp_ns = result
                decoded_frames.append((
                    stream.camera_id,
                    tensor,
                    timestamp_ns,
                    stream.stream_type,
                    stream.session_id,
                    stream.session_start_ns,
                    None,  # rtp_timestamp (not available from avdec_h264)
                ))
                total_frames += 1

        return total_frames, decoded_frames

    def get_camera_ids_for_decoder(self, decoder_idx: int) -> List[str]:
        """Get camera IDs assigned to a decoder slot."""
        if decoder_idx >= self.actual_pool_size:
            return []
        return [s.camera_id for s in self.streams_per_decoder[decoder_idx]]

    def get_source_fps_for_decoder(self, decoder_idx: int) -> float:
        """Get average source FPS for streams on this decoder."""
        if decoder_idx >= self.actual_pool_size:
            return 30.0
        streams = self.streams_per_decoder[decoder_idx]
        if not streams:
            return 30.0
        return sum(s.source_fps for s in streams) / len(streams)

    def release(self):
        """Stop all subprocess pipelines."""
        for camera_id, stream in self._all_streams.items():
            self._stop_pipeline(stream)
            logger.info(f"[Orin] {camera_id}: Pipeline stopped")
        self._all_streams.clear()

    def __del__(self):
        try:
            self.release()
        except Exception:
            pass
