import os
import sys
import types
import importlib.metadata as importlib_metadata
from pathlib import Path

_FAIRCHEM_AVAILABLE = False
_FAIRCHEM_BACKEND = None
try:
    # fairchem-core -> torchtnt imports pkg_resources. Some envs lack it even with setuptools.
    # Provide a minimal shim so fairchem-core can import.
    try:
        import pkg_resources  # noqa: F401
    except Exception:
        pkg_resources = types.ModuleType("pkg_resources")

        class _Dist:
            def __init__(self, version: str):
                self.version = version

        def get_distribution(name: str):
            return _Dist(importlib_metadata.version(name))

        pkg_resources.get_distribution = get_distribution  # type: ignore[attr-defined]
        sys.modules["pkg_resources"] = pkg_resources

    try:
        from fairchem.core import OCPCalculator
        from fairchem.core.models import available_pretrained_models
        _FAIRCHEM_AVAILABLE = True
        _FAIRCHEM_BACKEND = "ocp"
    except Exception:
        from fairchem.core.calculate.ase_calculator import FAIRChemCalculator
        from fairchem.core.calculate import pretrained_mlip
        available_pretrained_models = getattr(pretrained_mlip, "available_models", ())
        _FAIRCHEM_AVAILABLE = True
        _FAIRCHEM_BACKEND = "fairchem"
except Exception:
    available_pretrained_models = ()


DEFAULT_FAIRCHEM_MODEL = "EquiformerV2-31M-S2EF-OC20-All+MD"

def get_fairchem_calculator(model_path=DEFAULT_FAIRCHEM_MODEL, device="cpu", **kwargs):
    """
    Constructs a FAIRChem OCP calculator using fairchem-core.
    """
    if not _FAIRCHEM_AVAILABLE:
        raise RuntimeError("fairchem-core is not installed. Please reinstall with 'pip install macer'.")

    local_cache = os.path.join(Path.home(), ".macer", "fairchem")
    os.makedirs(local_cache, exist_ok=True)

    cpu = device == "cpu"
    checkpoint_path = None
    model_name = None

    if model_path and os.path.exists(model_path):
        checkpoint_path = model_path
    else:
        model_name = model_path
        if available_pretrained_models and model_name not in available_pretrained_models:
            print(
                f"WARNING: Unknown FAIRChem model name '{model_name}'. "
                f"Falling back to '{DEFAULT_FAIRCHEM_MODEL}'."
            )
            model_name = DEFAULT_FAIRCHEM_MODEL

    try:
        if _FAIRCHEM_BACKEND == "ocp":
            if model_name:
                print(f"Loading FAIRChem model: {model_name}")
                return OCPCalculator(model_name=model_name, local_cache=local_cache, cpu=cpu)
            print(f"Loading FAIRChem checkpoint: {checkpoint_path}")
            return OCPCalculator(checkpoint_path=checkpoint_path, cpu=cpu)

        if _FAIRCHEM_BACKEND == "fairchem":
            if model_name:
                print(f"Loading FAIRChem model: {model_name}")
                return FAIRChemCalculator.from_model_checkpoint(
                    model_name,
                    device="cpu" if cpu else "cuda",
                )
            print(f"Loading FAIRChem checkpoint: {checkpoint_path}")
            return FAIRChemCalculator.from_model_checkpoint(
                checkpoint_path,
                device="cpu" if cpu else "cuda",
            )

        raise RuntimeError("Unknown FAIRChem backend.")
    except Exception as e:
        raise RuntimeError(
            f"Failed to load FAIRChem model '{model_path}'. "
            f"Original error: {e}"
        )
