
import os
import sys

_ORB_AVAILABLE = False
try:
    from orb_models.forcefield import pretrained
    from orb_models.forcefield.calculator import ORBCalculator
    _ORB_AVAILABLE = True
except Exception:
    pass

from macer.defaults import resolve_model_path

def get_orb_calculator(model_path="orb_v2", device="cpu", **kwargs):
    """
    Constructs ORB calculator from a pretrained model name.
    """
    if not _ORB_AVAILABLE:
        raise RuntimeError("orb-models is not installed. Please reinstall with 'pip install macer'.")

    # Fix for macOS: Disable torch.compile to avoid library loading errors (@rpath/libc++.1.dylib)
    if sys.platform == "darwin" and "TORCH_COMPILE_DISABLE" not in os.environ:
        os.environ["TORCH_COMPILE_DISABLE"] = "1"

    # Normalize name for fuzzy matching (ignore hyphens/underscores and case)
    def normalize(s):
        return s.replace("-", "").replace("_", "").lower()

    target = normalize(model_path)
    available_names = dir(pretrained)
    
    # Try to find a match in pretrained models (using full string)
    match = None
    for name in available_names:
        if not name.startswith("_") and normalize(name) == target:
            match = name
            break
            
    # If not found, try using the basename of the path
    # (Because macer md resolves paths to absolute paths like .../orb_v2)
    if not match:
        target_base = normalize(os.path.basename(model_path))
        for name in available_names:
            if not name.startswith("_") and normalize(name) == target_base:
                match = name
                break
    
    if match:
        model_name = match
    else:
        # Try resolving as local path
        resolved = resolve_model_path(model_path)
        if resolved and os.path.exists(resolved):
            model_name = resolved
        else:
            model_name = model_path

    # Final check and loading
    try:
        model_func = getattr(pretrained, model_name)
        print(f"Loading ORB pretrained model: {model_name}")
        orbff = model_func(device=device)
    except AttributeError:
        # One last check: if the original model_path (before resolution) was a match
        # (should have been caught by the loop, but being extra safe)
        if match:
             model_func = getattr(pretrained, match)
             orbff = model_func(device=device)
        else:
             available_models = [name for name in available_names if not name.startswith('_') and callable(getattr(pretrained, name))]
             raise ValueError(f"Unsupported ORB model: '{model_name}'.\nAvailable models are: {available_models}")

    return ORBCalculator(orbff, device=device)
