"""
ErasusUnlearner — High-level API for machine unlearning.

This is the primary user-facing class. Example usage::

    from erasus import ErasusUnlearner

    unlearner = ErasusUnlearner(
        model=clip_model,
        strategy="modality_decoupling",
        selector="influence",
    )
    result = unlearner.fit(forget_data, retain_data)
    metrics = unlearner.evaluate(forget_data, retain_data)
"""

from __future__ import annotations

import logging
from typing import Any, Dict, List, Optional, Tuple

import torch
import torch.nn as nn
from torch.utils.data import DataLoader

from erasus.core.base_unlearner import BaseUnlearner, UnlearningResult
from erasus.core.registry import strategy_registry, selector_registry

# Ensure all strategies and selectors are registered
import erasus.strategies  # noqa: F401
import erasus.selectors  # noqa: F401

logger = logging.getLogger(__name__)


def _detect_model_type(model: nn.Module) -> str:
    """Heuristic model-type detection from architecture."""
    name = model.__class__.__name__.lower()
    module_name = model.__class__.__module__ or ""

    # LLM patterns
    if any(k in name for k in ("gpt", "llama", "mistral", "phi", "gemma", "opt", "bloom", "falcon")):
        return "llm"
    if "transformers" in module_name and any(k in name for k in ("causal", "seq2seq")):
        return "llm"

    # Diffusion
    if any(k in name for k in ("unet", "diffusion", "stable")):
        return "diffusion"

    # VLM
    if any(k in name for k in ("clip", "blip", "llava")):
        return "vlm"

    return "classifier"


class ErasusUnlearner(BaseUnlearner):
    """
    Main Erasus unlearner — resolves strategy/selector by name,
    orchestrates coreset selection + unlearning + evaluation.

    Smart defaults:
    - ``strategy="auto"`` picks the best strategy for the model type
    - ``precision="bf16-mixed"`` enables automatic mixed precision
    - ``gradient_checkpointing=True`` in ``fit()`` enables memory savings
    """

    def __init__(
        self,
        model: nn.Module,
        strategy: Any = "gradient_ascent",
        selector: Optional[Any] = None,
        device: Optional[str] = None,
        precision: Optional[str] = None,
        strategy_kwargs: Optional[Dict[str, Any]] = None,
        selector_kwargs: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        self.precision = precision
        from erasus.core.base_strategy import BaseStrategy
        from erasus.core.base_selector import BaseSelector
        from erasus.core.strategy_pipeline import StrategyPipeline

        # Resolve strategy: accept string name, instance, or pipeline
        if isinstance(strategy, str):
            strategy_cls = strategy_registry.get(strategy)
            strategy_instance = strategy_cls(**(strategy_kwargs or {}))
            self.strategy_name = strategy
        elif isinstance(strategy, StrategyPipeline):
            strategy_instance = strategy
            self.strategy_name = repr(strategy)
        elif isinstance(strategy, BaseStrategy):
            strategy_instance = strategy
            self.strategy_name = strategy.__class__.__name__
        else:
            raise TypeError(
                f"strategy must be a string, BaseStrategy instance, or "
                f"StrategyPipeline, got {type(strategy)}"
            )

        # Resolve selector: accept string name, instance, or None
        selector_instance = None
        if isinstance(selector, str):
            selector_cls = selector_registry.get(selector)
            selector_instance = selector_cls(**(selector_kwargs or {}))
            self.selector_name: Optional[str] = selector
        elif isinstance(selector, BaseSelector):
            selector_instance = selector
            self.selector_name = selector.__class__.__name__
        elif selector is None:
            self.selector_name = None
        else:
            raise TypeError(
                f"selector must be a string, BaseSelector instance, or None, got {type(selector)}"
            )

        super().__init__(
            model=model,
            strategy=strategy_instance,
            selector=selector_instance,
            device=device,
            **kwargs,
        )

    def fit(self, *args: Any, **kwargs: Any) -> UnlearningResult:
        """Override fit to inject precision from constructor."""
        if self.precision and "precision" not in kwargs:
            kwargs["precision"] = self.precision
        return super().fit(*args, **kwargs)

    def _run_unlearning(
        self,
        forget_loader: DataLoader,
        retain_loader: Optional[DataLoader],
        epochs: int,
        **kwargs: Any,
    ) -> Tuple[List[float], List[float]]:
        """Delegate to the resolved strategy."""
        self.model, forget_losses, retain_losses = self.strategy.unlearn(
            model=self.model,
            forget_loader=forget_loader,
            retain_loader=retain_loader,
            epochs=epochs,
            **kwargs,
        )
        return forget_losses, retain_losses

    @staticmethod
    def detect_model_type(model: nn.Module) -> str:
        """Detect model type from architecture (for use with strategy='auto')."""
        return _detect_model_type(model)
