import os
import gc
import cv2
cv2.setNumThreads(0)
import time
import yaml
import uuid
import logging
import platform
import threading
import numpy as np

from .utils import is_jetson

# Obtain the "nxva.streaming" logger, which will naturally propagate to the "nxva" logger
_logger = logging.getLogger(__name__)
_logger.addHandler(logging.NullHandler())  # Avoid the "No handler" warning


class MultiStreaming(object):
    """
    A class to handle multi-camera streaming with automatic reconnection, frame resizing and gstreamer support.

    Attributes:
        reconnect_interval (int): Interval in seconds to attempt reconnection to a camera.
        logger (Logger): Logger instance for logging messages.
        cfg (dict): Configuration settings for the cameras which read from a YAML file or dictionary.
        w (int): Width of the frame.
        h (int): Height of the frame.
        gst_enable (bool): Whether to use GStreamer for streaming.
        codec (str): Codec to use for streaming, default is 'h264'.
        cameras (dict): Dictionary containing camera configurations, where each key is a unique item_id 
                        and the value is a dictionary with the following keys:
            - item_id (str): Unique identifier for the camera.  Generated by uuid if not provided.
            - rtsp (list): List of dictionaries containing RTSP stream information:
                - ip (str): IP address of the camera.
                - user (str, optional): Username for the camera.
                - pw (str, optional): Password for the camera.
            - codec (str): Codec to use to generate gstreamer pipeline, default is 'h264'.
            - status (bool): Whether the camera is currently active.
            - rtsp_strings (list): List of RTSP string candidates
            - rtsp_reconnect (str): RTSP string to use for reconnection.
            - cap (cv2.VideoCapture): OpenCV VideoCapture object for the camera.
            - cap_info (dict): Dictionary containing camera properties like width, height, FPS, and codec.
        num (int): Number of cameras configured.
        running (bool): Flag to indicate if the streaming is running.
        mutex (threading.Lock): Mutex for thread-safe access to camera frames.
        threads (list): List of threads for each camera stream.
        frames (list): List of frames from each camera.
        id_to_idx (dict): Mapping from camera item_id to index in the frames list.
    """
    def __init__(
        self, 
        config=None, 
        reconnect_interval=30,
        logger=None,
        verbose=True
    ):
        """
        Initialize the MultiStreaming object with the given configuration settings.
        
        Args:
            config (str or dict): Configuration settings for the cameras, 
                which can be provided as a path to a YAML file or directly as a dictionary. 
                The configuration should include the following sections:
                
                'parameter': Dictionary containing:
                    - 'WIDTH' (int): Width of the frame.
                    - 'HEIGHT' (int): Height of the frame.
                    - 'GSTREAMER' (bool): Whether to use GStreamer for streaming.
                    - 'CODEC' (str, optional): Global codec to use for gstreamer
                
                'cameras': List of dictionaries, each containing:
                    - 'item_id' (str, optional): Unique identifier for the camera. If not provided, a UUID will be generated.
                    - 'ip' (str): IP address of the camera.
                    - 'user' (str, optional): Username for the camera.
                    - 'pw' (str, optional): Password for the camera.
                
                Example configuration:
                ```
                parameter:
                    WIDTH: 1920
                    HEIGHT: 1080
                    GSTREAMER: True
                    CODEC: h264
                cameras:
                    - ip: '192.168.1.100'
                      user: 'admin'
                      pw: 'password'
                ```
                or 
                ```
                cameras:
                    - item_id: 'CAMERA-1'
                      rtsp:
                        - ip: '192.168.1.100'
                          user: 'admin'
                          pw: 'password'
                      codec: 'h264'
                ```
            reconnect_interval (int): Interval in seconds to attempt reconnection to a camera.
            verbose (bool): Enable verbose output to help with debugging.
        """
        self.reconnect_interval = reconnect_interval
        if not verbose:
            _logger.setLevel(logging.ERROR)
        self.logger = logger or _logger
        
        if config:
            self._load_and_validate_config(config)
            self._generate_rtsp_strings()
        else:
            raise ValueError('Configuration must be provided, either as a dictionary or a YAML file path.')
        
        self.running = True
        self.mutex = threading.Lock()
        self.threads = []  # store streaming threads for each camera
        self.frames = [np.zeros((self.h, self.w, 3), dtype=np.uint8) for i in range(self.num)]
        self.id_to_idx = {cam['item_id']: i for i, cam in enumerate(self.cameras.values())}

    def _load_and_validate_config(self, config):
        if isinstance(config, str):
            assert config.endswith(('.yaml', '.yml'))
            with open(config, 'r') as f:
                self.cfg = yaml.safe_load(f)
        elif isinstance(config, dict):
            self.cfg = config
        else:
            raise TypeError('Config must be provided as a yaml file or dictionary.')
        
        # Ensure 'cameras' section exists and is a list
        if 'cameras' not in self.cfg or not isinstance(self.cfg['cameras'], list):
            raise ValueError("The 'cameras' section must be a list and cannot be empty.")
        
        param = self.cfg.get('parameter', {}) or {}
        self.w = int(param.get('WIDTH', 1920))
        self.h = int(param.get('HEIGHT', 1080))
        self.candidate = bool(param.get('CANDIDATE', False))
        self.gst_enable = bool(param.get('GSTREAMER', False))
        self.backend = param.get('BACKEND', None)  # 'opencv', 'gst' or 'ffmpeg'
        # Backward compatibility: if backend is not set, use 'gst' if GStreamer is enabled
        if self.backend is None:
            self.backend = 'gst' if self.gst_enable else 'opencv'
        self.codec = param.get('CODEC', 'h264')  # global setting for codec
        # FFmpeg exclusive parameters (optional)
        ffmpeg_param = param.get('FFMPEG', {}) or {}
        self.ffmpeg_fps = int(ffmpeg_param.get('FPS', 5))
        self.ffmpeg_hwaccel = bool(ffmpeg_param.get('HWACCEL', True))
        self.ffmpeg_fallback = bool(ffmpeg_param.get('FALLBACK', False))
        self.logger.info(f"Streaming global settings: width={self.w}, height={self.h}, backend={self.backend}, codec={self.codec}, "
                         f"ffmpeg_fps={self.ffmpeg_fps}, ffmpeg_hwaccel={self.ffmpeg_hwaccel}, ffmpeg_fallback={self.ffmpeg_fallback}")

        self.cameras = {}
        for cam in self.cfg['cameras']:
            item_id = cam.get('item_id', str(uuid.uuid4()))
            cam['item_id'] = item_id  # confirm item_id is set
            if item_id in self.cameras:
                raise ValueError(f"Duplicate item_id found: {item_id}. Each camera must have a unique item_id.")
            
            # Ensure global codec in camera configuration
            # If codec is not specified, use the global codec
            if 'codec' not in cam:
                cam['codec'] = self.codec
            cam['codec'] = cam['codec'].lower()
            if cam['codec'] not in ['h264', 'h265']:
                raise ValueError(f"Unsupported codec '{cam['codec']}' for camera {item_id}. Supported codecs are: h264, h265.")
            
            if 'rtsp' in cam:
                if not isinstance(cam['rtsp'], list):
                    raise ValueError(f"Camera {item_id} 'rtsp' must be a list of dictionary with 'ip', 'user', and 'pw'.")
                else:
                    # Validate each RTSP entry
                    for rtsp in cam['rtsp']:
                        if not isinstance(rtsp, dict):
                            raise ValueError(f"Each RTSP entry for camera {item_id} must be a dictionary with 'ip', 'user', and 'pw'.")
                        if 'ip' not in rtsp:
                            raise ValueError(f"RTSP entry for camera {item_id} must have 'ip' defined.")
                        if 'user' not in rtsp:
                            rtsp['user'] = None
                        if 'pw' not in rtsp:
                            rtsp['pw'] = None
                        if 'codec' not in rtsp:
                            rtsp['codec'] = cam['codec']
                        rtsp['codec'] = rtsp['codec'].lower()
                        if rtsp['codec'] not in ['h264', 'h265']:
                            raise ValueError(f"Unsupported codec '{rtsp['codec']}' for camera {item_id}. Supported codecs are: h264, h265.")
            else:
                ip = cam.get('ip', None)
                if not ip:
                    raise ValueError(f"Camera {item_id} must have 'ip' or 'rtsp' defined.")
                cam['rtsp'] = [{'ip': ip, 'user': cam.get('user', None), 'pw': cam.get('pw', None), 'codec': cam['codec']}]
            
            cam['status'] = False
            cam['rtsp_strings'] = []
            cam['rtsp_reconnect'] = None
            cam['cap'] = cv2.VideoCapture()
            cam['cap_info'] = {}
            self.cameras[item_id] = cam

        self.num = len(self.cameras)

    def _generate_rtsp_strings(self):
        cpu_arch = platform.machine()
        # Collect RTSP strings
        for item_id, cam in self.cameras.items():
            for rtsp in cam['rtsp']:
                ip = rtsp['ip']
                user = rtsp['user']
                pw = rtsp['pw']
                codec = rtsp['codec']

                if self.backend == 'gst':
                    gst_pipeline = self._gst_pipeline_template(cpu_arch, need_user_pw=bool(user and pw))
                    if gst_pipeline:
                        rtsp_string = gst_pipeline.format(ip=ip, user=user, pw=pw, codec=codec)
                        # Gstreamer pipeline put in front of the default rtsp
                        if len(cam['rtsp_strings']) > 1:
                            index = len(cam['rtsp_strings']) // 2
                            cam['rtsp_strings'].insert(index, rtsp_string)
                        else:
                            cam['rtsp_strings'].append(rtsp_string)
                
                # if Gstreamer is not used, use the default rtsp
                # if Gstreamer is used, this will be backup
                if user and pw:
                    cam['rtsp_strings'].append(f'rtsp://{user}:{pw}@{ip}')
                else:
                    cam['rtsp_strings'].append(f'rtsp://{ip}')

        log_string = (
            f"RTSP strings for {len(self.cameras)} cameras:\n"
            + "\n".join(f"{item_id}: {cam['rtsp_strings']}" for cam in self.cameras.values())
        )
        self.logger.debug(log_string)

    def _gst_pipeline_template(self, cpu_arch, need_user_pw=False):
        rtspsrc_base  = 'rtspsrc location=rtsp://{ip}'
        if need_user_pw:
            rtspsrc_base += ' user-id={user} user-pw={pw}'

        # https://forums.developer.nvidia.com/t/deepstream-losing-relevance-because-of-outdated-dependecies/203426/7?utm_source=chatgpt.com 
        rtspsrc_params = [
            'short-header=true',
            'is-live=true',
            'protocols=tcp',
            'user-agent="LibVLC/3.0.18"',
            'do-rtsp-keep-alive=true',
            'timeout=5000000',  # 5s
            'tcp-timeout=5000000'
            # 'latency=200',
            # 'drop-on-latency=true',
            # 'timeout=10',
            # 'tcp-timeout=5',
            # 'retry=3',
            # 'do-retransmission=false'
        ]
        rtspsrc = rtspsrc_base + ' ' + ' '.join(rtspsrc_params)
        
        if cpu_arch == 'x86_64':
            decode_chain = [
                f'queue max-size-buffers=0 max-size-bytes={1024*1024} leaky=downstream',  # 1MB buffer to keep GOP not be truncated
                'rtp{codec}depay', 
                '{codec}parse', 
                'avdec_{codec}', 
                'videoconvert n-threads=1', 
                'appsink drop=true sync=false'
            ]
            # decode_chain = [
            #     f'queue max-size-buffers=0 max-size-bytes={1024*1024} leaky=downstream',
            #     'rtp{codec}depay', 
            #     '{codec}parse', 
            #     'nv{codec}dec',
            #     'nvvideoconvert'  # it is important to use nvvideoconvert after nv{codec}dec for performance
            #     'queue max-size-buffers=1 leaky=downstream',
            #     'videoconvert n-threads=1',  # it is important to add convert after nvvideoconvert for performance
            #     'video/x-raw, format=RGBA',
            #     'appsink caps="video/x-raw,format=BGR drop=true sync=false'
            # ]
        elif cpu_arch == 'aarch64':
            if is_jetson():
                decode_chain = [
                    f'queue max-size-buffers=0 max-size-bytes={1024*1024} leaky=downstream',
                    'rtp{codec}depay', 
                    '{codec}parse config-interval=1', 
                    'nvv4l2decoder',  #  drop-frame has some problem, disable-dpb=true need the stream with no B-frames
                    'nvvidconv', 
                    'queue max-size-buffers=1 leaky=downstream',
                    'video/x-raw,format=RGBA',          # explicitly set RGBA format for system memory
                    'videoconvert n-threads=1',         
                    'video/x-raw,format=BGR',           # cap filter, ensure the output is in BGR format
                    'appsink drop=true sync=false'
                ]
                # decode_chain = [
                #     'rtp{codec}depay', '{codec}parse', 'nvv4l2decoder', 
                #     'nvvidconv', 'appsink drop=true sync=false'
                # ]
            else:
                decode_chain = [
                'rtp{codec}depay', 
                '{codec}parse', 
                'avdec_{codec}', 
                'videoconvert n-threads=1', 
                'appsink drop=true sync=false'
            ]
        else:
            decode_chain = []
            self.logger.warning('Not allow to use Gstreamer under the current platform! (Only support x86_64 and Jetson)')

        if not decode_chain:
            return ''
        return ' ! '.join([rtspsrc] + decode_chain)
    
    def __len__(self):
        return self.num
    
    def run(self):
        """Start the camera streaming threads."""
        # for i in range(self.num):
        for cam in self.cameras.values():
            t = threading.Thread(
                target=self._run_one,
                args=(cam, ),
                daemon=True
                )
            t.start()
            self.threads.append(t)

    def get_frames(self):
        """
        Retrieve the latest frames from all cameras.
        """
        with self.mutex:
            frames = self.frames.copy()
            return frames
        
    def get_status(self):
        """
        Retrieve the status of all cameras.
        Returns a list of booleans indicating whether each camera is active.
        """
        with self.mutex:
            status_list = [cam['status'] for cam in self.cameras.values()]
            return status_list
        
    def stop(self):
        """Stop all camera threads."""
        self.running = False

    def release(self):
        """Stop all camera threads and release resources."""
        self.running = False
        time.sleep(0.1)
        for cam in self.cameras.values():
            if 'cap' in cam and cam['cap'].isOpened():
                cam['cap'].release()

        # Wait for threads to finish
        for thread in self.threads:
            thread.join()
        self.threads.clear()
        self.logger.info("All camera threads stopped and resources released.")
            
    def _run_one(self, cam):
        """Handle the streaming of one individual camera."""
        self._init_camera(cam)
        cap = cam['cap']
        resize = (cam['cap_info']['w'], cam['cap_info']['h']) != (self.w, self.h)
        reconnect_time = time.time() - self.reconnect_interval
        
        try:
            fail = 0
            while self.running:
                # successfully read frame
                if cap.isOpened():
                    ret, frame = cap.read()
                    if ret:
                        if not cam['status']:
                            self.logger.info(f"Camera {cam['item_id']} reconnects successfully with RTSP \"{cam['rtsp_reconnect']}\"")
                        if resize:
                            frame = cv2.resize(frame, (self.w, self.h))
                        with self.mutex:
                            self.frames[self.id_to_idx[cam['item_id']]] = frame
                            cam['status'] = True
                        time.sleep(0.02)
                        continue
                    else:
                        fail += 1
                else:
                    fail += 1
                
                # camera is not open or failed to read frame
                with self.mutex:
                    self.frames[self.id_to_idx[cam['item_id']]] = np.zeros((self.h, self.w, 3), dtype=np.uint8)
                    cam['status'] = False
                    
                if time.time() - reconnect_time > self.reconnect_interval and fail > 10:
                    reconnect_time = time.time()
                    self.logger.warning(f"Camera {cam['item_id']} is not open, try reconnecting with RTSP \"{cam['rtsp_reconnect']}\"")
                    if cap.isOpened():
                        cap.release()
                    del cap
                    gc.collect()
                    time.sleep(0.2)  
                    if cam['rtsp_reconnect']:
                        cap = self._open_camera(cam['rtsp_reconnect'])
                        if not cap.isOpened():
                            if self.candidate:
                                cam['rtsp_reconnect'] = None
                    else:
                        for rtsp in cam['rtsp_strings']:
                            cap = self._open_camera(rtsp)
                            if cap.isOpened():
                                cam['rtsp_reconnect'] = rtsp
                                break
                    cam['cap'] = cap
                    cam['cap_info'] = self._get_info(cam['cap'])
                    resize = (cam['cap_info']['w'], cam['cap_info']['h']) != (self.w, self.h)
                    fail = 0
                
                time.sleep(0.02)
        finally:
            if cap.isOpened():
                cap.release()
    
    def _init_camera(self, cam):
        """Initialize all cameras based on the configuration settings."""
        for rtsp in cam['rtsp_strings']:
            cam['cap'] = self._open_camera(rtsp)
            cam['cap_info'] = self._get_info(cam['cap'])
            
            if cam['cap'].isOpened():
                cam['status'] = True
                cam['rtsp_reconnect'] = rtsp
                self.logger.info(f"Camera {cam['item_id']}: opened successfully with")
                self.logger.info(f"    RTSP: {rtsp}")
                self.logger.info(f"    Info: {cam['cap_info']}")
                break
            else:
                cam['status'] = False
                cam['rtsp_reconnect'] = None if self.candidate else rtsp
                self.logger.warning(f"Camera {cam['item_id']}: failed to open initially with RTSP \"{rtsp}\"")
            
            if not self.candidate:
                # If not candidate mode, break after the first RTSP string
                break

    def _open_camera(self, rtsp):
        """Build capture object based on the backend and RTSP string."""
        if self.backend == 'ffmpeg':
            # Only accept 'rtsp://' strings.
            # If a GStreamer pipeline is passed, it will be ignored and replaced by the upper layer.
            if not rtsp.startswith('rtsp://'):
                cap = type('DummyCap', (), {'isOpened': lambda s: False})()
                return cap
            from .ffmpeg_backend import FFmpegCapture
            cap = FFmpegCapture(self.w, self.h, fps=self.ffmpeg_fps, hwaccel=self.ffmpeg_hwaccel, fallback=self.ffmpeg_fallback)
            cap.open(rtsp)
            return cap
        else:
            cap = self._open_cv2_cap(rtsp, gst=(self.backend == 'gst'))
            return cap
    
    @staticmethod
    def _open_cv2_cap(rtsp, gst=False):
        """Open a camera using OpenCV VideoCapture with optional GStreamer support."""
        cap = cv2.VideoCapture()
        if gst and rtsp.startswith('rtspsrc'):
            t = threading.Thread(
                target=lambda c, u: c.open(u, cv2.CAP_GSTREAMER), 
                args=(cap, rtsp), 
                daemon=True
            )
        else:
            t = threading.Thread(
                target=lambda c, u: c.open(u), 
                args=(cap, rtsp), 
                daemon=True
            )
        t.start()
        t.join(15)
        return cap
    
    def _get_info(self, cap):
        """Return camera properties such as width, height, FPS, and codec."""
        try:
            if cap.isOpened():
                w = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
                h = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
                fps = cap.get(cv2.CAP_PROP_FPS)
                fourcc = int(cap.get(cv2.CAP_PROP_FOURCC))
                codec_format = ''.join([chr((fourcc >> 8 * i) & 0xFF) for i in range(4)])
                return {
                    'w': int(w),
                    'h': int(h),
                    'fps': int(fps),
                    'codec': codec_format
                }
        except Exception as e:
            self.logger.exception(f"Error getting camera properties: {e}")

        return {
            'w': 0,
            'h': 0,
            'fps': 0,
            'codec': ''
        }