import asyncio
import base64
import logging
import queue
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from typing import Any, Dict, Optional
from matrice_inference.server.stream.worker_metrics import WorkerMetrics

class InferenceWorker:
    """Handles inference processing with per-camera queue routing for order preservation.

    Uses hash-based camera assignment:
    - Each worker is assigned cameras via: hash(camera_id) % num_workers
    - Same camera always routes to same worker (preserves ordering)
    - Different cameras processed in parallel (enables 100x throughput)
    """

    def __init__(
        self,
        worker_id: int,
        num_workers: int,
        inference_queues: Dict[str, queue.Queue],
        postproc_queues: Dict[str, queue.Queue],
        pipeline: Any,
        inference_executor: ThreadPoolExecutor,
        message_timeout: float,
        inference_timeout: float,
        inference_interface: Optional[Any] = None,
        use_shared_metrics: Optional[bool] = True,
        use_async_inference: bool = False,
    ):
        self.worker_id = worker_id
        self.num_workers = num_workers
        self.inference_queues = inference_queues
        self.postproc_queues = postproc_queues
        self.pipeline = pipeline
        self.inference_executor = inference_executor
        self.message_timeout = message_timeout
        self.inference_timeout = inference_timeout
        self.inference_interface = inference_interface
        self.use_async_inference = use_async_inference
        self.running = False
        # self.metrics = WorkerMetrics(
        #     worker_id=f"inference_worker_{worker_id}",
        #     worker_type="inference"
        # )
        # self.metrics = WorkerMetrics.get_shared("inference")
        if use_shared_metrics:
            self.metrics = WorkerMetrics.get_shared("inference")
        else:
            self.metrics = WorkerMetrics(
                worker_id=f"inference_worker_{worker_id}",
                worker_type="inference"
            )

        self.logger = logging.getLogger(f"{__name__}.inference.{worker_id}")
    
    def start(self) -> threading.Thread:
        """Start the inference worker in a separate thread."""
        self.running = True
        self.metrics.mark_active() # ADD : Mark worker as active
        thread = threading.Thread(
            target=self._run,
            name=f"InferenceWorker-{self.worker_id}",
            daemon=False
        )
        thread.start()
        return thread
    
    def stop(self):
        """Stop the inference worker."""
        self.running = False
        self.metrics.mark_inactive()  # ADD: Mark as inactive

    
    def _run(self) -> None:
        """Main inference dispatcher loop with per-camera queue routing.

        Polls all camera queues assigned to this worker (via hash routing).
        Each camera's frames are processed in order, while different cameras
        are processed in parallel across workers.
        """
        self.logger.info(f"Started inference worker {self.worker_id}/{self.num_workers}")

        try:
            while self.running:
                task = self._get_task_from_queue()
                if task:
                    camera_id, task_data = task
                    self._process_inference_task(camera_id, task_data)
        except Exception as e:
            self.logger.error(f"Fatal error in inference worker: {e}")
        finally:
            self.logger.info(f"Inference worker {self.worker_id} stopped")

    def _get_task_from_queue(self) -> Optional[tuple]:
        """Get task from assigned camera queues using hash-based routing.

        Polls all camera queues, but only processes cameras assigned to this worker.
        Assignment: hash(camera_id) % num_workers == worker_id

        Returns:
            Tuple of (camera_id, task_data) or None if no tasks available
        """
        try:
            # Iterate through all camera queues (shallow copy to avoid modification during iteration)
            camera_ids = list(self.inference_queues.keys())

            for camera_id in camera_ids:
                # Check if this camera is assigned to this worker
                if hash(camera_id) % self.num_workers != self.worker_id:
                    continue

                # Try to get task from this camera's queue (non-blocking)
                try:
                    task_data = self.inference_queues[camera_id].get_nowait()
                    return (camera_id, task_data)
                except queue.Empty:
                    continue

            # No tasks available from any assigned camera
            time.sleep(0.01)  # Small sleep to prevent busy-waiting
            return None
        except Exception as e:
            self.logger.error(f"Error getting task from queue: {e}")
            return None
        
    def _process_inference_task(self, camera_id: str, task_data: Dict[str, Any]) -> None:
        """Process a single inference task with per-camera queue routing.

        Args:
            camera_id: Camera identifier for routing results to correct postproc queue
            task_data: Inference task data
        """
        start_time = time.time()
        try:
            if not self._validate_task_data(task_data):
                return

            result = self._execute_inference(task_data)
            processing_time = time.time() - start_time

            if result["success"]:
                postproc_task = self._create_postprocessing_task(
                    task_data, result, processing_time
                )
                # Route to per-camera postproc queue (preserves ordering)
                postproc_queue = self.pipeline.get_postproc_queue(camera_id)
                postproc_queue.put(postproc_task)
            else:
                self.logger.error(f"Inference failed: {result['error']}")

        except Exception as e:
            self.logger.error(f"Inference task error: {e}")

        finally:
            # ADD: Record metrics regardless of success/failure
            latency_ms = (time.time() - start_time) * 1000
            self.metrics.record_latency(latency_ms)
            self.metrics.record_throughput(count=1)

    def _validate_task_data(self, task_data: Dict[str, Any]) -> bool:
        """Validate that task data contains required fields."""
        required_fields = ["message", "input_stream", "stream_key", "camera_config"]
        for field in required_fields:
            if field not in task_data:
                self.logger.error(f"Missing required field '{field}' in task data")
                return False
        return True

    def _execute_inference(self, task_data: Dict[str, Any]) -> Dict[str, Any]:
        """Execute inference task in thread pool."""
        future = self.inference_executor.submit(self._run_inference, task_data)
        return future.result(timeout=self.inference_timeout)

    def _create_postprocessing_task(
        self, task_data: Dict[str, Any], result: Dict[str, Any], processing_time: float
    ) -> Dict[str, Any]:
        """Create post-processing task from inference result, preserving frame_id."""
        postproc_task = {
            "original_message": task_data["message"],
            "model_result": result["model_result"],
            "metadata": result["metadata"],
            "processing_time": processing_time,
            "input_stream": task_data["input_stream"],
            "stream_key": task_data["stream_key"],
            "camera_config": task_data["camera_config"]
        }
        
        # Preserve frame_id from task_data (critical for cache retrieval)
        if "frame_id" in task_data:
            postproc_task["frame_id"] = task_data["frame_id"]
            self.logger.debug(f"Preserved frame_id in postproc task: {task_data['frame_id']}")
        else:
            self.logger.warning("No frame_id in task_data to preserve")
        
        return postproc_task

    def _run_inference(self, task_data: Dict[str, Any]) -> Dict[str, Any]:
        """Run inference in thread pool with proper error handling and cleanup."""
        try:
            # Extract task data - handle camera streamer format
            input_stream_data = task_data.get("input_stream", {})
            stream_key = task_data.get("stream_key")
            stream_info = input_stream_data.get("stream_info", {})
            camera_info = input_stream_data.get("camera_info", {})
            extra_params = task_data.get("extra_params", {})
            
            # Ensure extra_params is a dictionary
            if not isinstance(extra_params, dict):
                logging.warning(f"extra_params is not a dict in inference worker, converting from {type(extra_params)}: {extra_params}")
                if isinstance(extra_params, list):
                    # Convert list to dict if possible
                    if len(extra_params) == 0:
                        extra_params = {}
                    elif all(isinstance(item, dict) for item in extra_params):
                        # Merge all dictionaries in the list
                        merged_params = {}
                        for item in extra_params:
                            merged_params.update(item)
                        extra_params = merged_params
                    else:
                        extra_params = {}
                else:
                    extra_params = {}
            
            if not self.inference_interface:
                raise ValueError("Inference interface not initialized")

            inference_params = self._extract_inference_params(task_data)
            loop = self._get_or_create_event_loop()

            # Use async_inference if available, otherwise use sync inference
            if self.use_async_inference:
                model_result, metadata = loop.run_until_complete(
                    self.inference_interface.async_inference(**inference_params)
                )
            else:
                model_result, metadata = loop.run_until_complete(
                    self.inference_interface.inference(**inference_params)
                )

            return self._create_success_result(model_result, metadata)

        except Exception as e:
            self.logger.error(f"Inference execution error: {e}", exc_info=True)
            return self._create_error_result(str(e))

    def _extract_inference_params(self, task_data: Dict[str, Any]) -> Dict[str, Any]:
        """Extract and validate inference parameters from task data."""
        input_stream_data = task_data.get("input_stream", {})
        # Prefer decoded bytes if provided by upstream stages
        input_bytes = task_data.get("decoded_input_bytes")
        if not isinstance(input_bytes, (bytes, bytearray)):
            content = input_stream_data.get("content")
            if isinstance(content, str) and content:
                try:
                    input_bytes = base64.b64decode(content)
                except Exception as e:
                    self.logger.warning(f"Failed to decode base64 content for inference: {e}")
                    input_bytes = None
            elif isinstance(content, (bytes, bytearray)):
                input_bytes = content
            else:
                input_bytes = None

        extra_params = self._normalize_extra_params(task_data.get("extra_params", {}))

        return {
            "input": input_bytes,
            "extra_params": extra_params,
            "apply_post_processing": False,
            "stream_key": task_data.get("stream_key"),
            "stream_info": input_stream_data.get("stream_info", {}),
            "camera_info": input_stream_data.get("camera_info", {})
        }

    def _decode_input_content(self, content: Any) -> Any:
        """Decode base64 content if it's a string."""
        if content and isinstance(content, str):
            try:
                return base64.b64decode(content)
            except Exception as e:
                self.logger.warning(f"Failed to decode base64 input: {e}")
        return content

    def _normalize_extra_params(self, extra_params: Any) -> Dict[str, Any]:
        """Normalize extra_params to ensure it's a dictionary."""
        if isinstance(extra_params, dict):
            return extra_params
        elif isinstance(extra_params, list):
            return self._merge_list_params(extra_params)
        else:
            self.logger.warning(f"Invalid extra_params type {type(extra_params)}, using empty dict")
            return {}

    def _merge_list_params(self, params_list: list) -> Dict[str, Any]:
        """Merge list of dictionaries into single dictionary."""
        if not params_list:
            return {}

        if all(isinstance(item, dict) for item in params_list):
            merged = {}
            for item in params_list:
                merged.update(item)
            return merged

        return {}

    def _get_or_create_event_loop(self) -> asyncio.AbstractEventLoop:
        """Get existing event loop or create a new one for this thread."""
        try:
            return asyncio.get_event_loop()
        except RuntimeError:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            return loop

    def _create_success_result(self, model_result: Any, metadata: Any) -> Dict[str, Any]:
        """Create successful inference result."""
        return {
            "model_result": model_result,
            "metadata": metadata,
            "success": True,
            "error": None
        }

    def _create_error_result(self, error_message: str) -> Dict[str, Any]:
        """Create error inference result."""
        return {
            "model_result": None,
            "metadata": None,
            "success": False,
            "error": error_message
        }

