"""Stub file for server.model directory."""
from typing import Any, Dict, List, Optional, Tuple, Union

from PIL import Image
from datetime import datetime, timezone
from io import BytesIO
from matrice.action_tracker import ActionTracker
from matrice_common.utils import dependencies_check
from matrice_inference.server.model.model_manager import ModelManager
from matrice_inference.server.model.triton_model_manager import TritonModelManager
from matrice_inference.server.model.triton_server import TritonServer, TritonInference
from ultralytics import YOLO
import asyncio
import cv2
import gc
import httpx
import importlib.util
import logging
import numpy as np
import onnx
import os
import requests
import shlex
import shutil
import subprocess
import tempfile
import tensorrt as trt
import threading
import time
import torch
import torchvision
import tritonclient.grpc as grpcclient
import tritonclient.grpc as tritonclientclass
import tritonclient.http as tritonclientclass
import zipfile

# Constants
BASE_PATH: Any = ...  # From triton_server

# Classes
# From model_manager
class ModelManager:
    """
    Minimal ModelManager that focuses on model lifecycle and prediction calls.
    """

    def __init__(self, action_tracker, load_model: Optional[Callable] = None, predict: Optional[Callable] = None, batch_predict: Optional[Callable] = None, async_predict: Optional[Callable] = None, async_batch_predict: Optional[Callable] = None, async_load_model: Optional[Callable] = None, num_model_instances: int = 1, model_path: Optional[str] = None) -> None: ...
        """
        Initialize the ModelManager
        
                Args:
                    action_tracker: Tracker for monitoring actions.
                    load_model: Function to load the model (synchronous).
                    predict: Function to run single predictions (sync).
                    batch_predict: Function to run batch predictions (sync).
                    async_predict: Function to run single predictions (async).
                    async_batch_predict: Function to run batch predictions (async).
                    async_load_model: Function to load the model asynchronously (loaded lazily in worker thread's event loop).
                    num_model_instances: Number of model instances to create.
                    model_path: Path to the model directory.
        """

    async def async_batch_inference(self, input: List[bytes], extra_params: Dict[str, Any] = None, stream_key: Optional[str] = None, stream_info: Optional[Dict[str, Any]] = None) -> Tuple[List[dict], bool]: ...
        """
        Run asynchronous batch inference on the provided input data.
        
                Args:
                    input: List of input data (e.g., image bytes)
                    extra_params: Additional parameters for inference (optional)
                    stream_key: Stream key for the inference
                    stream_info: Stream info for the inference
                Returns:
                    Tuple of (results_list, success_flag)
        
                Raises:
                    ValueError: If input data is invalid
        """

    async def async_inference(self, input, extra_params: Dict[str, Any] = None, stream_key: Optional[str] = None, stream_info: Optional[Dict[str, Any]] = None) -> Tuple[dict, bool]: ...
        """
        Run asynchronous inference on the provided input data.
        
                Args:
                    input: Primary input data (can be image bytes or numpy array)
                    extra_params: Additional parameters for inference (optional)
                    stream_key: Stream key for the inference
                    stream_info: Stream info for the inference
                Returns:
                    Tuple of (results, success_flag)
        
                Raises:
                    ValueError: If input data is invalid
        """

    def batch_inference(self, input: List[bytes], extra_params: Dict[str, Any] = None, stream_key: Optional[str] = None, stream_info: Optional[Dict[str, Any]] = None) -> Tuple[List[dict], bool]: ...
        """
        Run synchronous batch inference on the provided input data.
        
                If batch_predict is not available, falls back to calling predict
                on each input individually (similar to async_inference fallback).
        
                Args:
                    input: List of input data (e.g., image bytes)
                    extra_params: Additional parameters for inference (optional)
                    stream_key: Stream key for the inference
                    stream_info: Stream info for the inference
                Returns:
                    Tuple of (results_list, success_flag)
        
                Raises:
                    ValueError: If input data is invalid
        """

    async def ensure_models_loaded(self) -> Any: ...
        """
        Ensure all model instances are loaded in the current event loop.
        
                This method MUST be called in the StreamingPipeline's event loop
                before inference begins. It loads models asynchronously if async_load_model
                is provided, otherwise loads synchronously.
        
                This ensures all models are loaded in the same event loop used for inference.
        
                Returns:
                    bool: True if all models loaded successfully, False otherwise
        """

    def get_model(self) -> Any: ...
        """
        Get the model instance in round-robin fashion.
        
                Models MUST be loaded before calling this method by calling ensure_models_loaded()
                in the StreamingPipeline's event loop. This ensures all async operations use
                the same event loop.
        
                Returns:
                    model: The loaded model instance
        
                Raises:
                    RuntimeError: If model is not loaded
        """

    def inference(self, input, extra_params: Dict[str, Any] = None, stream_key: Optional[str] = None, stream_info: Optional[Dict[str, Any]] = None) -> Tuple[dict, bool]: ...
        """
        Run inference on the provided input data.
        
                Args:
                    input: Primary input data (can be image bytes or numpy array)
                    extra_params: Additional parameters for inference (optional)
                    stream_key: Stream key for the inference
                    stream_info: Stream info for the inference
                Returns:
                    Tuple of (results, success_flag)
        
                Raises:
                    ValueError: If input data is invalid
        """

    def scale_down(self) -> Any: ...
        """
        Unload the model from memory (scale down)
        """

    def scale_up(self) -> Any: ...
        """
        Load the model into memory (scale up)
        """


# From model_manager_wrapper
class ModelManagerWrapper:
    """
    Wrapper class for ModelManager and TritonModelManager to provide a unified interface.
    """

    def __init__(self, action_tracker, test_env: bool = False, model_type: str = 'default', num_model_instances: Optional[int] = None, load_model: Optional[Callable] = None, predict: Optional[Callable] = None, batch_predict: Optional[Callable] = None, model_name: Optional[str] = None, model_path: Optional[str] = None, runtime_framework: Optional[str] = None, internal_server_type: Optional[str] = None, internal_port: Optional[int] = None, internal_host: Optional[str] = None, input_size: Optional[Union[int, List[int]]] = None, num_classes: Optional[int] = None, use_dynamic_batching: Optional[bool] = None, max_batch_size: Optional[int] = None, is_yolo: Optional[bool] = None, is_ocr: Optional[bool] = None, use_trt_accelerator: Optional[bool] = None, preprocess_fn: Optional[Callable] = None, postprocess_fn: Optional[Callable] = None, preprocess_params: Optional[Dict[str, Any]] = None, postprocess_params: Optional[Dict[str, Any]] = None, async_predict: Optional[Callable] = None, async_batch_predict: Optional[Callable] = None, async_load_model: Optional[Callable] = None) -> None: ...
        """
        Initialize the ModelManagerWrapper.
        
        Args:
            action_tracker: Action tracker for category mapping and configuration.
            test_env: If True, use provided parameters for testing; if False, extract from action_tracker.
            model_type: Type of model manager ("default" for ModelManager, "triton" for TritonModelManager).
            internal_server_type: Type of internal server (e.g., "rest", "grpc").
            internal_port: Internal port number.
            internal_host: Internal host address.
            num_model_instances: Number of model instances to create.
            load_model: Function to load the model (for ModelManager).
            predict: Function to run predictions (for ModelManager).
            batch_predict: Function to run batch predictions (for ModelManager).
            model_name: Name of the model (for TritonModelManager).
            model_path: Path to the model (for TritonModelManager).
            runtime_framework: Runtime framework for the model (for TritonModelManager).
            input_size: Input size for the model (for TritonModelManager).
            num_classes: Number of classes for the model (for TritonModelManager).
            use_dynamic_batching: Whether to use dynamic batching (for TritonModelManager).
            max_batch_size: Maximum batch size (for TritonModelManager).
            is_yolo: Whether the model is YOLO (for TritonModelManager).
            is_ocr: Whether the model is OCR (for TritonModelManager).
            use_trt_accelerator: Whether to use TensorRT accelerator (for TritonModelManager).
            preprocess_fn: User-provided preprocessing function (optional).
            postprocess_fn: User-provided postprocessing function (optional).
            preprocess_params: Parameters for the preprocessing function.
            postprocess_params: Parameters for the postprocessing function.
            async_predict: Function for single async predictions. Defaults to None.
            async_batch_predict: Function for batch async predictions. Defaults to None.
            async_load_model: Function to load model asynchronously (loaded lazily in worker thread's event loop). Defaults to None.
        """

    async def async_batch_inference(self, input: List[Any], extra_params: Optional[Dict[str, Any]] = None, stream_key: Optional[str] = None, stream_info: Optional[Dict[str, Any]] = None) -> Tuple[List[Any], bool]: ...
        """
        Perform asynchronous batch inference.
        
        Uses async_batch_predict if available, otherwise falls back to sync batch_inference.
        """

    async def async_inference(self, input: Union[bytes, np.ndarray], extra_params: Optional[Dict[str, Any]] = None, stream_key: Optional[str] = None, stream_info: Optional[Dict[str, Any]] = None) -> Tuple[Any, bool]: ...
        """
        Perform asynchronous single inference.
        """

    def batch_inference(self, input: List[Any], extra_params: Optional[Dict[str, Any]] = None, stream_key: Optional[str] = None, stream_info: Optional[Dict[str, Any]] = None) -> Tuple[List[Any], bool]: ...
        """
        Perform synchronous batch inference.
        """

    def inference(self, input, extra_params: Optional[Dict[str, Any]] = None, stream_key: Optional[str] = None, stream_info: Optional[Dict[str, Any]] = None) -> Tuple[Any, bool]: ...
        """
        Perform synchronous single inference.
        """


# From triton_model_manager
class TritonModelManager:
    """
    Model manager for Triton Inference Server, aligned with pipeline and inference interface.
    """

    def __init__(self, model_name: str, model_path: str, runtime_framework: str, internal_server_type: str, internal_port: int, internal_host: str, input_size: Union[int, List[int]] = 640, num_classes: int = 10, num_model_instances: int = 1, use_dynamic_batching: bool = False, max_batch_size: int = 8, is_yolo: bool = False, is_ocr: bool = False, use_trt_accelerator: bool = False, preprocess_fn = None, postprocess_fn = None, preprocess_params: Dict[str, Any] = None, postprocess_params: Dict[str, Any] = None) -> None: ...

    async def async_batch_inference(self, input: List[bytes]) -> Tuple[List[Any], bool]: ...
        """
        Perform asynchronous batch inference using TritonInference client.
        
                Args:
                    input: List of primary input data (e.g., image bytes).
        
                Returns:
                    Tuple of (results_list, success_flag).
        """

    async def async_inference(self, input: Union[bytes, np.ndarray]) -> Tuple[Any, bool]: ...
        """
        Perform asynchronous single inference using TritonInference client.
                Args:
                    input: Primary input data (Image bytes or numpy array).
        
                Returns:
                    Tuple of (results, success_flag).
        """

    def batch_inference(self, input: List[bytes]) -> Tuple[List[Any], bool]: ...
        """
        Perform synchronous batch inference using TritonInference client.
        
                Args:
                    input: List of primary input data (e.g., image bytes).
        
                Returns:
                    Tuple of (results_list, success_flag).
        """

    def inference(self, input) -> Tuple[Any, bool]: ...
        """
        Perform synchronous single inference using TritonInference client.
        
                Args:
                    input: Primary input data (e.g., image bytes).
        
                Returns:
                    Tuple of (results, success_flag).
        """


# From triton_server
class TritonInference:
    """
    Class for making Triton inference requests.
    """

    def __init__(self, server_type: str, model_name: str, internal_port: int = 80, internal_host: str = 'localhost', task_type: str = 'detection', runtime_framework: str = 'onnx', is_yolo: bool = False, is_ocr: bool = False, input_size: Union[int, List[int]] = (224, 224)) -> None: ...
        """
        Initialize Triton inference client.
        
                Args:
                    server_type: Type of server (grpc/rest)
                    model_name: Name of model to use
                    internal_port: Port number for internal API
                    internal_host: Hostname for internal API
                    task_type: Type of task (e.g., detection)
                    runtime_framework: Framework used for the model (e.g., onnx)
                    is_yolo: Boolean indicating if the model is YOLO
                    is_ocr: Boolean indicating if the model is an OCR model
                    input_size: Input size for the model (int or [height, width])
        """

    async def async_inference(self, input_data: Union[bytes, np.ndarray]) -> Any: ...
        """
        Make an asynchronous inference request (REST + gRPC).
        """

    def format_response(self, response) -> Dict[str, Any]: ...
        """
        Format model response for consistent logging.
        
                Args:
                    response: Raw model output
        
                Returns:
                    Formatted response dictionary
        """

    def inference(self, input_data: Union[bytes, np.ndarray]) -> Any: ...
        """
        Make a synchronous inference request.
        
                Args:
                    input_data: Input data as bytes or stacked numpy array
        
                Returns:
                    Model prediction as numpy array
        
                Raises:
                    Exception: If inference fails
        """


# From triton_server
class TritonServer:
    def __init__(self, model_name: str, model_path: str, runtime_framework: str, input_size: Union[int, List[int]] = 224, num_classes: int = 10, dynamic_batching: bool = False, num_model_instances: int = 1, max_batch_size: int = 8, connection_protocol: str = 'rest', is_yolo: bool = False, is_ocr: bool = False, use_trt_accelerator: bool = False, **kwargs) -> None: ...
        """
        Initialize the Triton server.
        
                Args:
                    model_name: Name of the model (used for Triton model repository).
                    model_path: Path to the model file on the local filesystem.
                    runtime_framework: Framework of the model ('onnx', 'pytorch', 'torchscript', 'yolo', 'tensorrt', 'openvino').
                    input_size: Input size for the model (int for square images or [height, width]).
                    num_classes: Number of output classes.
                    dynamic_batching: Enable dynamic batching for the model.
                    num_model_instances: Number of model instances to deploy.
                    max_batch_size: Maximum batch size for inference.
                    connection_protocol: Protocol for Triton server ('rest' or 'grpc').
                    use_trt_accelerator: Enable TensorRT acceleration for inference.
        
                    is_yolo: Boolean indicating if the model is a YOLO model.
                    is_ocr: Boolean indicating if the model is an OCR model.
        """

    def create_model_repository(self) -> Any: ...
        """
        Create the model repository directory structure
        """

    def get_config_params(self) -> Any: ...
        """
        Get configuration parameters for Triton config file
        """

    def prepare_model(self, model_version_dir: str) -> None: ...
        """
        Prepare the model file for Triton Inference Server.
        
                Copies model from self.model_path to model_version_dir and converts if necessary
                to the format expected by Triton (model.onnx, model.xml, model.plan).
        
                Args:
                    model_version_dir: Directory to store the model file
                    (e.g., '/models/<model_name>/1').
        """

    def setup(self, internal_port: int = 8000) -> Any: ...
        """
        Setup the Triton server with the provided model.
        
                Args:
                    internal_port: Port to expose the server on
        """

    def start_server(self, internal_port: int = 8000) -> Any: ...
        """
        Start the Triton Inference Server
        
                Args:
                    internal_port: Port to expose the server on (not relevant anymore here after ultimate dockerfile exposes it, needs to be taken care while launching the container)
        """

    def to_onnx(self, checkpoint_path: str, onnx_path: str, input_shape: Tuple[int, int, int, int]) -> None: ...
        """
        Export PyTorch or YOLO checkpoint to ONNX.
        """

    def write_config_file(self, model_dir: str, max_batch_size: int = 8, num_model_instances: int = 1, image_size: List[int] = [224, 224], num_classes: int = 10, input_data_type: str = 'TYPE_FP32', output_data_type: str = 'TYPE_FP32', dynamic_batching: bool = False, preferred_batch_size: list = [2, 4, 8], max_queue_delay_microseconds: int = 100, input_pinned_memory: bool = True, output_pinned_memory: bool = True, **kwargs) -> Any: ...
        """
        Write the model configuration file for Triton Inference Server.
        """


from . import model_manager, model_manager_wrapper, triton_model_manager, triton_server