"""
Single async event loop consumer manager for 1000 cameras.

Architecture:
- Single async event loop handles all 1000 camera streams
- Async Redis/Kafka operations (non-blocking)
- Direct frame bytes extraction and forwarding
- No codec-specific processing (simplified)
- No frame caching (moved to producer)
- Backpressure handling (drop frames if queue full)
"""

import asyncio
import base64
import logging
import time
import uuid
from typing import Dict, Any, Optional
from matrice_inference.server.stream.utils import CameraConfig, StreamMessage
from matrice_inference.server.stream.worker_metrics import WorkerMetrics


class AsyncConsumerManager:
    """
    Manages 1000 camera streams with single async event loop.

    Key Features:
    - Single event loop for all cameras (not 1000 threads)
    - Async stream reads (non-blocking)
    - Direct bytes extraction (no codec processing)
    - Backpressure handling (drop frames if queue full)
    - Dynamic camera add/remove support
    """

    def __init__(
        self,
        camera_configs: Dict[str, CameraConfig],
        stream_config: Dict[str, Any],
        input_topic: str,
        pipeline: Any,
        message_timeout: float = 2.0,
    ):
        self.camera_configs = camera_configs
        self.stream_config = stream_config
        self.input_topic = input_topic
        self.pipeline = pipeline
        self.message_timeout = message_timeout
        self.running = False

        # Camera streams (one per camera)
        self.streams: Dict[str, Any] = {}

        # Async tasks (one per camera)
        self.consumer_tasks: Dict[str, asyncio.Task] = {}

        # Initialize metrics (shared across all cameras)
        self.metrics = WorkerMetrics.get_shared("consumer")

        self.logger = logging.getLogger(f"{__name__}.AsyncConsumerManager")

    async def start(self):
        """Start async consumers for all cameras."""
        self.running = True
        self.metrics.mark_active()

        self.logger.info(f"Starting async consumer manager for {len(self.camera_configs)} cameras")

        # Initialize streams for all cameras
        await self._initialize_streams()

        # Create async task for each camera
        for camera_id, config in self.camera_configs.items():
            await self.add_camera(camera_id, config)

        self.logger.info(f"Created {len(self.consumer_tasks)} async camera consumers")

    async def stop(self):
        """Stop all consumers."""
        self.running = False
        self.metrics.mark_inactive()

        # Cancel all consumer tasks
        for task in self.consumer_tasks.values():
            task.cancel()

        # Wait for tasks to complete
        if self.consumer_tasks:
            await asyncio.gather(*self.consumer_tasks.values(), return_exceptions=True)

        # Close all streams
        for stream in self.streams.values():
            try:
                await stream.async_close()
            except Exception:
                pass

        self.logger.info("Stopped async consumer manager")

    async def add_camera(self, camera_id: str, config: CameraConfig):
        """
        Add a new camera dynamically.

        Args:
            camera_id: Unique camera identifier
            config: Camera configuration
        """
        if camera_id in self.consumer_tasks:
            self.logger.warning(f"Camera {camera_id} already exists, skipping add")
            return

        try:
            # Initialize stream if not exists
            if camera_id not in self.streams:
                await self._initialize_camera_stream(camera_id)

            # Create async task for this camera
            task = asyncio.create_task(
                self._consume_camera(camera_id, config),
                name=f"consumer_{camera_id}"
            )
            self.consumer_tasks[camera_id] = task

            self.logger.info(f"Added camera {camera_id} to consumer manager")

        except Exception as e:
            self.logger.error(f"Failed to add camera {camera_id}: {e}")

    async def remove_camera(self, camera_id: str):
        """
        Remove a camera dynamically.

        Args:
            camera_id: Unique camera identifier
        """
        if camera_id not in self.consumer_tasks:
            self.logger.warning(f"Camera {camera_id} not found, skipping remove")
            return

        try:
            # Cancel the task
            task = self.consumer_tasks[camera_id]
            task.cancel()

            # Wait for task to complete
            try:
                await asyncio.wait_for(task, timeout=5.0)
            except (asyncio.CancelledError, asyncio.TimeoutError):
                pass

            # Remove from tasks dict
            del self.consumer_tasks[camera_id]

            # Close and remove stream
            if camera_id in self.streams:
                try:
                    await self.streams[camera_id].async_close()
                except Exception:
                    pass
                del self.streams[camera_id]

            self.logger.info(f"Removed camera {camera_id} from consumer manager")

        except Exception as e:
            self.logger.error(f"Failed to remove camera {camera_id}: {e}")

    async def _initialize_streams(self):
        """Initialize streams for all cameras."""
        from matrice_common.stream.matrice_stream import MatriceStream, StreamType

        stream_type = self._get_stream_type()
        stream_params = self._build_stream_params(stream_type)

        for camera_id, camera_config in self.camera_configs.items():
            try:
                stream = MatriceStream(stream_type, **stream_params)
                consumer_group = f"inference_consumer_{camera_id}"
                # Use camera-specific input topic (not shared input_topic)
                camera_input_topic = camera_config.input_topic
                await stream.async_setup(camera_input_topic, consumer_group)

                self.streams[camera_id] = stream

                self.logger.info(f"✓ Initialized stream for camera {camera_id} on topic {camera_input_topic}")

            except Exception as e:
                self.logger.error(f"Failed to initialize stream for camera {camera_id}: {e}")

    async def _initialize_camera_stream(self, camera_id: str):
        """Initialize stream for a single camera."""
        from matrice_common.stream.matrice_stream import MatriceStream, StreamType

        stream_type = self._get_stream_type()
        stream_params = self._build_stream_params(stream_type)

        try:
            # Get camera-specific input topic
            camera_config = self.camera_configs.get(camera_id)
            if not camera_config:
                raise ValueError(f"No config found for camera {camera_id}")

            stream = MatriceStream(stream_type, **stream_params)
            consumer_group = f"inference_consumer_{camera_id}"
            # Use camera-specific input topic (not shared input_topic)
            camera_input_topic = camera_config.input_topic
            await stream.async_setup(camera_input_topic, consumer_group)

            self.streams[camera_id] = stream

            self.logger.info(f"✓ Initialized stream for camera {camera_id} on topic {camera_input_topic}")

        except Exception as e:
            self.logger.error(f"Failed to initialize stream for camera {camera_id}: {e}")
            raise

    def _get_stream_type(self):
        """Determine stream type from configuration."""
        from matrice_common.stream.matrice_stream import StreamType
        stream_type_str = self.stream_config.get("stream_type", "kafka").lower()
        return StreamType.KAFKA if stream_type_str == "kafka" else StreamType.REDIS

    def _build_stream_params(self, stream_type) -> Dict[str, Any]:
        """Build stream parameters based on type."""
        from matrice_common.stream.matrice_stream import StreamType

        if stream_type == StreamType.KAFKA:
            return {
                "bootstrap_servers": self.stream_config.get("bootstrap_servers", "localhost:9092"),
                "sasl_username": self.stream_config.get("sasl_username", "matrice-sdk-user"),
                "sasl_password": self.stream_config.get("sasl_password", "matrice-sdk-password"),
                "sasl_mechanism": self.stream_config.get("sasl_mechanism", "SCRAM-SHA-256"),
                "security_protocol": self.stream_config.get("security_protocol", "SASL_PLAINTEXT"),
            }
        else:
            return {
                "host": self.stream_config.get("host", "localhost"),
                "port": self.stream_config.get("port", 6379),
                "password": self.stream_config.get("password"),
                "username": self.stream_config.get("username"),
                "db": self.stream_config.get("db", 0),
                "connection_timeout": self.stream_config.get("connection_timeout", 120),
            }

    async def _consume_camera(self, camera_id: str, config: CameraConfig):
        """
        Async consumer for single camera.

        This runs concurrently with all other camera consumers in the same event loop.
        Each camera has its own async while loop for continuous frame processing.
        """
        stream = self.streams.get(camera_id)
        if not stream:
            self.logger.error(f"No stream for camera {camera_id}")
            return

        self.logger.info(f"Started consumer for camera {camera_id}")

        while self.running and config.enabled:
            try:
                # Non-blocking stream read
                message_data = await stream.async_get_message(self.message_timeout)

                if not message_data:
                    continue

                # Process message - simplified to just extract bytes and forward
                await self._process_message(camera_id, config, message_data)

            except asyncio.CancelledError:
                break
            except Exception as e:
                self.logger.error(f"Error consuming camera {camera_id}: {e}")
                await asyncio.sleep(1.0)

        self.logger.info(f"Stopped consumer for camera {camera_id}")

    async def _process_message(
        self,
        camera_id: str,
        camera_config: CameraConfig,
        message_data: Dict[str, Any]
    ):
        """
        Process incoming message and enqueue for inference.

        Simplified version - just extracts bytes and forwards to inference queue.
        No codec-specific processing.
        """
        start_time = time.time()
        try:
            # Extract basic metadata
            message_key = self._extract_message_key(message_data)
            data = self._parse_message_data(message_data)

            # Reconstruct input_stream with binary content from flattened Redis fields
            input_stream = self._reconstruct_input_stream_content(data)
            extra_params = self._normalize_extra_params(data)
            frame_id = self._determine_frame_id(data, message_data, camera_id)

            # Enrich input_stream with frame_id
            self._enrich_input_stream(input_stream, frame_id)

            # Extract frame bytes (no codec processing)
            frame_bytes = self._extract_frame_bytes(input_stream)

            if not frame_bytes:
                self.logger.warning(f"No frame bytes for camera {camera_id}, frame_id {frame_id}")
                return

            # Create stream message
            stream_msg = self._create_stream_message(camera_id, message_key, data)

            # Build task data with simplified structure
            task_data = {
                "camera_id": camera_id,
                "frame_bytes": frame_bytes,  # Direct bytes for inference
                "frame_id": frame_id,
                "message": stream_msg,
                "input_stream": input_stream,
                "stream_key": camera_id,
                "extra_params": extra_params,
                "camera_config": camera_config,
            }

            # Enqueue directly to inference queue
            await self._enqueue_task(camera_id, task_data)

            # Record metrics for successfully processed message
            latency_ms = (time.time() - start_time) * 1000
            self.metrics.record_latency(latency_ms)
            self.metrics.record_throughput(count=1)

        except Exception as e:
            self.logger.error(f"Error processing message for camera {camera_id}: {e}")

    def _extract_frame_bytes(self, input_stream: Dict[str, Any]) -> Optional[bytes]:
        """
        Extract frame bytes from input_stream.

        Handles both raw bytes and base64-encoded strings.
        Also handles Redis flattened format where binary content is in 'input_stream__content'.
        """
        if not isinstance(input_stream, dict):
            return None

        content = input_stream.get("content")

        # Handle raw bytes
        if isinstance(content, bytes) and content:
            return content

        # Handle base64-encoded strings
        elif isinstance(content, str) and content:
            try:
                return base64.b64decode(content)
            except Exception as e:
                self.logger.warning(f"Failed to decode base64 content: {e}")
                return None

        return None

    def _reconstruct_input_stream_content(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Reconstruct input_stream with binary content from flattened Redis fields.

        Redis stores binary content separately as 'input_stream__content'.
        This method reconstructs it back into 'input_stream.content'.
        """
        input_stream = data.get("input_stream", {})
        if not isinstance(input_stream, dict):
            return {}

        # Check for flattened binary content field
        if "input_stream__content" in data:
            binary_content = data["input_stream__content"]
            if isinstance(binary_content, bytes):
                # Reconstruct: put binary content back into input_stream
                input_stream["content"] = binary_content

        return input_stream

    async def _enqueue_task(self, camera_id: str, task_data: Dict[str, Any]):
        """
        Enqueue task for inference with camera-based routing to specific worker queue.

        Routes frame to worker queue based on hash(camera_id) % num_workers.
        This ensures:
        - Same camera always goes to same worker queue
        - Frames are processed in FIFO order within each queue
        - No re-queuing → No race conditions → Order preserved per camera
        """
        try:
            # Get per-worker inference queues from pipeline
            inference_queues = self.pipeline.inference_queues

            if not inference_queues:
                self.logger.error("No inference queues available")
                return

            # Route to specific worker queue based on camera hash
            worker_id = hash(camera_id) % len(inference_queues)
            target_queue = inference_queues[worker_id]

            # Run blocking mp.Queue.put() in executor to avoid blocking event loop
            loop = asyncio.get_running_loop()
            await loop.run_in_executor(
                None,  # Use default executor
                target_queue.put,
                task_data,
                True,  # block=True
                1.0  # timeout (seconds)
            )

            self.logger.debug(
                f"Routed frame for camera {camera_id} to inference worker {worker_id} "
                f"(queue size: {target_queue.qsize()})"
            )

        except Exception as e:
            # Queue full or other error - drop frame with backpressure
            self.logger.warning(f"Dropped frame for camera {camera_id}: {e}")

    # Helper methods
    def _extract_message_key(self, message_data: Dict[str, Any]) -> str:
        """Extract message key."""
        return message_data.get("key", "")

    def _parse_message_data(self, message_data: Dict[str, Any]) -> Dict[str, Any]:
        """Parse message data."""
        return message_data.get("data", {})

    def _extract_input_stream(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """Extract input_stream from data."""
        input_stream = data.get("input_stream", {})
        if not isinstance(input_stream, dict):
            input_stream = {}
        return input_stream

    def _normalize_extra_params(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """Normalize extra_params."""
        extra_params = data.get("extra_params", {})
        if isinstance(extra_params, dict):
            return extra_params
        elif isinstance(extra_params, list):
            merged = {}
            for item in extra_params:
                if isinstance(item, dict):
                    merged.update(item)
            return merged
        return {}

    def _determine_frame_id(
        self,
        data: Dict[str, Any],
        message_data: Dict[str, Any],
        camera_id: str
    ) -> str:
        """Determine unique frame ID."""
        # Priority: upstream frame_id > message key > generated
        frame_id = data.get("frame_id")
        if frame_id and isinstance(frame_id, str):
            return frame_id

        message_key = message_data.get("key", "")
        if message_key and isinstance(message_key, str) and len(message_key) > 10:
            return message_key

        # Generate unique ID
        return f"{camera_id}_{uuid.uuid4()}"

    def _enrich_input_stream(self, input_stream: Dict[str, Any], frame_id: str):
        """Enrich input_stream with frame_id."""
        if isinstance(input_stream, dict):
            input_stream["frame_id"] = frame_id

    def _create_stream_message(self, camera_id: str, message_key: str, data: Dict[str, Any]) -> StreamMessage:
        """Create StreamMessage object."""
        from datetime import datetime

        # Get timestamp from data or use current time
        timestamp = data.get("timestamp")
        if not isinstance(timestamp, datetime):
            timestamp = datetime.now()

        return StreamMessage(
            camera_id=camera_id,
            message_key=message_key,
            data=data,
            timestamp=timestamp,
            priority=1,
        )
