#!/usr/bin/python3 -I
"""freesurfer-data-install: fetch FreeSurfer's atlases, template subjects and models on demand.

FreeSurfer keeps this data (about 11 GB per version) in git-annex, not in its source tarball, so
no RPM carries it. This tool downloads each file from MGH's public annex, checks it against the
sha256 and size in its annex key (both recorded in the manifest, generated from the release
tarball by ci/freesurfer-data-manifest.py), and installs it under a state directory, linked into
$FREESURFER_HOME where FreeSurfer reads it.

  freesurfer-data-install --list                    sets, sizes, and whether each is installed
  freesurfer-data-install SET [SET...]              download, verify, install
  freesurfer-data-install recon-all                 what a default recon-all needs (core + synthmorph)
  freesurfer-data-install --all                     every set
  freesurfer-data-install --remove SET              remove the links and the data
  freesurfer-data-install --from DIR SET            use local copies (air-gapped hosts): DIR holds
                                                    files named by annex key, or by their name
  freesurfer-data-install --status [SET...]         exit 0 only if every SET given (or all) is installed

--fs VERSION picks the FreeSurfer version (default: the newest with a manifest installed). State
lives in /var/lib/freesurfer-data/VERSION (override: --state), links go in
/usr/lib/freesurfer-VERSION (--fshome). The command line and exit codes follow fsl-data-install:
64 usage, 65 bad data, 73 a path is in the way, 77 needs root, 28 no space.

Neurofedora packaging tool, MIT. The data is MGH's, under the FreeSurfer licence; it is not rehosted.
"""
import argparse, glob, hashlib, os, shutil, sys, tarfile, urllib.request

MANIFEST_DIR = os.environ.get("FREESURFER_DATA_MANIFESTS", "/usr/share/freesurfer-data-fetch")
COLS = ("set", "name", "dest", "action", "bytes", "sha256", "key", "url")
ALIASES = {"recon-all": ["core", "synthmorph"]}
DESCRIPTIONS = {
    "core": "average/ atlases and colortables, fsafd, lib/bem icosahedra, the fsaverage subject, the SynthStrip models",
    "synthmorph": "SynthMorph models; recon-all's default (v8) registration needs deform.3 (3.5 GB)",
    "subjects": "fsaverage3-6, fsaverage_sym, cvs_avg35(_inMNI152), fsmni152, V1/EC_average",
    "models": "SynthSeg, SynthSR, SynthStrip, EasyReg, thalseg, topofit and other model files",
    "trctrain": "TRACULA training set",
    "other": "fsfast docs, matlab example data",
}


def die(msg, code=1):
    print("freesurfer-data-install: " + msg, file=sys.stderr)
    sys.exit(code)


def versions():
    return sorted((os.path.basename(os.path.dirname(p)) for p in glob.glob(os.path.join(MANIFEST_DIR, "*", "manifest.tsv"))),
                  key=lambda v: [int(x) if x.isdigit() else x for x in v.split(".")])


def load_manifest(path):
    sets = {}
    try:
        with open(path) as f:
            for line in f:
                if not line.strip() or line.startswith("#"):
                    continue
                parts = line.rstrip("\n").split("\t")
                if parts[0] == "set":
                    continue
                if len(parts) != len(COLS):
                    die(f"{path}: bad row (expected {len(COLS)} fields): {line.strip()[:80]}")
                r = dict(zip(COLS, parts)); r["bytes"] = int(r["bytes"])
                if r["action"] not in ("place", "unpack"):
                    die(f"{path}: bad action {r['action']!r} for {r['dest']}")
                if os.path.isabs(r["dest"]) or ".." in r["dest"].split("/"):
                    die(f"{path}: unsafe destination {r['dest']}", 65)
                sets.setdefault(r["set"], []).append(r)
    except OSError as e:
        die(f"cannot read manifest {path}: {e}")
    return sets


def sha256_of(path):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def download(url, dest, size):
    """Resumable: a partial file is continued with a Range request when the server allows it."""
    have = os.path.getsize(dest) if os.path.exists(dest) else 0
    if have > size:
        os.remove(dest); have = 0
    if have == size:
        return
    req = urllib.request.Request(url)
    if have and url.startswith(("http://", "https://")):
        req.add_header("Range", f"bytes={have}-")
    try:
        resp = urllib.request.urlopen(req, timeout=60)
    except Exception as e:
        die(f"download failed: {url}: {e}")
    mode = "ab" if have and getattr(resp, "status", 200) == 206 else "wb"
    done = have if mode == "ab" else 0
    last = -1
    with open(dest, mode) as out:
        while True:
            chunk = resp.read(1 << 20)
            if not chunk:
                break
            out.write(chunk); done += len(chunk)
            pct = int(100 * done / size) if size else 0
            if pct != last and sys.stderr.isatty():
                print(f"\r  {done >> 20} / {size >> 20} MB ({pct}%)", end="", file=sys.stderr); last = pct
    if sys.stderr.isatty():
        print(file=sys.stderr)


def link_targets(r, state):
    """The (link in FREESURFER_HOME, target in state) pairs one row installs, relative paths."""
    if r["action"] == "place":
        return [(r["dest"], os.path.join(state, r["dest"]))]
    done = os.path.join(state, ".unpacked", r["key"])
    if not os.path.exists(done):
        return []
    with open(done) as f:
        tops = [t for t in f.read().split("\n") if t]
    base = "" if r["dest"] == "." else r["dest"]
    return [(os.path.join(base, t), os.path.join(state, base, t)) for t in tops]


def row_status(r, a):
    if r["action"] == "unpack" and not os.path.exists(os.path.join(a.state, ".unpacked", r["key"])):
        return "missing"
    for link, target in link_targets(r, a.state):
        p = os.path.join(a.fshome, link)
        if os.path.islink(p) and os.path.realpath(p) == os.path.realpath(target) and os.path.exists(target):
            continue
        if os.path.exists(p) and not os.path.islink(p):
            continue                                   # shipped by an RPM: counts as present
        return "missing"
    return "ok"


def set_status(rows, a):
    n = sum(1 for r in rows if row_status(r, a) == "ok")
    return "installed" if n == len(rows) else ("not installed" if n == 0 else f"partial ({n}/{len(rows)} files)")


def place_link(link, target, a):
    p = os.path.join(a.fshome, link)
    os.makedirs(os.path.dirname(p), exist_ok=True)
    if os.path.islink(p):
        os.remove(p)
    elif os.path.exists(p):
        return False                                   # an RPM's file; leave it
    os.symlink(target, p)
    return True


def install_row(r, a, local):
    if row_status(r, a) == "ok":
        return 0
    state = a.state
    src = None
    if local:
        for cand in (os.path.join(local, r["key"]), os.path.join(local, r["name"])):
            if os.path.isfile(cand):
                src = cand; break
        if src is None:
            die(f"--from {local}: no {r['key']} or {r['name']}", 66)
    if src is None:
        os.makedirs(os.path.join(state, ".download"), exist_ok=True)
        src = os.path.join(state, ".download", r["key"])
        print(f"  {r['dest'] if r['action'] == 'place' else r['name']}: {r['bytes'] >> 20} MB")
        download(r["url"], src, r["bytes"])
    got_size = os.path.getsize(src)
    if got_size != r["bytes"]:
        if not local:
            os.remove(src)
        die(f"{r['name']}: size {got_size}, manifest says {r['bytes']}; nothing installed", 65)
    got = sha256_of(src)
    if got != r["sha256"]:
        if not local:
            os.remove(src)
        die(f"{r['name']}: sha256 mismatch (got {got}, manifest says {r['sha256']}); nothing installed", 65)
    if r["action"] == "place":
        dst = os.path.join(state, r["dest"])
        os.makedirs(os.path.dirname(dst), exist_ok=True)
        if local:
            shutil.copyfile(src, dst + ".part")
        else:
            os.replace(src, dst + ".part")
        os.replace(dst + ".part", dst)
        place_link(r["dest"], dst, a)
        return 1
    base = os.path.join(state, "" if r["dest"] == "." else r["dest"])
    os.makedirs(base, exist_ok=True)
    tops = set()
    try:
        with tarfile.open(src, "r:*") as tf:
            members = []
            for m in tf.getmembers():
                name = os.path.normpath(m.name)
                if os.path.isabs(m.name) or name.startswith(".."):
                    die(f"{r['name']}: unsafe path in archive: {m.name}", 65)
                if name in (".", ""):
                    continue
                tops.add(name.split("/")[0]); members.append(m)
            tf.extractall(base, members=members, filter="data")
    except (tarfile.TarError, OSError, EOFError) as e:
        die(f"{r['name']}: cannot unpack: {e}", 65)
    os.makedirs(os.path.join(state, ".unpacked"), exist_ok=True)
    with open(os.path.join(state, ".unpacked", r["key"]), "w") as f:
        f.write("\n".join(sorted(tops)) + "\n")
    for link, target in link_targets(r, state):
        place_link(link, target, a)
    if not local and os.path.exists(src):
        os.remove(src)
    return 1


def install_set(name, rows, a, local):
    st = set_status(rows, a)
    if st == "installed":
        print(f"{name}: already installed"); return
    os.makedirs(a.state, exist_ok=True)
    if not os.access(a.state, os.W_OK) or (os.path.isdir(a.fshome) and not os.access(a.fshome, os.W_OK)):
        die(f"cannot write to {a.state} or {a.fshome}: run as root (sudo freesurfer-data-install {name}) "
            f"or pass --state and --fshome", 77)
    need = int(sum(r["bytes"] for r in rows if row_status(r, a) != "ok") * 2.2)
    free = shutil.disk_usage(a.state).free
    if free < need:
        die(f"{name}: needs about {need >> 20} MB free in {a.state}, only {free >> 20} MB available", 28)
    print(f"{name}: {len(rows)} files, {sum(r['bytes'] for r in rows) >> 20} MB")
    n = sum(install_row(r, a, local) for r in rows)
    print(f"{name}: installed ({n} fetched, {len(rows) - n} already present)")


def remove_set(name, rows, a):
    for r in rows:
        for link, target in link_targets(r, a.state):
            p = os.path.join(a.fshome, link)
            if os.path.islink(p) and os.path.realpath(p) == os.path.realpath(target):
                os.remove(p)
            if os.path.isdir(target) and not os.path.islink(target):
                shutil.rmtree(target, ignore_errors=True)
            elif os.path.lexists(target):
                os.remove(target)
        u = os.path.join(a.state, ".unpacked", r["key"])
        if os.path.exists(u):
            os.remove(u)
    print(f"{name}: removed")


def main():
    p = argparse.ArgumentParser(description="Fetch FreeSurfer data on demand.")
    p.add_argument("names", nargs="*"); p.add_argument("--list", action="store_true")
    p.add_argument("--all", action="store_true"); p.add_argument("--remove", action="store_true")
    p.add_argument("--status", action="store_true"); p.add_argument("--from", dest="local", metavar="DIR")
    p.add_argument("--fs", metavar="VERSION"); p.add_argument("--manifest")
    p.add_argument("--state"); p.add_argument("--fshome")
    a = p.parse_args()
    if a.manifest is None:
        vs = versions()
        if not vs:
            die(f"no manifests in {MANIFEST_DIR}", 66)
        a.fs = a.fs or vs[-1]
        if a.fs not in vs:
            die(f"no manifest for FreeSurfer {a.fs} (have: {', '.join(vs)})", 64)
        a.manifest = os.path.join(MANIFEST_DIR, a.fs, "manifest.tsv")
    a.fs = a.fs or "unknown"
    a.state = a.state or os.environ.get("FREESURFER_DATA_STATE", f"/var/lib/freesurfer-data/{a.fs}")
    a.fshome = a.fshome or f"/usr/lib/freesurfer-{a.fs}"
    sets = load_manifest(a.manifest)
    if a.list or (not a.names and not a.all and not a.status):
        print(f"FreeSurfer {a.fs}: {a.manifest}\n{'SET':<11} {'FILES':>5} {'SIZE':>9}  {'STATE':<24} DESCRIPTION")
        for s, rows in sets.items():
            print(f"{s:<11} {len(rows):>5} {str(sum(r['bytes'] for r in rows) >> 20) + ' MB':>9}  "
                  f"{set_status(rows, a):<24} {DESCRIPTIONS.get(s, '')}")
        for al, members in ALIASES.items():
            print(f"{al:<11} {'':>5} {'':>9}  {'':<24} = {' + '.join(members)}")
        return
    names = []
    for n in (list(sets) if (a.all or (a.status and not a.names)) else a.names):
        names += ALIASES.get(n, [n])
    for n in names:
        if n not in sets:
            die(f"unknown set '{n}' (see --list)", 64)
    if a.status:
        bad = [n for n in names if set_status(sets[n], a) != "installed"]
        for n in names:
            print(f"{n}: {set_status(sets[n], a)}")
        sys.exit(1 if bad else 0)
    for n in names:
        if a.remove:
            remove_set(n, sets[n], a)
        else:
            install_set(n, sets[n], a, a.local)


if __name__ == "__main__":
    main()
