"""Model loading utilities for quantized models."""

from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Optional, Union

import numpy as np


@dataclass
class QuantizationConfig:
    """
    Configuration for model quantization.

    Attributes:
        quantization_type: Type of quantization ('int8', 'uint8', 'int4')
        per_channel: Whether to use per-channel quantization
        symmetric: Whether to use symmetric quantization
        optimize_for_inference: Apply inference-specific optimizations
        calibration_samples: Number of samples for calibration
    """
    quantization_type: str = "int8"
    per_channel: bool = True
    symmetric: bool = True
    optimize_for_inference: bool = True
    calibration_samples: int = 100


class QuantizedModel:
    """
    Wrapper for quantized models with runtime optimization.

    This class provides a unified interface for loading and running quantized
    models regardless of the backend (ONNX, PyTorch, TensorFlow).

    Example:
        >>> model = QuantizedModel("models/quantized/jailbreak_int8.onnx")
        >>>
        >>> # Check model info
        >>> print(f"Model size: {model.model_size_mb:.1f} MB")
        >>> print(f"Quantization: {model.quantization_type}")
        >>>
        >>> # Run inference
        >>> output = model.predict(input_array)
    """

    def __init__(self, model_path: Union[str, Path], config: Optional[QuantizationConfig] = None):
        """
        Load a quantized model.

        Args:
            model_path: Path to the quantized model file
            config: Optional quantization configuration
        """
        self.model_path = Path(model_path)
        self.config = config or QuantizationConfig()
        self._model: Optional[Any] = None
        self._session: Optional[Any] = None

        if not self.model_path.exists():
            raise FileNotFoundError(f"Model not found: {model_path}")

        self._load_model()

    def _load_model(self) -> None:
        """Load the model based on file extension."""
        suffix = self.model_path.suffix.lower()

        if suffix == ".onnx":
            self._load_onnx()
        elif suffix in [".pt", ".pth"]:
            self._load_pytorch()
        elif suffix in [".pb", ".h5"]:
            self._load_tensorflow()
        else:
            raise ValueError(f"Unsupported model format: {suffix}")

    def _load_onnx(self) -> None:
        """Load ONNX quantized model."""
        try:
            import onnxruntime as ort
        except ImportError:
            raise ImportError(
                "onnxruntime is required for ONNX models. "
                "Install with: pip install onnxruntime"
            )

        # Configure session options for optimal performance
        session_options = ort.SessionOptions()
        session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        session_options.intra_op_num_threads = 4

        # Create inference session
        self._session = ort.InferenceSession(
            str(self.model_path),
            sess_options=session_options,
            providers=["CPUExecutionProvider"]  # Quantized models run on CPU
        )

        # Store input/output names
        self.input_names = [inp.name for inp in self._session.get_inputs()]
        self.output_names = [out.name for out in self._session.get_outputs()]

    def _load_pytorch(self) -> None:
        """Load PyTorch quantized model."""
        try:
            import torch
        except ImportError:
            raise ImportError(
                "torch is required for PyTorch models. "
                "Install with: pip install torch"
            )

        self._model = torch.load(self.model_path, map_location="cpu")
        self._model.eval()

    def _load_tensorflow(self) -> None:
        """Load TensorFlow quantized model."""
        try:
            import tensorflow as tf
        except ImportError:
            raise ImportError(
                "tensorflow is required for TensorFlow models. "
                "Install with: pip install tensorflow"
            )

        self._model = tf.saved_model.load(str(self.model_path))

    def predict(self, input_data: Union[np.ndarray, Dict[str, np.ndarray]]) -> np.ndarray:
        """
        Run inference on input data.

        Args:
            input_data: Input array or dict of arrays for named inputs

        Returns:
            Model output as numpy array
        """
        if self._session:  # ONNX
            if isinstance(input_data, dict):
                return self._session.run(self.output_names, input_data)[0]
            else:
                return self._session.run(
                    self.output_names,
                    {self.input_names[0]: input_data}
                )[0]

        elif self._model:  # PyTorch or TensorFlow
            import torch
            if isinstance(self._model, torch.nn.Module):
                with torch.no_grad():
                    tensor_input = torch.from_numpy(input_data) if isinstance(input_data, np.ndarray) else input_data
                    return self._model(tensor_input).numpy()
            else:  # TensorFlow
                return self._model(input_data).numpy()

        raise RuntimeError("Model not loaded")

    @property
    def model_size_mb(self) -> float:
        """Get model file size in MB."""
        return self.model_path.stat().st_size / (1024 * 1024)

    @property
    def quantization_type(self) -> str:
        """Get quantization type."""
        return self.config.quantization_type

    def __repr__(self) -> str:
        return (
            f"QuantizedModel(path={self.model_path.name}, "
            f"size={self.model_size_mb:.1f}MB, "
            f"type={self.quantization_type})"
        )


def load_quantized_model(
    model_path: Union[str, Path],
    config: Optional[QuantizationConfig] = None
) -> QuantizedModel:
    """
    Load a quantized model.

    Args:
        model_path: Path to the quantized model file
        config: Optional quantization configuration

    Returns:
        QuantizedModel instance

    Example:
        >>> model = load_quantized_model("models/jailbreak_int8.onnx")
        >>> result = model.predict(input_data)
    """
    return QuantizedModel(model_path, config)


__all__ = ["QuantizedModel", "QuantizationConfig", "load_quantized_model"]
