"""
Quantization Engine

Configuration advisor that selects the optimal quantization format from the
compatibility matrix for the given hardware backend, then dispatches to torchao
for the actual quantization operation. TorchBridge adds the matrix selection and
fallback chain; quantization is performed by torchao.
"""

from __future__ import annotations

import copy
import logging
import warnings
from dataclasses import dataclass, field
from typing import Any

import torch
import torch.nn as nn

from torchbridge.core.config import (
    AMDConfig,
    HardwareBackend,
    HardwareConfig,
    NVIDIAConfig,
    TPUConfig,
    TrainiumConfig,
)

from .compatibility import Architecture, QuantizationCompatibilityMatrix
from .formats import QuantizationFormat
from .torchao_integration import TORCHAO_AVAILABLE, TorchAOBackend

logger = logging.getLogger(__name__)


@dataclass
class QuantizationResult:
    """Result from a quantization operation."""
    success: bool
    model: nn.Module | None
    format_applied: QuantizationFormat
    format_requested: QuantizationFormat
    used_fallback: bool = False
    fallback_chain: list[QuantizationFormat] = field(default_factory=list)
    memory_before_mb: float = 0.0
    memory_after_mb: float = 0.0
    warnings: list[str] = field(default_factory=list)
    errors: list[str] = field(default_factory=list)

    @property
    def memory_reduction_pct(self) -> float:
        """Actual memory reduction percentage."""
        if self.memory_before_mb <= 0:
            return 0.0
        return (
            (self.memory_before_mb - self.memory_after_mb)
            / self.memory_before_mb
            * 100
        )

    def to_dict(self) -> dict[str, Any]:
        """Serialize to dictionary."""
        return {
            "success": self.success,
            "format_applied": self.format_applied.value,
            "format_requested": self.format_requested.value,
            "used_fallback": self.used_fallback,
            "fallback_chain": [f.value for f in self.fallback_chain],
            "memory_before_mb": self.memory_before_mb,
            "memory_after_mb": self.memory_after_mb,
            "memory_reduction_pct": self.memory_reduction_pct,
            "warnings": self.warnings,
            "errors": self.errors,
        }


def _model_size_mb(model: nn.Module) -> float:
    """Calculate model size in MB."""
    param_bytes = sum(p.numel() * p.element_size() for p in model.parameters())
    buffer_bytes = sum(b.numel() * b.element_size() for b in model.buffers())
    return (param_bytes + buffer_bytes) / (1024 * 1024)


class QuantizationEngine:
    """Backend-aware quantization engine.

    Detects the current hardware backend and architecture, then selects
    the optimal quantization format or applies a user-specified format
    with automatic fallback.

    Args:
        backend: Optional BaseBackend instance. If ``None``, the engine
            detects the backend from ``HardwareConfig``.
    """

    def __init__(self, backend: Any | None = None) -> None:
        self._backend = backend
        self._hw_backend, self._architecture = self._detect_hardware(backend)

    def quantize(
        self,
        model: nn.Module,
        format: str | QuantizationFormat = "auto",
        calibration_data: Any | None = None,
        in_place: bool = False,
    ) -> QuantizationResult:
        """Quantize a model.

        Args:
            model: PyTorch model to quantize.
            format: Quantization format string or enum. ``"auto"`` selects
                the optimal format for the detected backend.
            calibration_data: Reserved for future use (not currently consumed).
            in_place: If ``False`` (default), quantizes a deep copy.

        Returns:
            QuantizationResult with the quantized model and metadata.
        """
        if not isinstance(model, nn.Module):
            raise TypeError(
                f"model must be an nn.Module, got {type(model).__name__}"
            )

        # Resolve format
        is_auto = False
        if isinstance(format, str):
            if format.lower() == "auto":
                is_auto = True
                requested = QuantizationFormat.NONE  # sentinel for auto
            elif format.lower() == "none":
                requested = QuantizationFormat.NONE
            else:
                requested = QuantizationFormat.from_string(format)
        else:
            requested = format

        # Handle NONE explicitly — return model unchanged
        if not is_auto and requested == QuantizationFormat.NONE:
            memory = _model_size_mb(model)
            return QuantizationResult(
                success=True,
                model=model,
                format_applied=QuantizationFormat.NONE,
                format_requested=QuantizationFormat.NONE,
                memory_before_mb=memory,
                memory_after_mb=memory,
            )

        # Auto-select optimal format
        if is_auto:
            target_format = QuantizationCompatibilityMatrix.get_optimal_format(
                self._hw_backend, self._architecture
            )
            requested_for_result = target_format
        else:
            target_format = requested
            requested_for_result = requested

        # Check support and get fallback chain
        fallback_chain = QuantizationCompatibilityMatrix.get_fallback_chain(
            target_format, self._hw_backend, self._architecture
        )
        used_fallback = fallback_chain[0] != target_format if fallback_chain else False
        actual_format = fallback_chain[0] if fallback_chain else target_format

        if used_fallback:
            msg = (
                f"{target_format.value} not supported on "
                f"{self._hw_backend.value}/{self._architecture}; "
                f"falling back to {actual_format.value}"
            )
            logger.warning(msg)

        # Prepare model
        memory_before = _model_size_mb(model)
        work_model = model if in_place else copy.deepcopy(model)

        result = QuantizationResult(
            success=False,
            model=None,
            format_applied=actual_format,
            format_requested=requested_for_result,
            used_fallback=used_fallback,
            fallback_chain=fallback_chain,
            memory_before_mb=memory_before,
        )

        if used_fallback:
            result.warnings.append(
                f"Requested {target_format.value} unavailable; "
                f"using {actual_format.value}"
            )

        # Apply quantization
        try:
            quantized = self._apply_format(
                work_model, actual_format, calibration_data
            )
            result.success = True
            result.model = quantized
            result.memory_after_mb = _model_size_mb(quantized)
        except Exception as e:
            result.errors.append(str(e))
            result.model = work_model
            result.memory_after_mb = memory_before
            logger.error("Quantization failed: %s", e, exc_info=True)

        return result

    def get_optimal_format(self) -> QuantizationFormat:
        """Return the optimal format for the detected backend."""
        return QuantizationCompatibilityMatrix.get_optimal_format(
            self._hw_backend, self._architecture
        )

    def get_supported_formats(self) -> list[QuantizationFormat]:
        """Return all supported formats for the detected backend."""
        return QuantizationCompatibilityMatrix.get_supported_formats(
            self._hw_backend, self._architecture
        )

    @property
    def backend_name(self) -> str:
        """Human-readable backend name."""
        return self._hw_backend.value

    @property
    def architecture_name(self) -> str:
        """Human-readable architecture name."""
        return str(self._architecture) if self._architecture else "unknown"

    # ── private ───────────────────────────────────────────────────────────

    def _detect_hardware(
        self, backend: Any | None
    ) -> tuple[HardwareBackend, Architecture]:
        """Detect hardware backend and architecture."""
        if backend is not None:
            return self._detect_from_backend_object(backend)

        # Fall back to HardwareConfig auto-detection
        hw = HardwareConfig()
        arch: Architecture = None

        if hw.backend == HardwareBackend.CUDA:
            arch = hw.nvidia.architecture
        elif hw.backend == HardwareBackend.AMD:
            arch = hw.amd.architecture
        elif hw.backend == HardwareBackend.TRAINIUM:
            arch = hw.trainium.architecture
        elif hw.backend == HardwareBackend.TPU:
            arch = hw.tpu.version

        return hw.backend, arch

    def _detect_from_backend_object(
        self, backend: Any
    ) -> tuple[HardwareBackend, Architecture]:
        """Extract backend type and architecture from a BaseBackend instance."""
        name = getattr(backend, "BACKEND_NAME", "cpu").lower()
        backend_map = {
            "nvidia": HardwareBackend.CUDA,
            "cuda": HardwareBackend.CUDA,
            "amd": HardwareBackend.AMD,
            "rocm": HardwareBackend.AMD,
            "trainium": HardwareBackend.TRAINIUM,
            "neuron": HardwareBackend.TRAINIUM,
            "tpu": HardwareBackend.TPU,
            "xla": HardwareBackend.TPU,
        }
        hw_backend = backend_map.get(name, HardwareBackend.CPU)

        # Try to get architecture from backend
        arch: Architecture = None
        if hw_backend == HardwareBackend.CUDA:
            arch = NVIDIAConfig().architecture
        elif hw_backend == HardwareBackend.AMD:
            arch = AMDConfig().architecture
        elif hw_backend == HardwareBackend.TRAINIUM:
            arch = TrainiumConfig().architecture
        elif hw_backend == HardwareBackend.TPU:
            arch = TPUConfig().version

        return hw_backend, arch

    def _torchao_backend_str(self) -> str:
        """Return the torchao backend identifier for the detected hardware."""
        if self._hw_backend == HardwareBackend.AMD:
            return "rocm"
        if self._hw_backend == HardwareBackend.CUDA:
            return "cuda"
        return "cpu"

    def _apply_format(
        self,
        model: nn.Module,
        fmt: QuantizationFormat,
        calibration_data: Any | None,
    ) -> nn.Module:
        """Dispatch to the appropriate quantization implementation."""
        if fmt == QuantizationFormat.NONE:
            return model

        if fmt == QuantizationFormat.BF16:
            return model.to(dtype=torch.bfloat16)

        if fmt == QuantizationFormat.INT8_DYNAMIC:
            return self._apply_int8_dynamic(model)

        if fmt == QuantizationFormat.INT8_DYNAMIC_ACTIVATIONS:
            return self._apply_int8_dynamic_activations(model)

        if fmt == QuantizationFormat.INT4_WEIGHT_ONLY:
            return self._apply_int4_weight_only(model)

        if fmt in (QuantizationFormat.FP8_E4M3, QuantizationFormat.FP8_E5M2):
            return self._apply_fp8(model, fmt)

        if fmt == QuantizationFormat.NVFP4:
            return self._apply_nvfp4(model)

        raise ValueError(f"Unhandled quantization format: {fmt.value}")

    # ── format implementations ────────────────────────────────────────────

    def _apply_int8_dynamic(self, model: nn.Module) -> nn.Module:
        """INT8 dynamic quantization (PyTorch native, no torchao needed)."""
        if TorchAOBackend.is_available_on_backend(self._torchao_backend_str()):
            try:
                return TorchAOBackend.quantize_int8_dynamic(model)
            except Exception as e:
                logger.warning("torchao INT8 failed, using PyTorch native: %s", e)
        elif TORCHAO_AVAILABLE:
            logger.info(
                "torchao not supported on %s backend; using PyTorch native INT8",
                self._hw_backend.value,
            )

        # PyTorch native fallback
        try:
            return torch.quantization.quantize_dynamic(
                model, {nn.Linear}, dtype=torch.qint8
            )
        except RuntimeError as e:
            if "NoQEngine" in str(e) or "quantized" in str(e).lower():
                logger.warning(
                    "PyTorch quantization engine not available: %s. "
                    "Falling back to BF16.",
                    e,
                )
                return model.to(dtype=torch.bfloat16)
            raise

    def _apply_int8_dynamic_activations(self, model: nn.Module) -> nn.Module:
        """INT8 dynamic activation quantization via torchao."""
        if TorchAOBackend.is_available_on_backend(self._torchao_backend_str()):
            return TorchAOBackend.quantize_int8_dynamic_activations(model)
        # Fallback to regular INT8
        warnings.warn(
            "torchao not available for INT8 dynamic activations; falling back to INT8 dynamic",
            stacklevel=2,
        )
        return self._apply_int8_dynamic(model)

    def _apply_int4_weight_only(self, model: nn.Module) -> nn.Module:
        """INT4 weight-only via torchao."""
        if TorchAOBackend.is_available_on_backend(self._torchao_backend_str()):
            return TorchAOBackend.quantize_int4_weight_only(model)
        # Fallback to INT8
        warnings.warn(
            "torchao not available for INT4; falling back to INT8 dynamic",
            stacklevel=2,
        )
        return self._apply_int8_dynamic(model)

    def _apply_fp8(
        self, model: nn.Module, fmt: QuantizationFormat
    ) -> nn.Module:
        """FP8 quantization via torchao. Falls back to INT8 if unavailable."""
        if TorchAOBackend.is_available_on_backend(self._torchao_backend_str()):
            try:
                return TorchAOBackend.quantize_fp8(model)
            except Exception as e:
                logger.debug("TorchAO FP8 failed: %s", e)
        warnings.warn(
            f"FP8 ({fmt.value}) not available on this backend; falling back to INT8 dynamic",
            stacklevel=2,
        )
        return self._apply_int8_dynamic(model)

    def _apply_nvfp4(self, model: nn.Module) -> nn.Module:
        """NVFP4 quantization — falls back to FP8 until native FP4 support lands."""
        warnings.warn(
            "NVFP4 native support not yet available; falling back to FP8 E4M3",
            stacklevel=2,
        )
        return self._apply_fp8(model, QuantizationFormat.FP8_E4M3)

