"""Shared utilities for uma-ase workflows."""

from __future__ import annotations

import argparse
from contextlib import suppress
import json
from functools import lru_cache
from importlib import resources
import logging
import os
import re
import shutil
import tempfile
import sys
from dataclasses import dataclass, field
from difflib import get_close_matches
from pathlib import Path
from typing import Any, Iterable, List, Mapping, Optional

import numpy as np
from ase import Atoms

# Lazy FairChem imports keep the web server process light and avoid importing
# torch elastic/distributed helpers until a real UMA calculation path needs it.
pretrained_mlip = None
FAIRChemCalculator = None
load_predict_unit = None
FAIR_CHEM_IMPORT_ERROR: Optional[Exception] = None
_FAIR_CHEM_IMPORT_ATTEMPTED = False


def ensure_fairchem_loaded() -> tuple[Any, Any, Any]:
    """Import FairChem lazily and cache the resolved symbols."""
    global pretrained_mlip
    global FAIRChemCalculator
    global load_predict_unit
    global FAIR_CHEM_IMPORT_ERROR
    global _FAIR_CHEM_IMPORT_ATTEMPTED

    if _FAIR_CHEM_IMPORT_ATTEMPTED:
        return pretrained_mlip, FAIRChemCalculator, load_predict_unit

    _FAIR_CHEM_IMPORT_ATTEMPTED = True

    try:
        from fairchem.core import FAIRChemCalculator as _FAIRChemCalculator
        from fairchem.core import pretrained_mlip as _pretrained_mlip
    except ModuleNotFoundError as exc:  # pragma: no cover - environment dependent
        FAIR_CHEM_IMPORT_ERROR = exc
        pretrained_mlip = None
        FAIRChemCalculator = None
        load_predict_unit = None
        return pretrained_mlip, FAIRChemCalculator, load_predict_unit

    pretrained_mlip = _pretrained_mlip
    FAIRChemCalculator = _FAIRChemCalculator
    FAIR_CHEM_IMPORT_ERROR = None

    try:  # FairChem >=1.0
        from fairchem.core.units.mlip_unit import load_predict_unit as _load_predict_unit
    except ModuleNotFoundError:  # pragma: no cover - fallback for old versions
        _load_predict_unit = None

    load_predict_unit = _load_predict_unit
    return pretrained_mlip, FAIRChemCalculator, load_predict_unit


def get_fairchem_import_error() -> Optional[Exception]:
    """Return the cached FairChem import error after attempting a lazy import."""
    ensure_fairchem_loaded()
    return FAIR_CHEM_IMPORT_ERROR


@dataclass
class ConformerEntry:
    """One structure (xyz/cif/etc.) within a conformer set."""

    isomer_id: str
    conf_id: str
    title: Optional[str] = None
    energy_uff: Optional[float] = None
    content: Optional[str] = None
    fmt: str = "xyz"


@dataclass
class Conformers:
    """Collection of conformers for a single connectivity."""

    source: Optional[str] = None
    cid: Optional[str] = None
    iupac_name: Optional[str] = None
    conformers: List[ConformerEntry] = field(default_factory=list)
    base_mol: Optional[Any] = None
    active_conf_id: Optional[str] = None


@dataclass
class Isomers:
    """Collection of isomers (different connectivity, same stoichiometry)."""

    formula: Optional[str] = None
    isomers: List[Conformers] = field(default_factory=list)
    active_isomer_id: Optional[str] = None


GROUND_STATE_MULTIPLICITY_RESOURCE = "data/ground_state_atomic_multiplicities.json"


@lru_cache(maxsize=1)
def load_ground_state_atomic_multiplicity_data() -> dict[str, Any]:
    """Load the packaged neutral-atom ground-state multiplicity table."""
    resource = resources.files("uma_ase").joinpath(GROUND_STATE_MULTIPLICITY_RESOURCE)
    with resource.open("r", encoding="utf-8") as handle:
        payload = json.load(handle)
    if not isinstance(payload, dict):
        raise ValueError("Invalid ground-state multiplicity payload.")
    elements = payload.get("elements")
    if not isinstance(elements, dict):
        raise ValueError("Ground-state multiplicity table is missing 'elements'.")
    return payload


def get_ground_state_atomic_multiplicity_source() -> dict[str, str]:
    """Return the provenance metadata for the packaged multiplicity table."""
    payload = load_ground_state_atomic_multiplicity_data()
    source = payload.get("source")
    if not isinstance(source, dict):
        return {}
    return {str(key): str(value) for key, value in source.items()}


def get_ground_state_atomic_multiplicity(symbol: str) -> int:
    """Return the neutral isolated-atom ground-state multiplicity for *symbol*."""
    normalized = str(symbol or "").strip().capitalize()
    if not normalized:
        return 1
    payload = load_ground_state_atomic_multiplicity_data()
    elements = payload.get("elements", {})
    value = elements.get(normalized)
    try:
        parsed = int(value)
    except (TypeError, ValueError):
        return 1
    return max(1, parsed)


def Sum_of_atomic_energies(
    counts: Mapping[str, int],
    uma_checkpoint: str,
    uma_model: str,
    device: str,
    logger: Optional[logging.Logger] = None,
    predictor: Optional[Any] = None,
) -> float:
    """
    Compute the total energy by summing the energy of each atom type multiplied
    by its number of occurrences.

    Parameters
    ----------
    counts : dict[str, int]
        Dictionary mapping atom symbol -> number of occurrences.
    uma_checkpoint : str
        UMA checkpoint identifier.
    uma_model : str
        UMA task/model identifier.
    device : str
        Execution device (e.g. 'cpu', 'cuda').
    logger : logging.Logger, optional
        Logger used to report per-element energies.

    Returns
    -------
    float
        Total summed energy.
    """
    _, fairchem_calculator_cls, _ = ensure_fairchem_loaded()
    if pretrained_mlip is None or FAIRChemCalculator is None:
        raise RuntimeError("UMA FairChem libraries are unavailable; cannot compute atomic energies.")

    if predictor is None:
        try:
            predictor = get_predict_unit_with_local_fallback(
                uma_checkpoint,
                device=device,
                logger=logger,
            )
        except FileNotFoundError as exc:
            raise FileNotFoundError(str(exc)) from exc
        except KeyError as exc:
            raise KeyError(f"UMA checkpoint '{uma_checkpoint}' not found.") from exc

    sum_atomic_energies = 0.0
    # Reuse a single single-atom system and calculator across all atom types.
    # Only the species changes between iterations.
    a = 5.0
    atom = Atoms("H", positions=[[0, 0, 0]], cell=[a, a, a], pbc=True)
    atom.info["charge"] = 0
    atom.calc = fairchem_calculator_cls(predictor, task_name=uma_model)

    for atom_type, count in counts.items():
        atom[0].symbol = atom_type
        multiplicity = get_ground_state_atomic_multiplicity(atom_type)
        atom.info["spin"] = multiplicity
        energy = atom.get_total_energy()
        label = f"Atomic energy ({atom_type}, mult={multiplicity})"
        message = f"  {label:<28}: {energy:.12f} eV"
        if logger:
            logger.info(message)
        else:
            print(message)
        sum_atomic_energies += energy * count
    return sum_atomic_energies


def sum_of_atomic_energies(
    counts: Mapping[str, int],
    uma_checkpoint: str,
    uma_model: str,
    device: str,
    logger: Optional[logging.Logger] = None,
    predictor: Optional[Any] = None,
) -> float:
    """PEP 8 alias for :func:`Sum_of_atomic_energies`."""
    return Sum_of_atomic_energies(
        counts,
        uma_checkpoint,
        uma_model,
        device,
        logger=logger,
        predictor=predictor,
    )


class Tee:
    """
    A file-like object that writes simultaneously to a file and stdout.

    Usage:
        from log_utils import Tee
        tee = Tee("output.log")
        print("Hello", file=tee)
        tee.close()
    """
    def __init__(self, filename: str, mode: str = "w", stream=sys.stdout):
        """
        Args:
            filename (str): Path to the log file.
            mode (str): File open mode ('w' to overwrite, 'a' to append).
            stream: Output stream to duplicate (default: sys.stdout).
        """
        self.file = open(filename, mode)
        self.stream = stream

    def write(self, message: str) -> None:
        self.stream.write(message)
        self.file.write(message)

    def flush(self) -> None:
        self.stream.flush()
        self.file.flush()

    def close(self) -> None:
        """Close the underlying file."""
        self.file.close()


def rmsd_simple(atoms1: Atoms, atoms2: Atoms) -> float:
    """Compute RMSD between two ASE Atoms objects without alignment."""
    pos1 = atoms1.get_positions()
    pos2 = atoms2.get_positions()
    diff = pos1 - pos2
    return np.sqrt((diff ** 2).sum() / len(pos1))


def get_filename_only(filepath: str) -> str:
    """
    Given a file path like 'folder/filename.ext', return only 'filename'.

    Args:
        filepath (str): The full file name or path.

    Returns:
        str: The filename without its extension.
    """
    # Extract the base name (remove path)
    base = os.path.basename(filepath)
    # Split into (filename, extension)
    name, _ = os.path.splitext(base)
    return name


def delete_file_if_exists(filepath: str) -> bool:
    """
    Check if a file exists, and delete it if it does.

    Args:
        filepath (str): Path to the file to delete.

    Returns:
        bool: True if file existed and was deleted, False otherwise.
    """
    try:
        if os.path.isfile(filepath):
            os.remove(filepath)
            print(f"Deleted file: {filepath}")
            return True
        else:
            print(f"File not found: {filepath}")
            return False
    except PermissionError:
        print(f"Permission denied: {filepath}")
        return False
    except Exception as exc:
        print(f"Error deleting {filepath}: {exc}")
        return False


class ChargeAction(argparse.Action):
    """Track when the user explicitly provides -chg."""

    def __call__(self, parser, namespace, values, option_string=None):
        setattr(namespace, self.dest, values)
        setattr(namespace, "_chg_explicit", True)


class SpinAction(argparse.Action):
    """Track when the user explicitly provides -spin."""

    def __call__(self, parser, namespace, values, option_string=None):
        setattr(namespace, self.dest, values)
        setattr(namespace, "_spin_explicit", True)


@dataclass(frozen=True)
class XYZMetadata:
    comment: Optional[str]
    url: Optional[str]
    charge: Optional[int]
    spin: Optional[int]


_URL_PATTERN = re.compile(r"https?://\S+")
_SIGNED_INT_PATTERN = re.compile(r"(?<![\w.-])[+-]?\d+(?![\w.-])")
_CHARGE_PATTERN = re.compile(r"(?:^|\b)(?:charge|chg|q)\s*[:=\s]+([+-]?\d+)", re.IGNORECASE)
_SPIN_PATTERN = re.compile(r"(?:^|\b)(?:spin|multiplicity|mult)\s*[:=\s]+([+-]?\d+)", re.IGNORECASE)

MD_ENGINE_CHOICES = (
    "langevin",
    "velocity_verlet",
    "nvt_berendsen",
    "nose_hoover_chain",
    "bussi",
    "andersen",
    "langevin_baoab_nvt",
    "npt_berendsen",
    "isotropic_mtk_npt",
    "mtk_npt",
    "langevin_baoab_npt",
    "melchionna_npt",
)


def parse_cli_args(argv: Optional[Iterable[str]] = None) -> argparse.Namespace:
    """
    Parse command-line arguments shared across UMA utilities.

    Required arguments
    ------------------
    - input : Path to the input geometry readable by ASE (or a .use JSON recipe).

    Optional arguments
    ------------------
    - chg             : Total molecular charge (default 0; inferred from XYZ comment when present).
    - spin            : Spin multiplicity (default 1).
    - iter             : Maximum number of optimisation cycles (default 250).
    - grad             : Convergence threshold for maximum force (default 0.01).
    - mlff-chk         : UMA checkpoint identifier (default 'uma-s-1p1').
    - mlff-task        : UMA task/model identifier (default 'omol').
    - visualize        : Open the final trajectory in an interactive viewer.
    - optimizer        : ASE optimizer to use (default 'BFGS').
    - run-type         : Workflow to execute ('sp', 'geoopt', 'freqs', 'ts', 'md'; default 'sp').
    - cpu              : Force CPU execution even if CUDA is available.
    - temp             : Thermochemistry temperature in Kelvin (default 298.15).
    - press            : Thermochemistry pressure in Pascals (default 101325.0).
    - md-engine        : MD integrator/ensemble (default 'langevin').
    - md-steps         : MD steps (default 500).
    - md-timestep-fs   : MD timestep in femtoseconds (default 0.5).
    - md-temperature   : MD target temperature in Kelvin (default 300).
    - md-friction      : Langevin friction in 1/fs (default 0.002).
    - md-traj-interval : Trajectory sampling interval in steps (default 10).
    - md-log-interval  : Log sampling interval in steps (default 10).
    - md-relax-fs      : Thermostat relaxation time in fs (default 100).
    - md-pressure-bar  : Target pressure in bar (default 1.0).
    - md-compressibility-bar : Compressibility in 1/bar (default 4.57e-5).
    - md-barostat-relax-fs   : Barostat relaxation time in fs (default 1000).
    - md-nose-tdamp-fs : Nose-Hoover damping time in fs (default 100).
    - md-nose-tchain   : Nose-Hoover chain length (default 3).
    - md-nose-tloop    : Nose-Hoover sub-steps (default 1).
    - md-bussi-taut-fs : Bussi time constant in fs (default 100).
    - md-andersen-prob : Andersen collision probability (default 0.001).
    - md-mtk-tdamp-fs  : MTK thermostat damping (default 100).
    - md-mtk-pdamp-fs  : MTK barostat damping (default 1000).
    - md-mtk-tchain    : MTK thermostat chain length (default 3).
    - md-mtk-pchain    : MTK barostat chain length (default 3).
    - md-mtk-tloop     : MTK thermostat sub-steps (default 1).
    - md-mtk-ploop     : MTK barostat sub-steps (default 1).
    - md-baoab-ttau-fs : BAOAB thermostat time in fs (default 50).
    - md-baoab-ptau-fs : BAOAB barostat time in fs (default 1000).
    - md-baoab-pmass   : BAOAB cell mass (optional).
    - md-baoab-pmass-factor : BAOAB mass factor (default 1.0).
    - md-baoab-hydrostatic : Use hydrostatic BAOAB cell (flag).
    - md-melchionna-ttime-fs : Melchionna thermostat time (default 25).
    - md-melchionna-ptime-fs : Melchionna barostat time (default 75).
    - md-melchionna-bulk-gpa : Melchionna bulk modulus in GPa (default 100).
    - md-melchionna-mask     : Melchionna mask vector.
    - md-use-pbc      : Apply periodic boundary conditions (flag).
    - md-cell-a/b/c   : Unit cell lengths in Å.
    - md-cell-alpha/beta/gamma : Unit cell angles in degrees.
    - md-restart-traj : Trajectory file to restart from (uses last-frame velocities).
    """
    parser = argparse.ArgumentParser(
        prog="uma-ase",
        description=(
            "Provide the required -input value. Charge defaults to 0 or to the "
            "signed integer embedded in the XYZ comment line."
        ),
    )
    parser.set_defaults(_chg_explicit=False, _spin_explicit=False)

    # --- Required arguments ---
    parser.add_argument(
        "-input",
        type=str,
        required=True,
        help="Input geometry readable by ASE or a .use JSON recipe file (XYZ comment may append the signed charge & spin).",
    )

    parser.add_argument(
        "-chg",
        type=int,
        action=ChargeAction,
        default=0,
        help="Charge (default 0; inferred from XYZ comment when present).",
    )

    # --- Optional arguments ---
    parser.add_argument(
        "-spin",
        type=int,
        action=SpinAction,
        default=1,
        help="Spin multiplicity (default 1; inferred from XYZ comment when present)"
    )

    parser.add_argument(
        "-run-type",
        nargs="+",
        choices=("sp", "geoopt", "freqs", "ts", "md"),
        default=["sp"],
        help="Run type(s) to execute: 'sp', 'geoopt', 'freqs', 'ts', 'md' or any sequence thereof (default: sp).",
    )
    parser.add_argument(
        "-cpu",
        action="store_true",
        help="Force UMA calculations to run on CPU even when CUDA is available.",
    )
    parser.add_argument(
        "-iter",
        type=int,
        default=250,
        help="Max number of geometry optimization cycles. Default=250"
    )

    parser.add_argument(
        "-grad",
        type=float,
        default=0.01,
        help="Max grad for convergence. Default=0.01 eV/A"
    )

    parser.add_argument(
        "-optimizer",
        default="LBFGS",
        metavar="NAME",
        help="ASE optimizer (e.g. BFGS, LBFGS, FIRE, BFGSLineSearch,MDMin). Default='LBFGS'."
    )
    parser.add_argument(
        "--ts-displacement",
        type=float,
        default=0.1,
        help="Initial displacement (Å) applied along the selected mode before the dimer TS search.",
    )
    parser.add_argument(
        "--ts-dimer-separation",
        type=float,
        default=0.01,
        help="Separation (Å) between the two images in the dimer search.",
    )
    parser.add_argument(
        "--ts-trial-step",
        type=float,
        default=0.05,
        help="Trial translation step (Å) used by the dimer optimizer when curvature is negative.",
    )
    parser.add_argument(
        "--ts-max-step",
        type=float,
        default=0.2,
        help="Maximum translation step (Å) allowed per dimer iteration.",
    )
    parser.add_argument(
        "--ts-max-rot",
        type=int,
        default=5,
        help="Maximum number of rotational updates per dimer iteration.",
    )

    parser.add_argument(
        "-mlff-chk",
        default="uma-s-1p1",
        metavar="CHECKPOINT",
        help="UMA checkpoint. Default='uma-s-1p1'."
    )

    parser.add_argument(
        "-mlff-task",
        default="omol",
        metavar="TASK",
        help="UMA task/model. Default='omol'."
    )

    parser.add_argument(
        "-temp",
        "--temperature",
        dest="temp",
        type=float,
        metavar="T",
        default=298.15,
        help="Temperature in Kelvin for vibrational analysis (default 298.15 K).",
    )

    parser.add_argument(
        "-press",
        "--pressure",
        dest="press",
        type=float,
        metavar="P",
        default=101325.0,
        help="Pressure in Pascals for vibrational analysis (default 1 atm=101325.0 Pa).",
    )

    parser.add_argument(
        "-visualize",
        action="store_true",
        help="Open the trajectory of a geoopt run in the ASE interactive viewer."
    )

    # --- MD arguments ---
    parser.add_argument(
        "--md-engine",
        default="langevin",
        choices=MD_ENGINE_CHOICES,
        help="MD integrator/ensemble (default: langevin).",
    )
    parser.add_argument(
        "--md-steps",
        type=int,
        default=500,
        help="MD steps (default: 500).",
    )
    parser.add_argument(
        "--md-timestep-fs",
        type=float,
        default=0.5,
        help="MD timestep in femtoseconds (default: 0.5).",
    )
    parser.add_argument(
        "--md-temperature",
        type=float,
        default=300.0,
        help="MD target temperature in Kelvin (default: 300).",
    )
    parser.add_argument(
        "--md-friction",
        type=float,
        default=0.002,
        help="Langevin friction in 1/fs (default: 0.002).",
    )
    parser.add_argument(
        "--md-traj-interval",
        type=int,
        default=10,
        help="MD trajectory sampling interval in steps (default: 10).",
    )
    parser.add_argument(
        "--md-log-interval",
        type=int,
        default=10,
        help="MD log sampling interval in steps (default: 10).",
    )
    parser.add_argument(
        "--md-relax-fs",
        type=float,
        default=100.0,
        help="Thermostat relaxation time in fs (default: 100).",
    )
    parser.add_argument(
        "--md-pressure-bar",
        type=float,
        default=1.0,
        help="Target pressure in bar (default: 1.0).",
    )
    parser.add_argument(
        "--md-compressibility-bar",
        type=float,
        default=4.57e-5,
        help="Compressibility in 1/bar (default: 4.57e-5).",
    )
    parser.add_argument(
        "--md-barostat-relax-fs",
        type=float,
        default=1000.0,
        help="Barostat relaxation time in fs (default: 1000).",
    )
    parser.add_argument(
        "--md-nose-tdamp-fs",
        type=float,
        default=100.0,
        help="Nose-Hoover damping time in fs (default: 100).",
    )
    parser.add_argument(
        "--md-nose-tchain",
        type=int,
        default=3,
        help="Nose-Hoover chain length (default: 3).",
    )
    parser.add_argument(
        "--md-nose-tloop",
        type=int,
        default=1,
        help="Nose-Hoover sub-steps (default: 1).",
    )
    parser.add_argument(
        "--md-bussi-taut-fs",
        type=float,
        default=100.0,
        help="Bussi time constant in fs (default: 100).",
    )
    parser.add_argument(
        "--md-andersen-prob",
        type=float,
        default=0.001,
        help="Andersen collision probability (default: 0.001).",
    )
    parser.add_argument(
        "--md-mtk-tdamp-fs",
        type=float,
        default=100.0,
        help="MTK thermostat damping time in fs (default: 100).",
    )
    parser.add_argument(
        "--md-mtk-pdamp-fs",
        type=float,
        default=1000.0,
        help="MTK barostat damping time in fs (default: 1000).",
    )
    parser.add_argument(
        "--md-mtk-tchain",
        type=int,
        default=3,
        help="MTK thermostat chain length (default: 3).",
    )
    parser.add_argument(
        "--md-mtk-pchain",
        type=int,
        default=3,
        help="MTK barostat chain length (default: 3).",
    )
    parser.add_argument(
        "--md-mtk-tloop",
        type=int,
        default=1,
        help="MTK thermostat sub-steps (default: 1).",
    )
    parser.add_argument(
        "--md-mtk-ploop",
        type=int,
        default=1,
        help="MTK barostat sub-steps (default: 1).",
    )
    parser.add_argument(
        "--md-baoab-ttau-fs",
        type=float,
        default=50.0,
        help="BAOAB thermostat time in fs (default: 50).",
    )
    parser.add_argument(
        "--md-baoab-ptau-fs",
        type=float,
        default=1000.0,
        help="BAOAB barostat time in fs (default: 1000).",
    )
    parser.add_argument(
        "--md-baoab-pmass",
        type=float,
        default=None,
        help="BAOAB cell mass (optional).",
    )
    parser.add_argument(
        "--md-baoab-pmass-factor",
        type=float,
        default=1.0,
        help="BAOAB mass factor (default: 1.0).",
    )
    parser.add_argument(
        "--md-baoab-hydrostatic",
        action="store_true",
        help="Use hydrostatic BAOAB cell (flag).",
    )
    parser.add_argument(
        "--md-melchionna-ttime-fs",
        type=float,
        default=25.0,
        help="Melchionna thermostat time in fs (default: 25).",
    )
    parser.add_argument(
        "--md-melchionna-ptime-fs",
        type=float,
        default=75.0,
        help="Melchionna barostat time in fs (default: 75).",
    )
    parser.add_argument(
        "--md-melchionna-bulk-gpa",
        type=float,
        default=100.0,
        help="Melchionna bulk modulus in GPa (default: 100).",
    )
    parser.add_argument(
        "--md-melchionna-mask",
        type=str,
        default="",
        help="Melchionna mask vector (comma-separated).",
    )
    parser.add_argument(
        "--md-use-pbc",
        action="store_true",
        help="Apply periodic boundary conditions.",
    )
    parser.add_argument(
        "--md-cell-a",
        type=float,
        default=None,
        help="Unit cell length a in Å.",
    )
    parser.add_argument(
        "--md-cell-b",
        type=float,
        default=None,
        help="Unit cell length b in Å.",
    )
    parser.add_argument(
        "--md-cell-c",
        type=float,
        default=None,
        help="Unit cell length c in Å.",
    )
    parser.add_argument(
        "--md-cell-alpha",
        type=float,
        default=None,
        help="Unit cell angle alpha in degrees.",
    )
    parser.add_argument(
        "--md-cell-beta",
        type=float,
        default=None,
        help="Unit cell angle beta in degrees.",
    )
    parser.add_argument(
        "--md-cell-gamma",
        type=float,
        default=None,
        help="Unit cell angle gamma in degrees.",
    )
    parser.add_argument(
        "--md-restart-traj",
        type=str,
        default=None,
        help="Trajectory file to restart from (uses last-frame velocities).",
    )
    args = parser.parse_args(argv)
    # Provide backward-compatible aliases for temperature/pressure fields.
    if not hasattr(args, "temp"):
        setattr(args, "temp", 298.15)
    if not hasattr(args, "press"):
        setattr(args, "press", 101325.0)
    setattr(args, "temperature", args.temp)
    setattr(args, "pressure", args.press)
    return args


def _resolve_custom_checkpoint_dir() -> Path:
    """
    Resolve the directory housing custom UMA checkpoints.
    """
    custom_dir = os.environ.get("UMA_CUSTOM_CHECKPOINT_DIR")
    if custom_dir:
        return Path(custom_dir).expanduser().resolve()
    return Path.home() / ".uma_ase" / "checkpoints"


def _locate_local_checkpoint(custom_root: Path, checkpoint: str) -> Optional[Path]:
    """
    Locate a checkpoint file under *custom_root* matching *checkpoint*.
    """
    expanded_root = custom_root.expanduser()
    if not expanded_root.exists():
        return None

    candidates = [
        expanded_root / checkpoint,
        expanded_root / f"{checkpoint}.pt",
    ]

    checkpoint_dir = expanded_root / checkpoint
    if checkpoint_dir.is_dir():
        candidates.extend(sorted(checkpoint_dir.glob("*.pt")))

    candidates.extend(sorted(expanded_root.glob(f"{checkpoint}*.pt")))

    seen: set[Path] = set()
    for candidate in candidates:
        candidate = candidate.expanduser()
        if candidate in seen:
            continue
        seen.add(candidate)
        if candidate.is_file():
            return candidate
    return None


def extract_xyz_comment(path: Path) -> Optional[str]:
    """
    Return the second-line comment from an XYZ file, if present.

    Parameters
    ----------
    path : Path
        Path to the candidate XYZ file.

    Returns
    -------
    str | None
        The stripped second line when available, otherwise ``None``.
    """
    metadata = extract_xyz_metadata(path)
    return metadata.comment


def extract_xyz_metadata(path: Path) -> XYZMetadata:
    """
    Extract comment, URL, charge, and spin metadata from an XYZ geometry.
    """
    default = XYZMetadata(comment=None, url=None, charge=None, spin=None)
    if path.suffix.lower() != ".xyz":
        return default

    try:
        with path.open("r", encoding="utf-8", errors="ignore") as handle:
            handle.readline()
            comment = handle.readline()
    except OSError:
        return default

    if comment is None:
        return default

    comment = comment.strip()
    if not comment:
        return default

    url_matches = list(_URL_PATTERN.finditer(comment))
    url = url_matches[0].group(0) if url_matches else None

    cleaned_comment = comment
    if url_matches:
        cleaned_comment = _URL_PATTERN.sub(" ", comment)

    charge = None
    spin = None
    charge_match = _CHARGE_PATTERN.search(cleaned_comment)
    if charge_match:
        with suppress(ValueError):
            charge = int(charge_match.group(1))
    spin_match = _SPIN_PATTERN.search(cleaned_comment)
    if spin_match:
        with suppress(ValueError):
            spin = int(spin_match.group(1))

    if charge is None or spin is None:
        tokens = cleaned_comment.split()
        candidates = []
        if len(tokens) >= 2:
            candidates.append((tokens[0], tokens[1]))
        if len(tokens) >= 3:
            candidates.append((tokens[1], tokens[2]))
        for charge_token, spin_token in candidates:
            parsed_charge = None
            parsed_spin = None
            with suppress(ValueError):
                parsed_charge = int(charge_token)
            with suppress(ValueError):
                parsed_spin = int(spin_token)
            if charge is None and parsed_charge is not None:
                charge = parsed_charge
            if spin is None and parsed_spin is not None:
                spin = parsed_spin
            if charge is not None and spin is not None:
                break

    if spin is not None and spin <= 0:
        spin = None

    return XYZMetadata(
        comment=comment,
        url=url,
        charge=charge,
        spin=spin,
    )


def get_predict_unit_with_local_fallback(
    checkpoint: str,
    *,
    device: str,
    logger: Optional[logging.Logger] = None,
):
    """Load a UMA predictor, consulting local checkpoints when necessary."""
    ensure_fairchem_loaded()
    if pretrained_mlip is None or FAIRChemCalculator is None:
        raise RuntimeError("UMA FairChem libraries are unavailable; cannot load checkpoints.")

    custom_root = _resolve_custom_checkpoint_dir()
    local_file = _locate_local_checkpoint(custom_root, checkpoint)
    local_error: Optional[Exception] = None

    if local_file is not None:
        try:
            return _load_predict_unit_from_resolved_local(
                checkpoint, local_file, device, logger
            )
        except Exception as exc:  # pragma: no cover - depends on FairChem
            local_error = exc
            if logger:
                logger.debug(
                    "Unable to load local UMA checkpoint '%s' from %s: %s",
                    checkpoint,
                    local_file,
                    exc,
                )

    try:
        return pretrained_mlip.get_predict_unit(checkpoint, device=device)
    except KeyError as exc:
        if local_file is None:
            raise FileNotFoundError(
                _format_missing_checkpoint_error(checkpoint, custom_root)
            ) from exc

        if local_error is not None:
            raise RuntimeError(
                f"UMA checkpoint '{checkpoint}' found at {local_file} but failed to load: {local_error}"
            ) from local_error

        # If we have a local file but no error recorded, ensure a final attempt.
        return _load_predict_unit_from_resolved_local(
            checkpoint, local_file, device, logger
        )
    except Exception as exc:
        if local_file is not None and local_error is None:
            try:
                return _load_predict_unit_from_resolved_local(
                    checkpoint, local_file, device, logger
                )
            except Exception as local_exc:  # pragma: no cover - fallback path
                if logger:
                    logger.debug(
                        "Secondary local load attempt for checkpoint '%s' failed: %s",
                        checkpoint,
                        local_exc,
                    )
                raise RuntimeError(
                    f"Remote UMA checkpoint load failed ({exc}) and local fallback failed ({local_exc})."
                ) from local_exc
        raise


def _load_predict_unit_from_resolved_local(
    checkpoint: str,
    local_file: Path,
    device: str,
    logger: Optional[logging.Logger],
):
    if logger:
        logger.debug("Loading local UMA checkpoint '%s' from %s", checkpoint, local_file)

    registry = getattr(pretrained_mlip, "MODEL_REGISTRY", None)
    model_dir_raw = getattr(pretrained_mlip, "MODEL_DIR", None)
    if registry is not None and model_dir_raw is not None:
        # Legacy FairChem version: mirror registry + model dir behaviour.
        entry = registry.get(checkpoint)
        if entry is None or entry.get("filename") != local_file.name:
            registry[checkpoint] = {
                "filename": local_file.name,
                "url": None,
                "md5": None,
                "license": "local",
            }

        model_dir = Path(model_dir_raw)
        model_dir.mkdir(parents=True, exist_ok=True)
        target = model_dir / local_file.name
        if target.exists():
            try:
                if os.path.samefile(target, local_file):
                    return pretrained_mlip.get_predict_unit(checkpoint, device=device)
            except OSError:
                pass
            try:
                target.unlink()
            except OSError:
                target.unlink(missing_ok=True)  # type: ignore[arg-type]
        try:
            target.symlink_to(local_file)
        except OSError:
            shutil.copy2(local_file, target)
        return pretrained_mlip.get_predict_unit(checkpoint, device=device)

    return _load_predict_unit_from_local_checkpoint(local_file, device, logger)


def _format_missing_checkpoint_error(checkpoint: str, custom_root: Path) -> str:
    """
    Build a descriptive error message including nearby identifiers and local files.
    """
    available_local = sorted(path.stem for path in custom_root.glob("*.pt"))

    available_registry_set: set[str] = set()
    registry = getattr(pretrained_mlip, "MODEL_REGISTRY", None)
    if isinstance(registry, dict):
        available_registry_set.update(registry.keys())

    model_ckpts = getattr(pretrained_mlip, "_MODEL_CKPTS", None)
    checkpoints = getattr(model_ckpts, "checkpoints", None)
    if isinstance(checkpoints, dict):
        available_registry_set.update(checkpoints.keys())

    builtin_models = getattr(pretrained_mlip, "available_models", ())
    if isinstance(builtin_models, (list, tuple, set)):
        available_registry_set.update(str(item) for item in builtin_models)

    available_registry = sorted(available_registry_set)

    message = (
        f"UMA checkpoint '{checkpoint}' not found in FairChem registry "
        f"or in '{custom_root}'."
    )

    suggestions = []
    for pool in (available_local, available_registry):
        suggestions.extend(get_close_matches(checkpoint, pool, n=1, cutoff=0.6))

    if suggestions:
        unique = []
        for item in suggestions:
            if item not in unique:
                unique.append(item)
        message += f" Did you mean: {', '.join(unique)}?"
    else:
        if available_local or available_registry:
            message += " Available options: "
            if available_local:
                message += f"local[{', '.join(available_local)}]"
            if available_registry:
                if available_local:
                    message += "; "
                message += f"registry[{', '.join(available_registry)}]"

    return message


def _load_predict_unit_from_local_checkpoint(
    local_file: Path,
    device: str,
    logger: Optional[logging.Logger],
):
    if load_predict_unit is None:
        raise RuntimeError(
            "Local checkpoint fallback unavailable; FairChem version lacks load_predict_unit."
        )

    try:
        return load_predict_unit(local_file, device=device)
    except Exception as exc:
        if "model_version" not in str(exc):
            raise

        predictor = _load_legacy_umamodel(local_file, device)
        if predictor is not None:
            if logger:
                logger.debug(
                    "Patched legacy UMA checkpoint '%s' by dropping unsupported keys.",
                    local_file.name,
                )
            return predictor
        raise


def _load_legacy_umamodel(local_file: Path, device: str):
    try:
        from torch.serialization import add_safe_globals
    except Exception:  # pragma: no cover - torch absent
        return None

    try:
        from fairchem.core.units.mlip_unit.api.inference import MLIPInferenceCheckpoint
    except ModuleNotFoundError:  # pragma: no cover
        return None

    add_safe_globals([MLIPInferenceCheckpoint])

    import torch  # local import to avoid hard dependency when unused

    checkpoint = torch.load(local_file, map_location="cpu", weights_only=False)
    backbone_cfg = None
    if isinstance(getattr(checkpoint, "model_config", None), dict):
        backbone_cfg = checkpoint.model_config.get("backbone")

    if not isinstance(backbone_cfg, dict) or "model_version" not in backbone_cfg:
        return None

    backbone_cfg.pop("model_version", None)

    with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as tmp_file:
        tmp_path = Path(tmp_file.name)

    torch.save(checkpoint, tmp_path)

    try:
        return load_predict_unit(tmp_path, device=device)
    finally:
        try:
            tmp_path.unlink()
        except OSError:
            pass
