#!/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
  freesurfer-data-install -j N SET                  download N files at a time (default 4; -j 1 is serial)

--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
from concurrent.futures import ThreadPoolExecutor, as_completed

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 (recon-all -label-v1 needs V1_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, progress=True):
    """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 progress and pct != last and sys.stderr.isatty():
                print(f"\r  {done >> 20} / {size >> 20} MB ({pct}%)", end="", file=sys.stderr); last = pct
    if progress and 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 prefetch(rows, a):
    """Download the missing files a.jobs at a time into state/.download. install_row then finds
    each one complete there and verifies, places or unpacks it serially, as before."""
    todo = [r for r in rows if row_status(r, a) != "ok"]
    if a.jobs <= 1 or len(todo) < 2:
        return
    os.makedirs(os.path.join(a.state, ".download"), exist_ok=True)
    print(f"  downloading {len(todo)} files, {a.jobs} at a time")
    ex = ThreadPoolExecutor(a.jobs)
    try:
        futs = {ex.submit(download, r["url"], os.path.join(a.state, ".download", r["key"]), r["bytes"], False): r
                for r in todo}
        for i, f in enumerate(as_completed(futs), 1):
            f.result()                                 # a failed download exits here, via die()
            r = futs[f]
            print(f"  [{i}/{len(todo)}] {r['dest'] if r['action'] == 'place' else r['name']}: {r['bytes'] >> 20} MB")
    except BaseException:
        ex.shutdown(wait=False, cancel_futures=True)   # don't fetch the rest after a failure
        raise
    ex.shutdown()


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")
    if not local:
        prefetch(rows, a)
    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")
    p.add_argument("-j", "--jobs", type=int, default=4, metavar="N", help="parallel downloads (default 4)")
    a = p.parse_args()
    if a.jobs < 1:
        die("-j must be at least 1", 64)
    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()
