"""
High-performance multiprocessing + async inference worker for optimal GPU utilization.

Architecture (keeps ALL performance optimizations):
- Multiprocessing: Multiple separate processes (one per GPU/core) - TRUE PARALLELISM
- Async Event Loops: Each process runs its own async event loop - NON-BLOCKING I/O
- InferenceInterface: Each process recreates InferenceInterface locally - PROPER ABSTRACTIONS
- Async Gather: Multiple concurrent inference requests via asyncio.create_task - MAX GPU UTILIZATION
- Dynamic Batching: Accumulate requests until batch_size or batch_timeout_ms - THROUGHPUT
- Camera Routing: hash(camera_id) % num_workers for load balancing - ORDER PRESERVATION
- Backpressure: Queue-based flow control - PREVENTS OOM

Architecture Flow:
- InferenceInterface → ModelManagerWrapper → ModelManager → async_predict (from predict.py)
- Uses normal ModelManager (NOT Triton) with predict functions from deploy.py pattern
- Each worker imports predict functions and recreates full inference stack
- Preprocessing/postprocessing handled separately in pipeline (consumer/post-processor)

Performance Targets:
- 10,000 FPS throughput
- <500ms latency (P50)
- 85%+ GPU utilization
- True parallelism (bypasses Python GIL)
"""

import asyncio
import base64
import logging
import multiprocessing as mp
import time
from typing import Any, Dict, Optional
from matrice_inference.server.stream.worker_metrics import WorkerMetrics


def inference_worker_process(
    worker_id: int,
    num_workers: int,
    input_queue: mp.Queue,
    output_queue: mp.Queue,
    model_config: Dict[str, Any],
    batch_size: int = 4,
    batch_timeout_ms: float = 10.0,
    max_concurrent_requests: int = 3,
):
    """
    Worker process for GPU inference with async event loop.

    Each process:
    1. Recreates InferenceInterface with ModelManagerWrapper + ModelManager
    2. Runs its own async event loop
    3. Handles multiple concurrent inference requests
    4. Processes tasks from shared multiprocessing queue

    Args:
        worker_id: Worker process ID
        num_workers: Total number of worker processes
        input_queue: Multiprocessing queue for incoming inference tasks
        output_queue: Multiprocessing queue for inference results
        model_config: Model configuration (action_id, predict functions, model_path, etc.)
        batch_size: Target batch size for dynamic batching
        batch_timeout_ms: Max time to wait for batch accumulation
        max_concurrent_requests: Max concurrent inference requests per worker
    """
    # Set up logging for this process
    logger = logging.getLogger(f"inference_worker_{worker_id}")
    logger.setLevel(logging.INFO)

    try:
        # Import dependencies inside process to avoid pickle issues
        from matrice.action_tracker import ActionTracker
        from matrice_inference.server.model.model_manager_wrapper import ModelManagerWrapper
        from matrice_inference.server.inference_interface import InferenceInterface
        from matrice_inference.server.stream.worker_metrics import WorkerMetrics

        # ARCHITECTURE NOTE: Model Loading Strategy
        # ==========================================
        # Models are loaded TWICE - once in pipeline event loop, once in each worker process.
        # This is INTENTIONAL and necessary for:
        # 1. Process Isolation: Each worker needs its own model instance (no shared state)
        # 2. GIL-Free Parallelism: Separate processes bypass Python GIL for true parallelism
        # 3. GPU Utilization: Each process can fully utilize GPU without contention
        # 4. Fault Isolation: Worker crash doesn't affect other workers or main pipeline
        #
        # The slight startup time/memory overhead is acceptable for 10K+ FPS throughput.

        # Get predict functions from model_config (passed from MatriceDeployServer)
        # These are module-level functions that CAN be pickled by reference
        load_model_fn = model_config.get("load_model")
        predict_fn = model_config.get("predict")
        async_predict_fn = model_config.get("async_predict")
        async_load_model_fn = model_config.get("async_load_model")
        batch_predict_fn = model_config.get("batch_predict")

        # Create ActionTracker for this worker
        action_id = model_config.get("action_id")
        if not action_id:
            raise ValueError("action_id is required in model_config")

        action_tracker = ActionTracker(action_id)

        # Create ModelManagerWrapper with ModelManager (NOT Triton)
        model_manager_wrapper = ModelManagerWrapper(
            action_tracker=action_tracker,
            model_type="default",  # Use default ModelManager, NOT triton
            load_model=load_model_fn,
            predict=predict_fn,
            async_predict=async_predict_fn,
            async_load_model=async_load_model_fn,
            batch_predict=batch_predict_fn,
            num_model_instances=model_config.get("num_model_instances", 1),
            model_path=model_config.get("model_path"),
        )

        # Create InferenceInterface
        inference_interface = InferenceInterface(
            model_manager_wrapper=model_manager_wrapper,
            post_processor=None,  # Post-processing handled separately in pipeline
        )

        # Initialize metrics for this worker (shared across all inference workers)
        metrics = WorkerMetrics.get_shared("inference")
        metrics.mark_active()

        logger.info(
            f"Worker {worker_id}/{num_workers} initialized with InferenceInterface "
            f"(ModelManager with async_predict={'available' if async_predict_fn else 'not available'}, "
            f"num_instances={model_config.get('num_model_instances', 1)})"
        )

        # Run async event loop in this process
        try:
            asyncio.run(
                _async_inference_loop(
                    worker_id=worker_id,
                    num_workers=num_workers,
                    inference_interface=inference_interface,
                    input_queue=input_queue,
                    output_queue=output_queue,
                    batch_size=batch_size,
                    batch_timeout_ms=batch_timeout_ms,
                    max_concurrent_requests=max_concurrent_requests,
                    logger=logger,
                    metrics=metrics,
                )
            )
        finally:
            # Mark worker as inactive when stopping
            metrics.mark_inactive()

    except Exception as e:
        logger.error(f"Worker {worker_id} crashed: {e}", exc_info=True)
        raise


async def _async_inference_loop(
    worker_id: int,
    num_workers: int,
    inference_interface: Any,
    input_queue: mp.Queue,
    output_queue: mp.Queue,
    batch_size: int,
    batch_timeout_ms: float,
    max_concurrent_requests: int,
    logger: logging.Logger,
    metrics: Any,  # WorkerMetrics instance
):
    """
    Async event loop for inference worker process.

    Features:
    - Dynamic batching (accumulate requests for batch_timeout_ms or until batch_size)
    - Multiple concurrent inference requests via InferenceInterface (avoid GPU idle time)
    - Camera-based routing (only process tasks assigned to this worker)
    - Metrics tracking (latency and throughput)
    """
    batch_buffer = []
    batch_start_time = None
    active_requests = 0

    while True:
        try:
            # Try to get task from queue (non-blocking)
            try:
                # Use a short timeout to periodically check batch conditions
                task = input_queue.get(timeout=0.01)  # Increased from 0.001 to 0.01

                logger.debug(f"Worker {worker_id} received task for camera {task.get('camera_id')}")

                # Check if this task is assigned to this worker (camera-based routing)
                camera_id = task.get("camera_id")
                if camera_id and hash(camera_id) % num_workers != worker_id:
                    # Not for this worker, put it back
                    logger.debug(f"Worker {worker_id} re-queuing task for camera {camera_id} (assigned to worker {hash(camera_id) % num_workers})")
                    input_queue.put(task)
                    await asyncio.sleep(0.0001)
                    continue

                logger.info(f"Worker {worker_id} added task to batch (camera {camera_id}, batch size: {len(batch_buffer) + 1})")
                batch_buffer.append(task)
                if batch_start_time is None:
                    batch_start_time = time.time()

            except Exception as e:
                # Queue empty or timeout
                if "Empty" not in str(type(e)):
                    logger.debug(f"Worker {worker_id} queue read error: {e}")
                pass

            # Check if should process batch
            should_process = False
            if batch_buffer:
                elapsed_ms = (time.time() - batch_start_time) * 1000 if batch_start_time else 0
                should_process = (
                    len(batch_buffer) >= batch_size or
                    elapsed_ms > batch_timeout_ms
                )

            # Process batch if ready and we have capacity for more concurrent requests
            if should_process and active_requests < max_concurrent_requests:
                # Create async task for batch processing
                batch_to_process = batch_buffer.copy()
                logger.info(f"Worker {worker_id} processing batch of {len(batch_to_process)} frames")
                batch_buffer.clear()
                batch_start_time = None

                # Process batch concurrently
                active_requests += 1
                asyncio.create_task(
                    _process_batch_async(
                        worker_id=worker_id,
                        batch=batch_to_process,
                        inference_interface=inference_interface,
                        output_queue=output_queue,
                        logger=logger,
                        metrics=metrics,
                    )
                )

                # Decrement active_requests when done (handled in _process_batch_async)

            # Small sleep to prevent busy loop
            await asyncio.sleep(0.0001)

        except Exception as e:
            logger.error(f"Error in inference loop: {e}", exc_info=True)
            batch_buffer.clear()
            batch_start_time = None
            await asyncio.sleep(0.1)


async def _process_batch_async(
    worker_id: int,
    batch: list,
    inference_interface: Any,
    output_queue: mp.Queue,
    logger: logging.Logger,
    metrics: Any,  # WorkerMetrics instance
):
    """
    Process a batch of inference requests asynchronously.

    This runs concurrently with other batches in the same worker process,
    allowing multiple inference requests to be in-flight simultaneously.
    Records latency and throughput metrics.
    """
    start_time = time.time()

    try:
        # Create concurrent inference tasks via InferenceInterface
        inference_tasks = []
        for task_data in batch:
            inference_task = asyncio.create_task(
                _execute_single_inference(inference_interface, task_data, logger)
            )
            inference_tasks.append((task_data, inference_task))

        # Wait for all inferences to complete
        successful_count = 0
        for task_data, inference_task in inference_tasks:
            try:
                result = await inference_task

                # Put result in output queue for post-processing
                output_data = {
                    "camera_id": task_data.get("camera_id"),
                    "frame_id": task_data.get("frame_id"),
                    "original_message": task_data.get("message"),
                    "model_result": result.get("model_result"),
                    "metadata": result.get("metadata", {}),
                    "processing_time": time.time() - start_time,
                    "input_stream": task_data.get("input_stream", {}),
                    "stream_key": task_data.get("stream_key"),
                    "camera_config": task_data.get("camera_config"),
                }

                output_queue.put(output_data)
                successful_count += 1

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

        # Record metrics for successfully processed frames
        batch_latency_ms = (time.time() - start_time) * 1000
        if successful_count > 0:
            metrics.record_latency(batch_latency_ms)
            metrics.record_throughput(count=successful_count)

        logger.debug(
            f"Worker {worker_id}: Processed batch size={len(batch)}, "
            f"successful={successful_count}, latency={batch_latency_ms:.1f}ms"
        )

    except Exception as e:
        logger.error(f"Batch processing error: {e}", exc_info=True)


async def _execute_single_inference(
    inference_interface: Any,
    task_data: Dict[str, Any],
    logger: logging.Logger,
) -> Dict[str, Any]:
    """
    Execute async inference for a single task using InferenceInterface.

    Args:
        inference_interface: InferenceInterface instance (with ModelManagerWrapper → ModelManager)
        task_data: Task data containing input and parameters
        logger: Logger instance

    Returns:
        Dict with model_result and metadata
    """
    try:
        # Validate task data
        if not _validate_task_data(task_data, logger):
            return {"model_result": None, "metadata": {}, "success": False}

        # Extract inference parameters
        input_bytes = _extract_input_bytes(task_data, logger)
        if input_bytes is None:
            return {"model_result": None, "metadata": {}, "success": False}

        # Call InferenceInterface.async_inference()
        # Flow: InferenceInterface → ModelManagerWrapper → ModelManager → async_predict
        # Note: Raw bytes → model inference → raw results (no post-processing here)
        # Postprocessing handled separately by post_processing_manager
        model_result, metadata = await inference_interface.async_inference(
            input=input_bytes,
            extra_params=task_data.get("extra_params"),
            apply_post_processing=False,  # No post-processing in workers
            stream_key=task_data.get("stream_key"),
            stream_info=None,
        )

        return {
            "success": True,
            "model_result": model_result,
            "metadata": metadata or {},
        }

    except Exception as e:
        logger.error(f"Inference execution error: {e}", exc_info=True)
        return {
            "success": False,
            "error": str(e),
            "model_result": None,
            "metadata": {},
        }


def _validate_task_data(task_data: Dict[str, Any], logger: logging.Logger) -> bool:
    """
    Validate that task data contains required fields.

    Required fields (from consumer_manager):
    - camera_id: For routing and identification
    - frame_bytes or input_stream: Input data
    - message: Original stream message
    - stream_key: Stream identifier
    - camera_config: Camera configuration
    """
    required_fields = ["camera_id", "message", "stream_key", "camera_config"]
    for field in required_fields:
        if field not in task_data:
            logger.error(f"Missing required field '{field}' in task data")
            return False

    # Check that we have input data (frame_bytes or input_stream)
    has_input = (
        "frame_bytes" in task_data or
        "decoded_input_bytes" in task_data or
        "input_stream" in task_data
    )
    if not has_input:
        logger.error("No input data found (frame_bytes, decoded_input_bytes, or input_stream)")
        return False

    return True


def _extract_input_bytes(task_data: Dict[str, Any], logger: logging.Logger) -> Optional[bytes]:
    """
    Extract input bytes from task data.

    Supports multiple formats (priority order):
    1. task_data["frame_bytes"] - direct bytes from consumer_manager (NEW)
    2. task_data["decoded_input_bytes"] - decoded bytes
    3. task_data["input_stream"]["content"] - bytes or base64 string
    """
    # Priority 1: Direct frame_bytes from consumer_manager (simplified flow)
    frame_bytes = task_data.get("frame_bytes")
    if isinstance(frame_bytes, (bytes, bytearray)) and frame_bytes:
        return bytes(frame_bytes)

    # Priority 2: Decoded input bytes
    decoded_bytes = task_data.get("decoded_input_bytes")
    if isinstance(decoded_bytes, (bytes, bytearray)) and decoded_bytes:
        return bytes(decoded_bytes)

    # Priority 3: Extract from input_stream
    input_stream_data = task_data.get("input_stream", {})
    if not isinstance(input_stream_data, dict):
        logger.error(f"input_stream is not a dict: {type(input_stream_data)}")
        return None

    content = input_stream_data.get("content")

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

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

    logger.error("No valid input bytes found in task data")
    return None


class MultiprocessInferencePool:
    """
    Pool of multiprocessing inference workers with async event loops.

    Architecture:
    - Creates multiple worker processes (one per GPU/core)
    - Each process recreates InferenceInterface → ModelManagerWrapper → ModelManager
    - Uses normal ModelManager with async_predict from predict.py (NOT Triton)
    - Each process runs its own async event loop
    - Processes communicate via multiprocessing queues
    - Camera-based routing for task distribution
    """

    def __init__(
        self,
        num_workers: int,
        model_config: Dict[str, Any],
        input_queue: mp.Queue,
        output_queue: mp.Queue,
        batch_size: int = 4,
        batch_timeout_ms: float = 10.0,
        max_concurrent_requests: int = 3,
    ):
        self.num_workers = num_workers
        self.model_config = model_config
        self.batch_size = batch_size
        self.batch_timeout_ms = batch_timeout_ms
        self.max_concurrent_requests = max_concurrent_requests

        # Use shared multiprocessing queues from pipeline (not created here)
        self.input_queue = input_queue
        self.output_queue = output_queue

        self.processes = []
        self.running = False

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

    def start(self):
        """Start all worker processes."""
        self.running = True

        for worker_id in range(self.num_workers):
            process = mp.Process(
                target=inference_worker_process,
                args=(
                    worker_id,
                    self.num_workers,
                    self.input_queue,
                    self.output_queue,
                    self.model_config,
                    self.batch_size,
                    self.batch_timeout_ms,
                    self.max_concurrent_requests,
                ),
                daemon=True,
            )
            process.start()
            self.processes.append(process)

        self.logger.info(
            f"Started {self.num_workers} multiprocess inference workers "
            f"(batch_size={self.batch_size}, concurrent_requests={self.max_concurrent_requests})"
        )

    def stop(self):
        """Stop all worker processes."""
        self.running = False

        for process in self.processes:
            if process.is_alive():
                process.terminate()
                process.join(timeout=5)
                if process.is_alive():
                    process.kill()

        self.processes.clear()
        self.logger.info("Stopped all inference worker processes")

    def submit_task(self, task_data: Dict[str, Any], timeout: float = 0.1) -> bool:
        """
        Submit inference task to worker pool.

        Args:
            task_data: Task data with camera_id, frame, etc.
            timeout: Max time to wait if queue is full

        Returns:
            True if task was submitted, False if queue full (backpressure)
        """
        try:
            self.input_queue.put(task_data, timeout=timeout)
            return True
        except Exception:
            # Queue full - apply backpressure
            return False

    def get_result(self, timeout: float = 0.001) -> Optional[Dict[str, Any]]:
        """
        Get inference result from worker pool.

        Args:
            timeout: Max time to wait for result

        Returns:
            Result dict or None if no result available
        """
        try:
            return self.output_queue.get(timeout=timeout)
        except Exception:
            return None
