#!/usr/bin/env python3
"""GStreamer RTP Demuxer with ABSOLUTE RTP Timestamps.

This module provides RTSP and file demuxing using GStreamer, extracting H264 NAL units
with ABSOLUTE RTP timestamps that persist across reconnections.

Key features:
- Captures RAW 32-bit RTP timestamps via pad probe BEFORE depayloading
- RTCP Sender Reports for absolute NTP time mapping
- Outputs H264 NAL units for external decoder (PyNvVideoCodec)

Key difference from PyAV/FFmpeg:
- PyAV/FFmpeg normalizes PTS to start near 0 each session (relative timestamps)
- GStreamer provides the ACTUAL raw 32-bit RTP timestamp from the source
- Raw RTP timestamps are at 90kHz and persist across reconnections

Architecture (RTSP):
    RTSP -> rtspsrc -> [PAD PROBE captures raw timestamp] -> rtph264depay ->
         -> h264parse -> appsink (H264 NAL units)

Architecture (File):
    File -> filesrc -> qtdemux -> h264parse -> appsink (H264 NAL units)

Usage:
    from gstreamer_rtp_demuxer import GstRTPDemuxer

    demuxer = GstRTPDemuxer("rtsp://camera/stream")
    demuxer.open()
    for h264_bytes, rtp_ts, rtp_ts_ns, absolute_ns in demuxer.demux():
        # h264_bytes: H264 NAL unit bytes
        # rtp_ts: ABSOLUTE 32-bit RTP timestamp (NOT normalized!)
        # rtp_ts_ns: RTP timestamp converted to nanoseconds
        # absolute_ns: Absolute Unix time if RTCP SR available
        decoded_frame = nvc_decoder.Decode(h264_bytes)
    demuxer.close()
"""
from __future__ import annotations

import logging
import re
import threading
import time
import uuid
from collections import deque
from dataclasses import dataclass
from enum import Enum
from typing import Any, Iterator, Optional, Tuple

import gi
gi.require_version('Gst', '1.0')
gi.require_version('GstRtp', '1.0')
gi.require_version('GstApp', '1.0')
from gi.repository import Gst, GLib, GstRtp

Gst.init(None)

logger = logging.getLogger(__name__)

# RTP clock rate for video (H.264/H.265)
RTP_CLOCK_RATE = 90000


class DemuxerType(Enum):
    """Demuxer backend selection."""
    NVC = "nvc"           # PyNvVideoCodec demuxer (files, fastest)
    GSTREAMER = "gstreamer"  # GStreamer demuxer (RTSP with ABSOLUTE RTP timestamps)


@dataclass
class RTCPSenderReport:
    """Parsed RTCP Sender Report for RTP-to-NTP time mapping.

    Attributes:
        ssrc: Synchronization source identifier
        ntp_timestamp: 64-bit NTP timestamp from RTCP SR
        rtp_timestamp: 32-bit RTP timestamp at SR time
        unix_time: NTP timestamp converted to Unix time
        received_at_ns: Local time when SR was received (nanoseconds)
    """
    ssrc: int
    ntp_timestamp: int      # 64-bit NTP timestamp
    rtp_timestamp: int      # 32-bit RTP timestamp at SR time
    unix_time: float        # Converted to Unix timestamp
    received_at_ns: int     # Local time when SR was received


class RTCPTracker:
    """Tracks RTCP Sender Reports for RTP-to-absolute-time mapping.

    RTCP Sender Reports provide the mapping between RTP timestamps (at 90kHz)
    and absolute wall-clock time (NTP). This allows converting any RTP timestamp
    to Unix time for cross-camera synchronization.
    """

    NTP_UNIX_EPOCH_DIFF = 2208988800  # Seconds between 1900 and 1970

    def __init__(self, clock_rate: int = RTP_CLOCK_RATE):
        self.clock_rate = clock_rate
        self._latest_sr: Optional[RTCPSenderReport] = None
        self._lock = threading.Lock()
        self._have_sr = threading.Event()

    def update_from_session_stats(self, stats_str: str, ssrc: int) -> Optional[RTCPSenderReport]:
        """Parse RTCP SR from GStreamer session stats.

        Args:
            stats_str: GStreamer stats structure as string
            ssrc: Source SSRC identifier

        Returns:
            RTCPSenderReport if valid SR found, None otherwise
        """
        try:
            ntp_match = re.search(r'sr-ntptime=\([^)]+\)(\d+)', stats_str)
            rtp_match = re.search(r'sr-rtptime=\([^)]+\)(\d+)', stats_str)
            have_sr_match = re.search(r'have-sr=\([^)]+\)(true|false)', stats_str)

            if have_sr_match and have_sr_match.group(1) == 'true' and ntp_match and rtp_match:
                ntp_val = int(ntp_match.group(1))
                rtp_val = int(rtp_match.group(1))

                if ntp_val > 0:
                    # Convert 64-bit NTP to Unix time
                    ntp_seconds = ntp_val >> 32
                    ntp_fractions = ntp_val & 0xFFFFFFFF
                    unix_time = (ntp_seconds - self.NTP_UNIX_EPOCH_DIFF) + (ntp_fractions / (2**32))

                    sr = RTCPSenderReport(
                        ssrc=ssrc,
                        ntp_timestamp=ntp_val,
                        rtp_timestamp=rtp_val,
                        unix_time=unix_time,
                        received_at_ns=time.time_ns(),
                    )

                    with self._lock:
                        self._latest_sr = sr
                        self._have_sr.set()

                    return sr
        except Exception as e:
            logger.debug(f"Failed to parse RTCP SR: {e}")
        return None

    def rtp_to_unix_ns(self, rtp_timestamp: int) -> Optional[int]:
        """Convert RTP timestamp to absolute Unix time in nanoseconds.

        Uses the most recent RTCP Sender Report to map RTP timestamps to
        absolute wall-clock time. Handles 32-bit RTP timestamp rollover.

        Args:
            rtp_timestamp: 32-bit RTP timestamp to convert

        Returns:
            Unix time in nanoseconds, or None if no SR available
        """
        with self._lock:
            if self._latest_sr is None:
                return None
            sr = self._latest_sr

        # Handle 32-bit rollover with signed arithmetic
        diff = (rtp_timestamp - sr.rtp_timestamp) & 0xFFFFFFFF
        if diff > 0x7FFFFFFF:
            diff = diff - 0x100000000

        time_diff = diff / self.clock_rate
        unix_time = sr.unix_time + time_diff
        return int(unix_time * 1_000_000_000)

    def has_sender_report(self) -> bool:
        """Check if at least one RTCP Sender Report has been received."""
        return self._have_sr.is_set()

    @property
    def latest_sr(self) -> Optional[RTCPSenderReport]:
        """Get the most recent RTCP Sender Report."""
        with self._lock:
            return self._latest_sr


class GstRTPDemuxer:
    """GStreamer H264 demuxer with ABSOLUTE RTP timestamp capture.

    For RTSP streams:
    - Captures RAW 32-bit RTP timestamps via pad probe BEFORE depayloading
    - These are the ACTUAL timestamps from the RTP packet header
    - NOT normalized like PyAV (which starts near 0 each session)
    - Timestamps persist across reconnections (monotonically increasing from source)
    - Automatically converts localhost to 127.0.0.1 to avoid IPv6 issues
    - Uses TCP by default to avoid UDP IPv6 socket issues

    For files:
    - Uses container PTS (presentation timestamp)
    - No RTP timestamps available

    Outputs H264 NAL units for external decoder (e.g., PyNvVideoCodec).

    Args:
        source: RTSP URL or file path
        use_tcp: Use TCP for RTSP transport (default: True, more reliable)
        loop: Loop file playback (default: False)
    """

    # Class-level retry configuration
    MAX_RETRIES = 2
    RETRY_DELAY = 1.0  # seconds

    @staticmethod
    def _normalize_rtsp_url(url: str) -> str:
        """Convert localhost to 127.0.0.1 to avoid IPv6 issues.

        GStreamer's rtspsrc has issues with IPv6 UDP sockets on some systems.
        Using explicit IPv4 address avoids 'Invalid address family (got 10)' errors.
        """
        if url.lower().startswith(('rtsp://', 'rtsps://')):
            # Replace localhost with 127.0.0.1 (case-insensitive)
            url = re.sub(r'(rtsp://|rtsps://)localhost([:/])', r'\g<1>127.0.0.1\2', url, flags=re.IGNORECASE)
            url = re.sub(r'(rtsp://|rtsps://)localhost$', r'\g<1>127.0.0.1', url, flags=re.IGNORECASE)
        return url

    def __init__(self, source: str, use_tcp: bool = True, loop: bool = False):
        # Normalize RTSP URL to avoid IPv6 issues
        self.source = self._normalize_rtsp_url(source)
        self.use_tcp = use_tcp
        self.loop = loop
        self._original_source = source  # Keep original for logging

        # Detect source type
        self._is_rtsp = self.source.lower().startswith(('rtsp://', 'rtsps://'))

        # GStreamer state
        self._pipeline: Optional[Gst.Pipeline] = None
        self._appsink: Optional[Gst.Element] = None
        self._depay: Optional[Gst.Element] = None
        self._rtspsrc: Optional[Gst.Element] = None
        self._mainloop: Optional[GLib.MainLoop] = None
        self._mainloop_thread: Optional[threading.Thread] = None

        # RTCP tracking (for absolute time mapping)
        self._rtcp_tracker = RTCPTracker()

        # RTP timestamp queue for frame correlation
        # Stores (rtp_timestamp, sequence_number, is_marker) tuples
        self._timestamp_queue: deque[Tuple[int, int, bool]] = deque(maxlen=100)
        self._timestamp_lock = threading.Lock()

        # Frame accumulator for RTP packets
        self._current_frame_rtp_ts: Optional[int] = None
        self._current_frame_seq: int = 0

        # Stream info (populated after open)
        self._width: int = 0
        self._height: int = 0
        self._fps: float = 30.0

        # State
        self._is_open = False
        self._eof = False
        self._error: Optional[str] = None
        self._frames_demuxed = 0
        self._session_id: str = ""
        self._session_start_ns: int = 0
        self._first_rtp_ts: Optional[int] = None

        # Depay output mapping for accurate RTP correlation
        # These are set when the FIRST buffer comes out of the depayloader
        self._depay_first_pts: Optional[int] = None  # PTS of first depay output
        self._depay_first_rtp: Optional[int] = None  # RTP of first depay output

    def _create_pipeline(self) -> bool:
        """Create GStreamer pipeline for H264 extraction."""
        try:
            if self._is_rtsp:
                return self._create_rtsp_pipeline()
            else:
                return self._create_file_pipeline()
        except Exception as e:
            logger.error(f"Error creating pipeline: {e}")
            return False

    def _create_rtsp_pipeline(self) -> bool:
        """Create RTSP pipeline with RTP timestamp capture.

        Pipeline outputs H264 NAL units (not decoded frames):
        rtspsrc -> [PAD PROBE] -> rtph264depay -> h264parse -> appsink
        """
        protocol = "tcp" if self.use_tcp else "udp"
        pipeline_str = f"""
            rtspsrc location={self.source} latency=0 buffer-mode=0
                protocols={protocol} do-rtcp=true name=source !
            rtph264depay name=depay !
            h264parse config-interval=-1 !
            video/x-h264,stream-format=byte-stream,alignment=au !
            appsink name=sink emit-signals=false sync=false max-buffers=4 drop=true
        """

        self._pipeline = Gst.parse_launch(pipeline_str)
        if not self._pipeline:
            return False

        # Get element references
        self._rtspsrc = self._pipeline.get_by_name("source")
        self._depay = self._pipeline.get_by_name("depay")
        self._appsink = self._pipeline.get_by_name("sink")

        if not all([self._rtspsrc, self._depay, self._appsink]):
            return False

        # Setup RTP timestamp probe BEFORE depayloading (captures raw RTP header)
        self._setup_rtp_probe()

        # Setup RTCP tracking
        self._rtspsrc.connect("new-manager", self._on_new_manager)

        return True

    def _create_file_pipeline(self) -> bool:
        """Create file pipeline for local video H264 extraction.

        Pipeline outputs H264 NAL units:
        filesrc -> demuxer -> h264parse -> appsink
        """
        source_lower = self.source.lower()
        if source_lower.endswith(('.mp4', '.mov')):
            demux = "qtdemux"
        elif source_lower.endswith('.mkv'):
            demux = "matroskademux"
        elif source_lower.endswith('.avi'):
            demux = "avidemux"
        elif source_lower.endswith('.ts'):
            demux = "tsdemux"
        else:
            demux = "qtdemux"

        pipeline_str = f"""
            filesrc location={self.source} name=source !
            {demux} name=demux !
            h264parse config-interval=-1 !
            video/x-h264,stream-format=byte-stream,alignment=au !
            appsink name=sink emit-signals=false sync=false max-buffers=4 drop=false
        """

        self._pipeline = Gst.parse_launch(pipeline_str)
        if not self._pipeline:
            return False

        self._appsink = self._pipeline.get_by_name("sink")

        if self.loop:
            bus = self._pipeline.get_bus()
            bus.add_signal_watch()
            bus.connect("message::eos", self._on_file_eos)

        return True

    def _on_file_eos(self, bus, message):
        """Handle EOS for file looping."""
        if self.loop and self._pipeline:
            self._pipeline.seek_simple(
                Gst.Format.TIME,
                Gst.SeekFlags.FLUSH | Gst.SeekFlags.KEY_UNIT,
                0
            )
            return True
        return False

    def _setup_rtp_probe(self):
        """Add probes to capture RTP timestamps.

        Two probes:
        1. Depay SINK: Capture raw RTP timestamp when complete frame arrives (marker=1)
        2. Depay SOURCE: Capture PTS->RTP mapping for the first frame

        This establishes: first_frame_pts -> first_frame_rtp
        Then all other frames use the queue for correlation.
        """
        if self._depay:
            # Probe on depay SINK to capture raw RTP timestamps
            depay_sink = self._depay.get_static_pad("sink")
            if depay_sink:
                depay_sink.add_probe(Gst.PadProbeType.BUFFER, self._on_rtp_buffer)

            # Probe on depay SOURCE to capture PTS->RTP mapping
            depay_src = self._depay.get_static_pad("src")
            if depay_src:
                depay_src.add_probe(Gst.PadProbeType.BUFFER, self._on_depay_output)

    def _on_rtp_buffer(self, pad, info):
        """Capture RAW RTP timestamp from packet header.

        This extracts the TRUE 32-bit RTP timestamp directly from the RTP header.
        The timestamp is NOT modified or normalized - it's the raw value from the source.
        """
        buffer = info.get_buffer()

        try:
            success, rtp_buffer = GstRtp.RTPBuffer.map(buffer, Gst.MapFlags.READ)
            if success:
                # Get RAW 32-bit RTP timestamp from packet header
                rtp_ts = rtp_buffer.get_timestamp()
                seq = rtp_buffer.get_seq()
                marker = rtp_buffer.get_marker()
                rtp_buffer.unmap()

                # Track first RTP timestamp
                if self._first_rtp_ts is None:
                    self._first_rtp_ts = rtp_ts
                    logger.debug(f"First ABSOLUTE RTP timestamp: {rtp_ts:,}")

                # Accumulate timestamp - use timestamp from first packet of frame
                if self._current_frame_rtp_ts is None:
                    self._current_frame_rtp_ts = rtp_ts
                    self._current_frame_seq = seq

                # When marker bit is set, frame is complete
                if marker:
                    with self._timestamp_lock:
                        self._timestamp_queue.append((
                            self._current_frame_rtp_ts,
                            self._current_frame_seq,
                            marker
                        ))
                    self._current_frame_rtp_ts = None
        except Exception as e:
            logger.debug(f"RTP probe error: {e}")

        return Gst.PadProbeReturn.OK

    def _on_depay_output(self, pad, info):
        """Capture PTS at depay output to correlate with RTP timestamp.

        The depay outputs H.264 NAL units with PTS derived from RTP:
          pts = (rtp_ts - base_rtp) * 1e9 / 90000

        On the FIRST buffer, we record the PTS and correlate it with
        the RTP timestamp from the queue. This establishes the mapping
        between depay output PTS and actual RTP timestamps.
        """
        if self._depay_first_pts is not None:
            # Already have mapping, skip
            return Gst.PadProbeReturn.OK

        buffer = info.get_buffer()
        pts = buffer.pts

        if pts != Gst.CLOCK_TIME_NONE:
            # Get corresponding RTP from queue (the first complete frame)
            with self._timestamp_lock:
                if self._timestamp_queue:
                    rtp_ts, _, _ = self._timestamp_queue[0]  # Peek, don't pop
                    self._depay_first_pts = pts
                    self._depay_first_rtp = rtp_ts
                    logger.debug(f"Depay output mapping: PTS={pts} -> RTP={rtp_ts}")

        return Gst.PadProbeReturn.OK

    def _on_new_manager(self, rtspsrc, manager):
        """Setup RTCP tracking when rtpbin manager is created."""
        try:
            manager.connect("on-sender-ssrc-active", self._on_sender_ssrc_active)
        except Exception as e:
            logger.debug(f"Failed to connect RTCP handler: {e}")

    def _on_sender_ssrc_active(self, manager, session, ssrc):
        """Called when RTCP SR is received."""
        try:
            internal_session = manager.emit("get-internal-session", session)
            if internal_session:
                stats = internal_session.get_property("stats")
                if stats:
                    source_stats = stats.get_value("source-stats")
                    if source_stats:
                        for i in range(len(source_stats)):
                            sr = self._rtcp_tracker.update_from_session_stats(
                                str(source_stats[i]), ssrc
                            )
                            if sr:
                                logger.debug(f"RTCP SR: RTP {sr.rtp_timestamp} -> Unix {sr.unix_time:.3f}")
        except Exception as e:
            logger.debug(f"RTCP handler error: {e}")

    def _on_bus_message(self, bus, message):
        """Handle GStreamer bus messages."""
        if message.type == Gst.MessageType.ERROR:
            err, debug = message.parse_error()
            self._error = err.message
            logger.error(f"Pipeline error: {err.message}")
            if self._mainloop:
                self._mainloop.quit()
        elif message.type == Gst.MessageType.EOS:
            self._eof = True
            logger.debug("Pipeline EOS")
            if self._mainloop:
                self._mainloop.quit()
        return True

    def _run_mainloop(self):
        """Run GLib mainloop in separate thread."""
        self._mainloop = GLib.MainLoop()
        try:
            self._mainloop.run()
        except Exception as e:
            logger.debug(f"Mainloop error: {e}")

    def _cleanup_pipeline(self):
        """Clean up pipeline resources for retry."""
        if self._pipeline:
            self._pipeline.set_state(Gst.State.NULL)
        if self._mainloop:
            self._mainloop.quit()
        self._pipeline = None
        self._appsink = None
        self._depay = None
        self._rtspsrc = None
        self._mainloop = None
        self._mainloop_thread = None
        self._error = None
        self._width = 0
        self._height = 0

    def _try_open(self) -> bool:
        """Attempt to open the pipeline. Returns True on success."""
        self._eof = False
        self._error = None

        if not self._create_pipeline():
            return False

        bus = self._pipeline.get_bus()
        bus.add_signal_watch()
        bus.connect("message", self._on_bus_message)

        ret = self._pipeline.set_state(Gst.State.PLAYING)
        if ret == Gst.StateChangeReturn.FAILURE:
            return False

        self._mainloop_thread = threading.Thread(target=self._run_mainloop, daemon=True)
        self._mainloop_thread.start()

        # Wait for first sample to get stream info
        timeout = 5.0
        start = time.perf_counter()
        while self._width == 0 and time.perf_counter() - start < timeout:
            if self._error:
                return False

            sample = self._appsink.emit("try-pull-sample", int(0.1 * Gst.SECOND))
            if sample:
                caps = sample.get_caps()
                if caps:
                    struct = caps.get_structure(0)
                    _, self._width = struct.get_int("width")
                    _, self._height = struct.get_int("height")
                    fps_ok, fps_num, fps_den = struct.get_fraction("framerate")
                    if fps_ok and fps_den > 0:
                        self._fps = fps_num / fps_den
                break

        return self._width > 0

    def open(self, quiet: bool = False) -> 'GstRTPDemuxer':
        """Open source and initialize demuxer pipeline with retry logic.

        For RTSP streams, automatically retries with TCP if UDP fails,
        and converts localhost to 127.0.0.1 to avoid IPv6 issues.

        Args:
            quiet: Suppress info logging

        Returns:
            self (for chaining)

        Raises:
            RuntimeError: If pipeline creation fails after all retries
        """
        self._session_start_ns = time.time_ns()
        self._session_id = str(uuid.uuid4())[:8]
        last_error = None

        # Try with current settings first
        for attempt in range(self.MAX_RETRIES + 1):
            try:
                if self._try_open():
                    break  # Success!

                # Failed - prepare for retry
                last_error = self._error or "Failed to get video dimensions (0x0 resolution)"
                self._cleanup_pipeline()

                if attempt < self.MAX_RETRIES:
                    # On first retry, try switching to TCP if using UDP
                    if attempt == 0 and not self.use_tcp and self._is_rtsp:
                        logger.warning(f"Retry {attempt + 1}: Switching to TCP transport...")
                        self.use_tcp = True
                    else:
                        logger.warning(f"Retry {attempt + 1}: Retrying connection...")

                    time.sleep(self.RETRY_DELAY)

            except Exception as e:
                last_error = str(e)
                self._cleanup_pipeline()
                if attempt < self.MAX_RETRIES:
                    time.sleep(self.RETRY_DELAY)

        if self._width == 0:
            raise RuntimeError(f"Failed to open source after {self.MAX_RETRIES + 1} attempts: {last_error}")

        self._is_open = True

        if not quiet:
            src_type = "RTSP" if self._is_rtsp else "File"
            logger.info(f"GstRTPDemuxer ({src_type}): {self._width}x{self._height} @ {self._fps:.1f} FPS")
            if self._is_rtsp and self._first_rtp_ts:
                logger.info(f"First ABSOLUTE RTP timestamp: {self._first_rtp_ts:,}")

        return self

    def demux(self) -> Iterator[Tuple[bytes, int, int, Optional[int]]]:
        """Demux H264 NAL units with timestamps.

        Yields:
            (h264_bytes, timestamp, timestamp_ns, absolute_time_ns)

            For RTSP:
            - h264_bytes: H264 NAL unit as bytes
            - timestamp: RAW 32-bit RTP timestamp (ABSOLUTE, NOT normalized)
            - timestamp_ns: RTP timestamp converted to nanoseconds (at 90kHz)
            - absolute_time_ns: Unix time if RTCP SR available, None otherwise

            For files:
            - h264_bytes: H264 NAL unit as bytes
            - timestamp: Container PTS
            - timestamp_ns: PTS in nanoseconds
            - absolute_time_ns: None
        """
        if not self._is_open:
            self.open()

        frame_counter = 0

        while not self._eof:
            if self._error:
                logger.warning(f"Demux stopping due to error: {self._error}")
                break

            sample = self._appsink.emit("try-pull-sample", int(0.1 * Gst.SECOND))
            if sample is None:
                continue

            buffer = sample.get_buffer()
            caps = sample.get_caps()

            # Update dimensions if needed
            if self._width == 0 and caps:
                struct = caps.get_structure(0)
                _, self._width = struct.get_int("width")
                _, self._height = struct.get_int("height")

            # Get timestamp based on source type
            if self._is_rtsp:
                buffer_pts = buffer.pts

                # Use depay output mapping if available (most accurate)
                # This correlates buffer PTS with actual RTP timestamps
                if (buffer_pts != Gst.CLOCK_TIME_NONE and
                    self._depay_first_pts is not None and
                    self._depay_first_rtp is not None):
                    # PTS delta from first depay output -> RTP delta
                    pts_delta = buffer_pts - self._depay_first_pts
                    rtp_delta = int(pts_delta * RTP_CLOCK_RATE // 1_000_000_000)
                    rtp_ts = self._depay_first_rtp + rtp_delta
                else:
                    # Fallback to queue
                    rtp_ts = 0
                    with self._timestamp_lock:
                        if self._timestamp_queue:
                            rtp_ts, _, _ = self._timestamp_queue.popleft()

                timestamp = rtp_ts
                timestamp_ns = rtp_ts * 1_000_000_000 // RTP_CLOCK_RATE
                absolute_time_ns = self._rtcp_tracker.rtp_to_unix_ns(rtp_ts)
            else:
                # File: use buffer PTS
                pts = buffer.pts if buffer.pts != Gst.CLOCK_TIME_NONE else frame_counter * 1_000_000_000 // 30
                timestamp = pts
                timestamp_ns = pts
                absolute_time_ns = None

            # Map buffer and extract H264 bytes
            success, map_info = buffer.map(Gst.MapFlags.READ)
            if not success:
                continue

            try:
                # Extract H264 NAL unit as bytes
                h264_bytes = bytes(map_info.data)

                self._frames_demuxed += 1
                frame_counter += 1
                yield (h264_bytes, timestamp, timestamp_ns, absolute_time_ns)

            finally:
                buffer.unmap(map_info)

    def close(self):
        """Close demuxer and release resources."""
        if self._pipeline:
            self._pipeline.set_state(Gst.State.NULL)
        if self._mainloop:
            self._mainloop.quit()
        self._pipeline = None
        self._is_open = False
        logger.debug(f"GstRTPDemuxer closed, demuxed {self._frames_demuxed} frames")

    # Properties
    @property
    def width(self) -> int:
        """Video frame width."""
        return self._width

    @property
    def height(self) -> int:
        """Video frame height."""
        return self._height

    @property
    def fps(self) -> float:
        """Video frame rate."""
        return self._fps

    @property
    def session_id(self) -> str:
        """Unique session identifier."""
        return self._session_id

    @property
    def session_start_ns(self) -> int:
        """Session start time in nanoseconds."""
        return self._session_start_ns

    @property
    def first_rtp_timestamp(self) -> Optional[int]:
        """First ABSOLUTE RTP timestamp captured (RTSP only)."""
        return self._first_rtp_ts

    @property
    def has_rtcp_sr(self) -> bool:
        """Whether RTCP Sender Report is available for absolute time mapping."""
        return self._rtcp_tracker.has_sender_report()

    @property
    def latest_rtcp_sr(self) -> Optional[RTCPSenderReport]:
        """Most recent RTCP Sender Report."""
        return self._rtcp_tracker.latest_sr

    @property
    def frames_demuxed(self) -> int:
        """Number of frames demuxed in this session."""
        return self._frames_demuxed

    @property
    def is_rtsp(self) -> bool:
        """Whether source is RTSP stream."""
        return self._is_rtsp

    @property
    def is_eof(self) -> bool:
        """Whether end-of-stream has been reached."""
        return self._eof


if __name__ == "__main__":
    import argparse
    import sys

    logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')

    parser = argparse.ArgumentParser(description="GStreamer RTP Demuxer with ABSOLUTE Timestamps")
    parser.add_argument("source", nargs="?", help="RTSP URL or file path")
    parser.add_argument("--duration", type=float, default=5.0, help="Test duration in seconds")
    parser.add_argument("--max-frames", type=int, default=100, help="Max frames to demux (file only)")
    args = parser.parse_args()

    if not args.source:
        print("Usage: python gstreamer_rtp_demuxer.py <rtsp://url | /path/to/video.mp4>")
        print("\nExample (RTSP with ABSOLUTE RTP timestamps):")
        print("  python gstreamer_rtp_demuxer.py rtsp://camera/stream --duration 10")
        print("\nExample (file):")
        print("  python gstreamer_rtp_demuxer.py /path/to/video.mp4 --max-frames 50")
        sys.exit(1)

    demuxer = GstRTPDemuxer(args.source)
    demuxer.open()

    is_rtsp = demuxer.is_rtsp
    print(f"\n{'RTSP' if is_rtsp else 'File'}: {args.source}")
    print(f"Resolution: {demuxer.width}x{demuxer.height} @ {demuxer.fps:.1f} FPS")
    if is_rtsp and demuxer.first_rtp_timestamp:
        print(f"First ABSOLUTE RTP timestamp: {demuxer.first_rtp_timestamp:,}")
    print()

    if is_rtsp:
        print(f"{'#':>4} {'ABSOLUTE RTP':>18} {'RTP (ms)':>12} {'H264 size':>12}")
    else:
        print(f"{'#':>4} {'PTS':>15} {'PTS (ms)':>12} {'H264 size':>12}")
    print("-" * 60)

    frames = 0
    start = time.perf_counter()

    for h264_bytes, ts, ts_ns, absolute_ns in demuxer.demux():
        frames += 1

        if frames <= 10:
            ts_ms = ts_ns / 1_000_000
            size_kb = len(h264_bytes) / 1024
            if is_rtsp:
                print(f"{frames:4} {ts:18,} {ts_ms:12.1f} {size_kb:10.1f} KB")
            else:
                print(f"{frames:4} {ts:15} {ts_ms:12.1f} {size_kb:10.1f} KB")

        # Check stop condition
        if is_rtsp:
            if time.perf_counter() - start >= args.duration:
                break
        else:
            if frames >= args.max_frames:
                break

    elapsed = time.perf_counter() - start
    fps = frames / elapsed if elapsed > 0 else 0

    print("-" * 60)
    print(f"Demuxed {frames} frames in {elapsed:.2f}s = {fps:.1f} FPS")
    if is_rtsp:
        print(f"RTCP Sender Report available: {demuxer.has_rtcp_sr}")

    demuxer.close()
