"""
MultiSensor - Multi-process sensor management for parallel processing.

This module provides the MultiSensor class for running multiple Sensor
instances in parallel processes, enabling true parallel processing with
independent CUDA contexts.
"""

import os


def _sensor_worker(sensor_id, config, data_queue, stop_event, use_async=False):
    """
    Worker function that runs in a separate PROCESS for each sensor.
    Each process has its own Python interpreter and CUDA context.
    
    Args:
        use_async: If True, use 3-stage async pipeline within the process
    """
    os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
    
    from .sensor import Sensor
    
    print(f"[Sensor {sensor_id}] Process started (PID: {os.getpid()}), async={use_async}")
    
    sensor = Sensor(
        vid_src=config.get('vid_src'),
        isstitch=config.get('isstitch', True),
        cuda=config.get('cuda', True),
        verbose=config.get('verbose', False),
        config_name=config.get('config_name', 'prototype'),
        cal_path=config.get('cal_path', None)
    )
    
    # Enable async pipeline if requested
    if use_async:
        sensor.enable_async(queue_size=2)
    
    frame_count = 0
    
    def queue_put(data):
        try:
            while not data_queue.empty():
                try: data_queue.get_nowait()
                except: break
            data_queue.put_nowait(data)
        except: pass
    
    while not stop_event.is_set():
        try:
            if use_async:
                # Use async get
                result = sensor.get_async(timeout=0.05)
                if result is None:
                    continue
                data = {
                    'sensor_id': sensor_id,
                    'img': result['img'].copy() if result['img'] is not None else None,
                    'flow': result['flow'].copy() if result['flow'] is not None else None,
                    'normal': result['normal'].copy() if result['normal'] is not None else None,
                    'shear': result['shear'].copy() if result['shear'] is not None else None,
                    'depth': result['depth'].copy() if result['depth'] is not None else None,
                    'fps': result['fps'],
                    'fnormal': result['fnormal'],
                    'fshearx': result['fshearx'],
                    'fsheary': result['fsheary'],
                    'ftor': result['ftor'],
                    'frame_count': frame_count,
                }
            else:
                # Sync mode
                sensor.get_img()
                sensor.compute_deformation()
                data = {
                    'sensor_id': sensor_id,
                    'img': sensor.img.copy() if sensor.img is not None else None,
                    'flow': sensor.flow.copy() if sensor.flow is not None else None,
                    'normal': sensor.vnormal.copy() if sensor.vnormal is not None else None,
                    'shear': sensor.shear.copy() if sensor.shear is not None else None,
                    'depth': sensor.depth.copy() if sensor.depth is not None else None,
                    'fps': sensor.fps,
                    'fnormal': sensor.fnormal,
                    'fshearx': sensor.fshearx,
                    'fsheary': sensor.fsheary,
                    'ftor': sensor.ftor,
                    'frame_count': frame_count,
                }
            
            queue_put(data)
            frame_count += 1
            
        except Exception as e:
            print(f"[Sensor {sensor_id}] Error: {e}")
            break
    
    sensor.disconnect()
    print(f"[Sensor {sensor_id}] Process stopped")


class MultiSensor:
    """
    Manages multiple Sensor instances running in parallel processes.
    
    Each sensor runs in its own process with independent CUDA context,
    enabling true parallel processing without GIL limitations.
    
    Example:
        # Simple usage with context manager
        configs = [
            {'vid_src': './s1.mp4', 'isstitch': True},
            {'vid_src': './s2.mp4', 'isstitch': True},
        ]
        
        with MultiSensor(configs) as multi:
            for all_data in multi:
                for data in all_data:
                    if data:
                        print(f"Sensor {data['sensor_id']}: FPS={data['fps']}")
                        cv2.imshow(f"Sensor {data['sensor_id']}", data['img'])
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    break
        
        # Or manual control
        multi = MultiSensor(configs)
        multi.start()
        
        data = multi.get(0)  # Get from sensor 0
        all_data = multi.get_all()  # Get from all sensors
        
        multi.stop()
    """
    
    def __init__(self, sensor_configs):
        """
        Initialize MultiSensor manager.
        
        Args:
            sensor_configs: List of dicts, each containing Sensor constructor args:
                - vid_src: Video source (required) - int for camera ID, str for file path
                - isstitch: Whether to stitch (default: True)
                - cuda: Use CUDA acceleration (default: True)
                - verbose: Print logs (default: False)
                - config_name: Config name (default: 'prototype')
                - cal_path: Calibration file path (default: None)
        
        Example:
            configs = [
                {'vid_src': 0, 'isstitch': True, 'cuda': True},
                {'vid_src': 1, 'isstitch': True, 'cuda': True},
            ]
            multi = MultiSensor(configs)
        """
        import multiprocessing as mp
        
        self.sensor_configs = sensor_configs
        self.num_sensors = len(sensor_configs)
        self._mp = mp
        self._use_async = False
        
        self.processes = []
        self.data_queues = []
        self.stop_events = []
        self._started = False
        
        # Create queues and events for each sensor
        for i in range(self.num_sensors):
            self.data_queues.append(mp.Queue(maxsize=2))
            self.stop_events.append(mp.Event())
    
    def start(self, use_async=False):
        """
        Start all sensor worker processes.
        
        Args:
            use_async: If True, each sensor process uses 3-stage async pipeline
                      for maximum throughput (Capture → Flow → HHD in parallel)
        """
        if self._started:
            print("MultiSensor already started")
            return
        
        self._use_async = use_async
        
        for i, config in enumerate(self.sensor_configs):
            process = self._mp.Process(
                target=_sensor_worker,
                args=(i, config, self.data_queues[i], self.stop_events[i], use_async),
                daemon=True
            )
            process.start()
            self.processes.append(process)
            print(f"Started sensor {i} (PID: {process.pid})")
        
        mode = "async (3-stage pipeline)" if use_async else "sync"
        self._started = True
        print(f"All {len(self.processes)} sensors started in {mode} mode")
    
    def get(self, sensor_idx, timeout=0.01):
        """
        Get the latest data from a specific sensor.
        
        Args:
            sensor_idx: Index of the sensor (0-based)
            timeout: How long to wait for data (seconds)
        
        Returns:
            dict: Sensor data or None if no data available
                - sensor_id: int
                - img: np.ndarray (H, W, 3)
                - flow: np.ndarray (H, W, 2)
                - normal: np.ndarray
                - shear: np.ndarray
                - depth: np.ndarray
                - fps: float
                - fnormal, fshearx, fsheary, ftor: float
                - frame_count: int
        """
        try:
            return self.data_queues[sensor_idx].get(timeout=timeout)
        except:
            return None
    
    def get_all(self, timeout=0.01):
        """
        Get the latest data from all sensors.
        
        Returns:
            list: List of sensor data dicts (None for sensors with no new data)
        """
        return [self.get(i, timeout) for i in range(self.num_sensors)]
    
    def __iter__(self):
        """
        Iterate over sensor data. Yields all sensor data on each iteration.
        
        Example:
            for all_data in multi:
                for data in all_data:
                    if data:
                        process(data)
        """
        return self
    
    def __next__(self):
        """Get next batch of sensor data."""
        if not self._started:
            raise StopIteration("MultiSensor not started")
        
        # Check if any process is still alive
        if not any(p.is_alive() for p in self.processes):
            raise StopIteration("All sensor processes have stopped")
        
        return self.get_all()
    
    def __len__(self):
        """Return number of sensors."""
        return self.num_sensors
    
    def __getitem__(self, idx):
        """Get latest data from sensor by index."""
        return self.get(idx)
    
    def stop(self):
        """Stop all sensor processes gracefully."""
        if not self._started:
            return
        
        print("Stopping all sensors...")
        
        # Signal all processes to stop
        for event in self.stop_events:
            event.set()
        
        # Wait for processes to finish
        for i, process in enumerate(self.processes):
            process.join(timeout=3.0)
            if process.is_alive():
                print(f"Force terminating sensor {i}")
                process.terminate()
        
        self.processes.clear()
        self._started = False
        print("All sensors stopped.")
    
    def __enter__(self):
        """Context manager entry - starts sensors (sync mode by default)."""
        self.start(use_async=False)
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit - stops sensors."""
        self.stop()
        return False
    
    def __del__(self):
        """Cleanup on deletion."""
        self.stop()
    
    @classmethod
    def async_context(cls, sensor_configs):
        """
        Create a MultiSensor that uses async mode when used as context manager.
        
        Example:
            with MultiSensor.async_context(configs) as multi:
                for all_data in multi:
                    ...
        """
        return _AsyncMultiSensorContext(sensor_configs)


class _AsyncMultiSensorContext(MultiSensor):
    """Helper class for async context manager."""
    def __enter__(self):
        self.start(use_async=True)
        return self