"""
Model utilities — framework detection and automatic extractor/absorber selection.
"""

from __future__ import annotations

from typing import Any, Optional

from model_osmosis.extractors.base import BaseExtractor
from model_osmosis.extractors.pytorch_extractor import PyTorchExtractor
from model_osmosis.extractors.sklearn_extractor import SklearnExtractor
from model_osmosis.extractors.generic_extractor import GenericExtractor
from model_osmosis.absorbers.base import BaseAbsorber
from model_osmosis.absorbers.pytorch_absorber import PyTorchAbsorber
from model_osmosis.absorbers.sklearn_absorber import SklearnAbsorber

# Ordered by specificity — more specific extractors first
_EXTRACTORS = [
    PyTorchExtractor(),
    SklearnExtractor(),
    GenericExtractor(),  # fallback
]

_ABSORBERS = [
    PyTorchAbsorber(),
    SklearnAbsorber(),
]


def detect_framework(model: Any) -> str:
    """Detect which ML framework a model belongs to.

    Returns:
        One of: 'pytorch', 'sklearn', 'generic', 'unknown'.
    """
    try:
        import torch.nn as nn
        if isinstance(model, nn.Module):
            return "pytorch"
    except ImportError:
        pass

    try:
        from sklearn.base import BaseEstimator
        if isinstance(model, BaseEstimator):
            return "sklearn"
    except ImportError:
        pass

    if hasattr(model, "predict") or hasattr(model, "__call__"):
        return "generic"

    return "unknown"


def auto_select_extractor(model: Any) -> BaseExtractor:
    """Automatically select the best extractor for a model.

    Tries framework-specific extractors first, falls back to generic.

    Raises:
        ValueError: If no extractor can handle the model.
    """
    for extractor in _EXTRACTORS:
        if extractor.can_handle(model):
            return extractor
    raise ValueError(
        f"No extractor can handle model of type {type(model).__name__}. "
        f"Model must have at least a predict() or __call__() method."
    )


def auto_select_absorber(model: Any) -> BaseAbsorber:
    """Automatically select the best absorber for a model.

    Raises:
        ValueError: If no absorber can handle the model.
    """
    for absorber in _ABSORBERS:
        if absorber.can_handle(model):
            return absorber
    raise ValueError(
        f"No absorber can handle model of type {type(model).__name__}. "
        f"Supported targets: PyTorch nn.Module, sklearn estimators."
    )
