"""
Adapter Engine

Injects adapter layers into models, manages base model quantization
for QLoRA/QDoRA, and provides merge-for-deployment functionality.
"""

from __future__ import annotations

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

import torch
import torch.nn as nn

from torchbridge.core.config import HardwareBackend
from torchbridge.precision.quantization.formats import QuantizationFormat
from torchbridge.precision.quantization.torchao_integration import (
    TORCHAO_AVAILABLE as _TORCHAO_AVAILABLE,
)

from .compatibility import AdapterCompatibilityMatrix, Architecture
from .config import AdapterConfig, AdapterMethod
from .layers import DoRALinear, LoRALinear, QDoRALinear, QLoRALinear
from .model_families import ModelFamily, detect_model_family, get_target_modules

logger = logging.getLogger(__name__)


@dataclass
class AdapterResult:
    """Result from adapter injection.

    Attributes:
        success: Whether injection completed successfully.
        method_applied: Adapter method actually used.
        method_requested: Adapter method originally requested.
        used_fallback: Whether a fallback method was used.
        modules_adapted: Number of modules replaced with adapters.
        trainable_params: Total trainable parameters after injection.
        total_params: Total parameters (base + adapters).
        trainable_ratio: Ratio of trainable to total parameters.
        base_quantized: Whether the base model was quantized.
        base_quant_format: Quantization format used for base model.
        memory_before_mb: Model memory before injection (MB).
        memory_after_mb: Model memory after injection (MB).
        warnings: Any warnings generated during injection.
    """

    success: bool
    method_applied: AdapterMethod
    method_requested: AdapterMethod
    used_fallback: bool = False
    modules_adapted: int = 0
    trainable_params: int = 0
    total_params: int = 0
    trainable_ratio: float = 0.0
    base_quantized: bool = False
    base_quant_format: QuantizationFormat | None = None
    memory_before_mb: float = 0.0
    memory_after_mb: float = 0.0
    warnings: list[str] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        """Serialize to dictionary."""
        return {
            "success": self.success,
            "method_applied": self.method_applied.value,
            "method_requested": self.method_requested.value,
            "used_fallback": self.used_fallback,
            "modules_adapted": self.modules_adapted,
            "trainable_params": self.trainable_params,
            "total_params": self.total_params,
            "trainable_ratio": round(self.trainable_ratio, 6),
            "base_quantized": self.base_quantized,
            "base_quant_format": (
                self.base_quant_format.value if self.base_quant_format else None
            ),
            "memory_before_mb": round(self.memory_before_mb, 1),
            "memory_after_mb": round(self.memory_after_mb, 1),
            "warnings": self.warnings,
        }


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 AdapterEngine:
    """Inject and manage adapters on any model.

    Handles module targeting, base model quantization for QLoRA/QDoRA,
    backend-aware method selection, and adapter lifecycle management.

    Args:
        config: Adapter configuration.
        backend: Target hardware backend.
        architecture: Specific hardware architecture (optional).
    """

    def __init__(
        self,
        config: AdapterConfig,
        backend: HardwareBackend = HardwareBackend.CPU,
        architecture: Architecture = None,
    ):
        self._config = config
        self._backend = backend
        self._architecture = architecture
        self._resolved_method = self._resolve_method()

    def _resolve_method(self) -> AdapterMethod:
        """Resolve the adapter method, falling back if unsupported."""
        requested = self._config.method
        chain = AdapterCompatibilityMatrix.get_fallback_chain(
            self._backend, self._architecture
        )

        if requested in chain:
            return requested

        fallback = chain[0]
        logger.warning(
            "%s not supported on %s, falling back to %s",
            requested.value,
            self._backend.value,
            fallback.value,
        )
        return fallback

    def inject(self, model: nn.Module) -> AdapterResult:
        """Inject adapters into the model in-place.

        Steps:
        1. Measure pre-injection memory.
        2. Identify target modules by name pattern.
        3. Replace nn.Linear with LoRALinear or DoRALinear.
        4. Freeze all base parameters.
        5. Return result with stats.

        Args:
            model: The model to inject adapters into (modified in-place).

        Returns:
            AdapterResult with injection statistics.
        """
        method = self._resolved_method
        used_fallback = method != self._config.method
        warnings_list: list[str] = []

        if used_fallback:
            warnings_list.append(
                f"Fell back from {self._config.method.value} to {method.value}"
            )

        # Auto-detect target modules based on model family (only when
        # the user hasn't explicitly set target_modules)
        effective_targets = self._config.target_modules
        if (
            self._config.auto_detect_targets
            and not self._config._targets_explicitly_set
        ):
            family = detect_model_family(model)
            if family != ModelFamily.UNKNOWN:
                detected = get_target_modules(model, family)
                if detected != effective_targets:
                    logger.info(
                        "Auto-detected model family %s, using target "
                        "modules: %s",
                        family.value,
                        detected,
                    )
                    effective_targets = detected

        self._effective_targets = effective_targets

        memory_before = _model_size_mb(model)

        # Check for existing adapters (double injection guard)
        existing_adapters = sum(
            1 for _, m in model.named_modules()
            if isinstance(m, (LoRALinear, DoRALinear, QLoRALinear, QDoRALinear))
        )
        if existing_adapters > 0:
            warnings_list.append(
                f"Model already has {existing_adapters} adapter(s). "
                "Call merge() first to remove existing adapters."
            )

        # Identify and replace target modules
        modules_adapted = 0
        has_linear = False
        replacements: list[tuple[nn.Module, str, LoRALinear | DoRALinear]] = []
        _base_quantized = False
        _base_quant_format = None

        for name, module in model.named_modules():
            if not isinstance(module, nn.Linear):
                continue
            has_linear = True
            if not self._is_target_module(name):
                continue
            # Find parent module and attribute name
            parent, attr = self._get_parent_and_attr(model, name)
            if parent is None:
                continue

            adapter_layer = self._create_adapter_layer(module, method)
            replacements.append((parent, attr, adapter_layer))
            modules_adapted += 1
            # Track quantization from the first adapter created
            if modules_adapted == 1:
                if isinstance(adapter_layer, (QLoRALinear, QDoRALinear)):
                    _base_quantized = True
                    _base_quant_format = adapter_layer._quant_format

        # Apply replacements
        for parent, attr, adapter_layer in replacements:
            setattr(parent, attr, adapter_layer)

        if modules_adapted == 0:
            if not has_linear and existing_adapters == 0:
                warnings_list.append(
                    "Model has no nn.Linear modules to adapt"
                )
            elif existing_adapters > 0:
                warnings_list.append(
                    "No new modules adapted (model may already be adapted)"
                )
            else:
                warnings_list.append(
                    f"No modules matched target patterns: {effective_targets}"
                )

        # Count params
        trainable = sum(
            p.numel() for p in model.parameters() if p.requires_grad
        )
        total = sum(p.numel() for p in model.parameters())
        ratio = trainable / total if total > 0 else 0.0

        memory_after = _model_size_mb(model)

        logger.info(
            "Injected %s adapters into %d modules "
            "(trainable: %d / %d = %.4f%%)",
            method.value,
            modules_adapted,
            trainable,
            total,
            ratio * 100,
        )

        return AdapterResult(
            success=modules_adapted > 0,
            method_applied=method,
            method_requested=self._config.method,
            used_fallback=used_fallback,
            modules_adapted=modules_adapted,
            trainable_params=trainable,
            total_params=total,
            trainable_ratio=ratio,
            base_quantized=_base_quantized,
            base_quant_format=_base_quant_format,
            memory_before_mb=memory_before,
            memory_after_mb=memory_after,
            warnings=warnings_list,
        )

    def merge(self, model: nn.Module) -> int:
        """Merge all adapters back into base weights.

        Replaces LoRALinear/DoRALinear with plain nn.Linear containing
        merged weights. Used for deployment (inference with zero overhead).

        Args:
            model: Model with adapter layers to merge.

        Returns:
            Number of modules merged.
        """
        merged_count = 0
        replacements: list[tuple[nn.Module, str, nn.Linear]] = []

        for name, module in model.named_modules():
            if isinstance(module, (LoRALinear, DoRALinear, QLoRALinear, QDoRALinear)):
                parent, attr = self._get_parent_and_attr(model, name)
                if parent is None:
                    continue
                merged_linear = module.merge()
                replacements.append((parent, attr, merged_linear))
                merged_count += 1

        for parent, attr, merged_linear in replacements:
            setattr(parent, attr, merged_linear)

        logger.info("Merged %d adapter modules", merged_count)
        return merged_count

    def get_adapter_params(
        self, model: nn.Module
    ) -> dict[str, torch.Tensor]:
        """Extract only adapter parameters (for saving).

        Returns parameters that belong to LoRALinear/DoRALinear modules
        (lora_A, lora_B, magnitude). Non-adapter trainable params like
        embeddings and heads are excluded.

        Args:
            model: Model with adapter layers.

        Returns:
            Dict of adapter parameter name → tensor.
        """
        # Collect prefixes of all adapter modules
        adapter_prefixes: list[str] = []
        for name, module in model.named_modules():
            if isinstance(module, (LoRALinear, DoRALinear, QLoRALinear, QDoRALinear)):
                adapter_prefixes.append(name + ".")

        adapter_params: dict[str, torch.Tensor] = {}
        for name, param in model.named_parameters():
            if any(name.startswith(prefix) for prefix in adapter_prefixes):
                if param.requires_grad:
                    adapter_params[name] = param.data.clone()
        return adapter_params

    def load_adapter_params(
        self, model: nn.Module, params: dict[str, torch.Tensor]
    ) -> int:
        """Load adapter parameters into an adapted model.

        Args:
            model: Model with adapter layers.
            params: Dict of adapter parameter name → tensor.

        Returns:
            Number of parameters loaded.
        """
        loaded = 0
        model_params = dict(model.named_parameters())
        for name, tensor in params.items():
            if name in model_params and model_params[name].requires_grad:
                model_params[name].data.copy_(tensor)
                loaded += 1
        return loaded

    def get_info(self, model: nn.Module) -> dict[str, Any]:
        """Return adapter injection statistics.

        Args:
            model: Model (with or without adapters).

        Returns:
            Dict with adapter statistics.
        """
        adapter_count = 0
        adapter_types: set[str] = set()
        for module in model.modules():
            if isinstance(module, QLoRALinear):
                adapter_count += 1
                adapter_types.add("qlora")
            elif isinstance(module, QDoRALinear):
                adapter_count += 1
                adapter_types.add("qdora")
            elif isinstance(module, LoRALinear):
                adapter_count += 1
                adapter_types.add("lora")
            elif isinstance(module, DoRALinear):
                adapter_count += 1
                adapter_types.add("dora")

        trainable = sum(
            p.numel() for p in model.parameters() if p.requires_grad
        )
        total = sum(p.numel() for p in model.parameters())

        return {
            "adapter_count": adapter_count,
            "adapter_types": sorted(adapter_types),
            "trainable_params": trainable,
            "total_params": total,
            "trainable_ratio": trainable / total if total > 0 else 0.0,
            "config": self._config.to_dict(),
            "backend": self._backend.value,
            "resolved_method": self._resolved_method.value,
        }

    # ── Internal helpers ─────────────────────────────────────────────────────

    def _is_target_module(self, name: str) -> bool:
        """Check if a module name matches any target pattern."""
        targets = getattr(self, "_effective_targets", self._config.target_modules)
        return any(name.endswith(t) for t in targets)

    def _get_parent_and_attr(
        self, root: nn.Module, target_name: str
    ) -> tuple[nn.Module | None, str]:
        """Find the parent module and attribute name for a given path."""
        parts = target_name.split(".")
        if len(parts) == 1:
            return root, parts[0]

        parent = root
        for part in parts[:-1]:
            if hasattr(parent, part):
                parent = getattr(parent, part)
            else:
                return None, ""
        return parent, parts[-1]

    def _create_adapter_layer(
        self,
        base_linear: nn.Linear,
        method: AdapterMethod,
    ) -> LoRALinear | DoRALinear | QLoRALinear | QDoRALinear:
        """Create the appropriate adapter layer."""
        if method == AdapterMethod.QLORA:
            return self._create_qlora_layer(base_linear)
        if method == AdapterMethod.QDORA:
            return self._create_qdora_layer(base_linear)
        if method == AdapterMethod.DORA:
            return DoRALinear(
                base_linear=base_linear,
                rank=self._config.rank,
                alpha=self._config.alpha,
                dropout=self._config.dropout,
                init_method=self._config.init_method,
            )
        return LoRALinear(
            base_linear=base_linear,
            rank=self._config.rank,
            alpha=self._config.alpha,
            dropout=self._config.dropout,
            init_method=self._config.init_method,
        )

    def _create_qlora_layer(self, base_linear: nn.Linear) -> LoRALinear | QLoRALinear:
        """Create a QLoRALinear, falling back to LoRALinear if torchao unavailable."""
        quant_format = AdapterCompatibilityMatrix.get_base_quant_format(self._backend)
        if quant_format is None or not _TORCHAO_AVAILABLE:
            logger.warning(
                "QLoRA not available on %s (torchao=%s, format=%s); "
                "falling back to LoRA",
                self._backend.value,
                _TORCHAO_AVAILABLE,
                quant_format,
            )
            return LoRALinear(
                base_linear=base_linear,
                rank=self._config.rank,
                alpha=self._config.alpha,
                dropout=self._config.dropout,
                init_method=self._config.init_method,
            )
        return QLoRALinear(
            base_linear=base_linear,
            rank=self._config.rank,
            alpha=self._config.alpha,
            dropout=self._config.dropout,
            init_method=self._config.init_method,
            quant_format=quant_format,
        )

    def _create_qdora_layer(self, base_linear: nn.Linear) -> DoRALinear | QDoRALinear:
        """Create a QDoRALinear, falling back to DoRALinear if torchao unavailable."""
        quant_format = AdapterCompatibilityMatrix.get_base_quant_format(self._backend)
        if quant_format is None or not _TORCHAO_AVAILABLE:
            logger.warning(
                "QDoRA not available on %s (torchao=%s, format=%s); "
                "falling back to DoRA",
                self._backend.value,
                _TORCHAO_AVAILABLE,
                quant_format,
            )
            return DoRALinear(
                base_linear=base_linear,
                rank=self._config.rank,
                alpha=self._config.alpha,
                dropout=self._config.dropout,
                init_method=self._config.init_method,
            )
        return QDoRALinear(
            base_linear=base_linear,
            rank=self._config.rank,
            alpha=self._config.alpha,
            dropout=self._config.dropout,
            init_method=self._config.init_method,
            quant_format=quant_format,
        )
