"""Regenerate the shipped modified-RNA .lib files from an external modXNA
checkout. Maintainer-only; not installed or used at runtime.

    python tools/generate.py --modxna /path/to/modxna [--amberhome DIR]

Requires AmberTools (cpptraj >= 6.26, tleap) reachable via --amberhome or the
AMBERHOME env var, or with cpptraj already on PATH.
"""
import argparse
import os
import re
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
from modrna_params._data import (  # noqa: E402
    RESIDUES, TERMINI, MODXNA_BASE, NET_CHARGE, _FORM_RECIPE, SUGAR,
    unit_name, expected_atom_count,
)


def amber_env(amberhome: str) -> dict:
    bindir = str(Path(amberhome) / "bin")
    env = dict(os.environ)
    env["AMBERHOME"] = amberhome
    env["PATH"] = bindir + os.pathsep + env.get("PATH", "")
    env["CPPTRAJ"] = str(Path(bindir) / "cpptraj")
    return env


def dedupe_atom_names(text: str, unit: str) -> tuple[str, list[tuple[str, str]]]:
    """Make atom names unique within a unit's .lib.

    Some modXNA fragments (e.g. the I6A isopentenyl group) define two atoms
    with the same name. tleap tolerates this but parmed and other tools reject
    it. Renaming the 2nd+ occurrence (e.g. C11 -> C11A) is a pure label fix:
    .lib connectivity is index-based and force-field terms key off atom types,
    so nothing chemical changes. Applied to both the `atoms` and
    `atomspertinfo` tables, which list atoms in the same order.
    """
    atoms_hdr = f"!entry.{unit}.unit.atoms table"
    perinfo_hdr = f"!entry.{unit}.unit.atomspertinfo table"
    lines = text.splitlines(keepends=True)

    def block_names(header: str) -> list[str]:
        names, inside = [], False
        for ln in lines:
            if ln.startswith(header):
                inside = True
                continue
            if inside:
                if ln.startswith("!entry."):
                    break
                m = re.match(r'\s*"([^"]*)"', ln)
                if m:
                    names.append(m.group(1))
        return names

    names = block_names(atoms_hdr)
    used = set(names)  # any name that appears at all is "taken"
    seen: set[str] = set()
    new_by_pos: dict[int, str] = {}
    renames: list[tuple[str, str]] = []
    for i, nm in enumerate(names):
        if nm not in seen:
            seen.add(nm)
            continue
        for suf in "ABCDEFGHIJKLMNOPQRSTUVWXYZ":
            cand = (nm[:3] + suf)[:4]
            if cand not in used and cand not in seen:
                break
        else:
            raise RuntimeError(f"{unit}: cannot make {nm!r} unique")
        new_by_pos[i] = cand
        used.add(cand)
        seen.add(cand)
        renames.append((nm, cand))

    if not new_by_pos:
        return text, []

    def rewrite(header: str) -> None:
        inside, pos = False, -1
        for idx, ln in enumerate(lines):
            if ln.startswith(header):
                inside, pos = True, -1
                continue
            if inside:
                if ln.startswith("!entry."):
                    inside = False
                    continue
                m = re.match(r'(\s*")([^"]*)(".*)', ln, re.S)
                if m:
                    pos += 1
                    if pos in new_by_pos:
                        lines[idx] = m.group(1) + new_by_pos[pos] + m.group(3)

    rewrite(atoms_hdr)
    rewrite(perinfo_hdr)
    return "".join(lines), renames


def _unit(lib: Path, name: str):
    import parmed
    off = parmed.amber.AmberOFFLibrary.parse(str(lib))
    if name not in off:
        raise RuntimeError(f"{name}: unit not found in {lib.name}")
    return off[name]


def generate_one(res: str, terminus: str, modxna: str, env: dict,
                 libdir: Path) -> tuple:
    base = MODXNA_BASE[res]
    backbone, flag = _FORM_RECIPE[terminus]
    name = unit_name(res, terminus)
    with tempfile.TemporaryDirectory() as td:
        Path(td, "in.txt").write_text(f"{backbone} {SUGAR} {base}\n")
        cmd = ["bash", str(Path(modxna) / "modxna.sh"),
               "-i", "in.txt", "-m", name, "--clean", "--nomin"]
        if flag:
            cmd.append(flag)
        proc = subprocess.run(cmd, cwd=td, env=env,
                              capture_output=True, text=True)
        lib = Path(td, f"{name}.lib")
        if not lib.exists():
            raise RuntimeError(
                f"{name}: modxna.sh produced no lib (rc={proc.returncode})\n"
                f"{proc.stdout[-2000:]}\n{proc.stderr[-1000:]}")
        fixed, renames = dedupe_atom_names(lib.read_text(), name)
        if renames:
            lib.write_text(fixed)
            for old, new in renames:
                print(f"     renamed duplicate atom {old!r} -> {new!r}")
        unit = _unit(lib, name)
        q = sum(a.charge for a in unit.atoms)
        if abs(q - NET_CHARGE[terminus]) > 5e-3:
            raise RuntimeError(
                f"{name}: net charge {q:.4f} != target {NET_CHARGE[terminus]}")
        n = len(unit.atoms)
        if n != expected_atom_count(res, terminus):
            raise RuntimeError(
                f"{name}: {n} atoms != expected "
                f"{expected_atom_count(res, terminus)}")
        if res == "PSU" and any(a.element == 16 for a in unit.atoms):
            raise RuntimeError(f"{name}: unexpected sulfur (thiouracil?)")
        shutil.copyfile(lib, libdir / f"{name}.lib")
        return name, base, terminus, n, q


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--modxna", default=os.environ.get("MODXNA_HOME"))
    ap.add_argument("--amberhome", default=os.environ.get("AMBERHOME"))
    args = ap.parse_args()
    if not args.modxna:
        sys.exit("--modxna PATH (or MODXNA_HOME) is required")
    if not args.amberhome:
        cpptraj = shutil.which("cpptraj")
        if not cpptraj:
            sys.exit("cpptraj not found; pass --amberhome or set AMBERHOME")
        args.amberhome = str(Path(cpptraj).resolve().parent.parent)
    env = amber_env(args.amberhome)

    data = Path(__file__).resolve().parent.parent / "src/modrna_params/data"
    libdir = data / "lib"
    libdir.mkdir(parents=True, exist_ok=True)
    shutil.copyfile(Path(args.modxna) / "dat/frcmod.modxna",
                    data / "frcmod.modxna")

    rows = []
    for res in RESIDUES:
        for term in TERMINI:
            name, base, t, n, q = generate_one(res, term, args.modxna, env,
                                               libdir)
            rows.append(name)
            print(f"  OK {name:6s} base={base:4s} {t:8s} atoms={n} q={q:+.4f}")
    print(f"generated {len(rows)} libs into {libdir}")


if __name__ == "__main__":
    main()
