import os
import sys
from importlib.util import find_spec
from unittest.mock import MagicMock

from macer.defaults import resolve_model_path

def get_m3gnet_calculator(model_path=None, device='cpu', elements=None, **kwargs):
    if find_spec("matgl") is None:
        raise RuntimeError('M3GNet (matgl) is not installed. Please reinstall with "pip install macer".')

    # Ensure DGL backend for M3GNet models before importing matgl
    os.environ.setdefault("MATGL_BACKEND", "DGL")

    # On some systems (e.g. macOS with Torch 2.4+),
    # DGL / Graphbolt might fail to load some libraries.
    # We try to mock dgl.graphbolt if it's the culprit.
    try:
        sys.modules.setdefault('dgl.graphbolt', MagicMock())
    except Exception:
        pass

    import matgl
    # Compatibility shim: older ASE may not expose ExpCellFilter in ase.constraints
    try:
        from ase.constraints import ExpCellFilter  # noqa: F401
    except Exception:
        try:
            from ase.filters import ExpCellFilter  # noqa: F401
            import ase.constraints as _ase_constraints
            _ase_constraints.ExpCellFilter = ExpCellFilter
        except Exception:
            raise RuntimeError(
                "ASE ExpCellFilter not available. Please upgrade ASE to a newer version "
                "or install a version that provides ExpCellFilter."
            )
    from matgl.ext.ase import PESCalculator

    try:
        if hasattr(matgl, "set_backend"):
            matgl.set_backend("DGL")
    except Exception:
        pass

    if model_path is None:
        # Use the known M3GNet PES model that includes Zn (and most elements).
        model_candidates = ["M3GNet-MP-2021.2.8-PES"]
    else:
        # Check if it's a known pretrained name or a local path
        resolved = resolve_model_path(model_path)
        if resolved and resolved != model_path:
             model_path = resolved

    last_err = None
    elements = set(elements or [])
    if model_path is None:
        for name in model_candidates:
            try:
                potential = matgl.load_model(name)
                if elements:
                    model_elements = set(getattr(potential.model, "element_types", []))
                    if not elements.issubset(model_elements):
                        last_err = KeyError(f"Missing elements in model '{name}': {sorted(elements - model_elements)}")
                        continue
                print(f"INFO: Using M3GNet model: {name}")
                return PESCalculator(potential=potential)
            except Exception as e:
                last_err = e
    else:
        try:
            potential = matgl.load_model(model_path)
            if elements:
                model_elements = set(getattr(potential.model, "element_types", []))
                if not elements.issubset(model_elements):
                    missing = sorted(elements - model_elements)
                    raise RuntimeError(
                        f"M3GNet model '{model_path}' does not support elements: {missing}. "
                        "Please choose a different pretrained model via --model."
                    )
            return PESCalculator(potential=potential)
        except Exception as e:
            last_err = e

    # Final fallback: try torch.hub universal potential
    try:
        import torch
        potential = torch.hub.load("materialsvirtuallab/matgl", "m3gnet_universal_potential")
        return PESCalculator(potential=potential)
    except Exception as e:
        last_err = e

    raise RuntimeError(
        "Failed to load M3GNet model. Ensure DGL backend is set "
        "and clear any stale cache via `python -c 'import matgl; matgl.clear_cache()'`. "
        "If you are on matgl>=2.x, the default pretrained model names may have changed. "
        "Consider passing a valid model name via `--model` or setting MATGL_BACKEND=DGL. "
        f"Original error: {last_err}"
    )
