"""Model quantization utilities."""

from pathlib import Path
from typing import Any, List, Optional, Union

import numpy as np

from .loader import QuantizationConfig


def quantize_model(
    model_path: Union[str, Path],
    output_path: Union[str, Path],
    config: Optional[QuantizationConfig] = None,
    calibration_data: Optional[List[np.ndarray]] = None,
) -> Path:
    """
    Quantize a model to int8 precision.

    This function quantizes ONNX, PyTorch, or TensorFlow models to reduce
    size and improve inference speed with minimal accuracy loss.

    Args:
        model_path: Path to the original model
        output_path: Path to save the quantized model
        config: Quantization configuration
        calibration_data: Optional calibration data for dynamic quantization

    Returns:
        Path to the quantized model

    Example:
        >>> quantize_model(
        ...     "models/jailbreak.onnx",
        ...     "models/jailbreak_int8.onnx",
        ...     config=QuantizationConfig(quantization_type="int8", per_channel=True)
        ... )
    """
    model_path = Path(model_path)
    output_path = Path(output_path)
    config = config or QuantizationConfig()

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

    suffix = model_path.suffix.lower()

    if suffix == ".onnx":
        return _quantize_onnx(model_path, output_path, config, calibration_data)
    elif suffix in [".pt", ".pth"]:
        return _quantize_pytorch(model_path, output_path, config)
    elif suffix in [".pb", ".h5"]:
        return _quantize_tensorflow(model_path, output_path, config)
    else:
        raise ValueError(f"Unsupported model format: {suffix}")


def _quantize_onnx(
    model_path: Path,
    output_path: Path,
    config: QuantizationConfig,
    calibration_data: Optional[List[np.ndarray]] = None,
) -> Path:
    """Quantize ONNX model."""
    try:
        from onnxruntime.quantization import (
            quantize_dynamic,
            quantize_static,
            QuantType,
            CalibrationDataReader,
        )
    except ImportError:
        raise ImportError(
            "onnxruntime is required for ONNX quantization. "
            "Install with: pip install onnxruntime"
        )

    # Determine quantization type
    quant_type_map = {
        "int8": QuantType.QInt8,
        "uint8": QuantType.QUInt8,
    }
    quant_type = quant_type_map.get(config.quantization_type, QuantType.QInt8)

    output_path.parent.mkdir(parents=True, exist_ok=True)

    if calibration_data:
        # Static quantization (more accurate but requires calibration data)
        class DataReader(CalibrationDataReader):
            def __init__(self, data: List[np.ndarray]):
                self.data = data
                self.idx = 0

            def get_next(self):
                if self.idx >= len(self.data):
                    return None
                result = {"input": self.data[self.idx]}
                self.idx += 1
                return result

        quantize_static(
            str(model_path),
            str(output_path),
            DataReader(calibration_data),
            quant_format=quant_type,
            per_channel=config.per_channel,
        )
    else:
        # Dynamic quantization (no calibration data needed)
        quantize_dynamic(
            str(model_path),
            str(output_path),
            weight_type=quant_type,
            per_channel=config.per_channel,
            reduce_range=not config.symmetric,
            optimize_model=config.optimize_for_inference,
        )

    return output_path


def _quantize_pytorch(
    model_path: Path,
    output_path: Path,
    config: QuantizationConfig,
) -> Path:
    """Quantize PyTorch model."""
    try:
        import torch
        from torch.quantization import quantize_dynamic, default_dynamic_qconfig
    except ImportError:
        raise ImportError(
            "torch is required for PyTorch quantization. "
            "Install with: pip install torch"
        )

    # Load model
    model = torch.load(model_path, map_location="cpu")
    model.eval()

    # Apply dynamic quantization
    quantized_model = quantize_dynamic(
        model,
        {torch.nn.Linear, torch.nn.LSTM, torch.nn.GRU},  # Quantize these layers
        dtype=torch.qint8,
    )

    # Save quantized model
    output_path.parent.mkdir(parents=True, exist_ok=True)
    torch.save(quantized_model, output_path)

    return output_path


def _quantize_tensorflow(
    model_path: Path,
    output_path: Path,
    config: QuantizationConfig,
) -> Path:
    """Quantize TensorFlow model."""
    try:
        import tensorflow as tf
    except ImportError:
        raise ImportError(
            "tensorflow is required for TensorFlow quantization. "
            "Install with: pip install tensorflow"
        )

    # Load model
    model = tf.saved_model.load(str(model_path))

    # Convert to TFLite with quantization
    converter = tf.lite.TFLiteConverter.from_saved_model(str(model_path))
    converter.optimizations = [tf.lite.Optimize.DEFAULT]

    if config.quantization_type == "int8":
        converter.target_spec.supported_types = [tf.int8]

    tflite_model = converter.convert()

    # Save quantized model
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_bytes(tflite_model)

    return output_path


def quantize_transformers_model(
    model_name: str,
    output_path: Union[str, Path],
    config: Optional[QuantizationConfig] = None,
) -> Path:
    """
    Quantize a Hugging Face Transformers model.

    Args:
        model_name: Name of the model on Hugging Face Hub
        output_path: Path to save the quantized model
        config: Quantization configuration

    Returns:
        Path to the quantized model

    Example:
        >>> quantize_transformers_model(
        ...     "distilbert-base-uncased",
        ...     "models/distilbert_int8.onnx"
        ... )
    """
    try:
        from transformers import AutoModel, AutoTokenizer
        from optimum.onnxruntime import ORTQuantizer
        from optimum.onnxruntime.configuration import AutoQuantizationConfig
    except ImportError:
        raise ImportError(
            "transformers and optimum are required. "
            "Install with: pip install transformers optimum[onnxruntime]"
        )

    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    config = config or QuantizationConfig()

    # Load model and tokenizer
    model = AutoModel.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    # Export to ONNX
    from transformers.onnx import export
    onnx_path = output_path.with_suffix(".onnx")
    export(
        preprocessor=tokenizer,
        model=model,
        config=model.config,
        opset=14,
        output=onnx_path,
    )

    # Quantize
    quantizer = ORTQuantizer.from_pretrained(model_name)

    qconfig = AutoQuantizationConfig.arm64() if config.optimize_for_inference else AutoQuantizationConfig.avx512()

    quantizer.quantize(
        save_dir=output_path.parent,
        quantization_config=qconfig,
    )

    return output_path


__all__ = ["quantize_model", "quantize_transformers_model"]
