#!/usr/bin/python3 -I
"""fsl-data-install: fetch FSL's large data sets on demand instead of shipping them in RPMs.

The big FSL data sets (FIRST subcortical models, the OMM-1 atlas, ...) are 0.3-1.1 GB each and are
published by FMRIB on their public conda channel. Rehosting them in a package repository would
multiply that by every chroot; this tool downloads the original archive, checks its sha256 against
the manifest, and unpacks only the data/ tree into a state directory, then links it into
$FSLDIR/data where FSL looks for it.

  fsl-data-install --list                 what is available, how big, and whether it is installed
  fsl-data-install NAME [NAME...]         download, verify, install
  fsl-data-install --all                  everything in the manifest
  fsl-data-install --remove NAME          remove the link and the unpacked data
  fsl-data-install --from FILE NAME       use a local copy of the archive (air-gapped hosts)
  fsl-data-install --status               exit 0 only if every NAME given (or all) is installed

State lives in /var/lib/fsl-data (override: --state), links in /usr/lib/fsl/data (--fsldir).
"""
import argparse, bz2, hashlib, os, shutil, sys, tarfile, urllib.request

DEFAULT_MANIFEST = os.environ.get("FSL_DATA_MANIFEST", "/usr/share/fsl-data-fetch/manifest.tsv")
COLS = ("name", "version", "url", "sha256", "bytes", "prefix", "link", "description")


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


def load_manifest(path):
    rows = {}
    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] == "name":
                    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"]); rows[r["name"]] = r
    except OSError as e:
        die(f"cannot read manifest {path}: {e}")
    return rows


def paths(r, a):
    state = os.path.join(a.state, r["name"])
    link = os.path.join(a.fsldir, "data", r["link"])
    target = os.path.join(state, r["prefix"])
    return state, link, target


def status(r, a):
    state, link, target = paths(r, a)
    if os.path.islink(link) and os.path.isdir(link) and os.path.realpath(link) == os.path.realpath(target) \
            and os.path.exists(os.path.join(state, ".complete")):
        return "installed"
    if os.path.isdir(link) and not os.path.islink(link) and os.listdir(link):
        return "installed (provided by an RPM)"
    if os.path.islink(link) and not os.path.exists(link):
        return "broken link"
    if os.path.isdir(state + ".partial"):
        return "partial download"
    return "not installed"


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 install(r, a, local):
    state, link, target = paths(r, a)
    st = status(r, a)
    if st != "not installed" and st != "partial download" and st != "broken link":
        print(f"{r['name']}: already {st}"); return
    os.makedirs(a.state, exist_ok=True)
    if not os.access(a.state, os.W_OK):
        die(f"cannot write to {a.state}: run as root (sudo fsl-data-install {r['name']}) or pass --state DIR", 77)
    need = int(r["bytes"] * 2.5)
    free = shutil.disk_usage(a.state).free
    if free < need:
        die(f"{r['name']}: needs about {need >> 20} MB free in {a.state}, only {free >> 20} MB available", 28)
    archive = local
    if archive is None:
        archive = os.path.join(a.state, r["name"] + ".download")
        print(f"{r['name']}: downloading {r['bytes'] >> 20} MB from {r['url']}")
        download(r["url"], archive, r["bytes"])
    got = sha256_of(archive)
    if got != r["sha256"]:
        if local is None:
            os.remove(archive)
        die(f"{r['name']}: sha256 mismatch (got {got}, manifest says {r['sha256']}); nothing installed", 65)
    part = state + ".partial"
    shutil.rmtree(part, ignore_errors=True); os.makedirs(part)
    pre = r["prefix"].rstrip("/") + "/"
    n = 0
    try:
        with tarfile.open(archive, "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 + "/").startswith(pre) or pre.startswith(name + "/"):
                    if m.issym() or m.islnk():
                        die(f"{r['name']}: unexpected link in archive: {m.name}", 65)
                    members.append(m)
            tf.extractall(part, members=members, filter="data")
            n = sum(1 for m in members if m.isfile())
    except (tarfile.TarError, OSError, EOFError) as e:
        shutil.rmtree(part, ignore_errors=True); die(f"{r['name']}: cannot unpack: {e}", 65)
    if n == 0 or not os.path.isdir(os.path.join(part, r["prefix"])):
        shutil.rmtree(part, ignore_errors=True)
        die(f"{r['name']}: archive has no {r['prefix']}/ tree; nothing installed", 65)
    open(os.path.join(part, ".complete"), "w").write(f"{r['version']} {r['sha256']}\n")
    shutil.rmtree(state, ignore_errors=True); os.rename(part, state)
    os.makedirs(os.path.dirname(link), exist_ok=True)
    if os.path.islink(link):
        os.remove(link)
    if os.path.lexists(link):
        die(f"{link} exists and is not a link; remove it first", 73)
    os.symlink(target, link)
    if local is None and os.path.exists(archive):
        os.remove(archive)
    print(f"{r['name']}: installed {n} files, linked {link}")


def remove(r, a):
    state, link, target = paths(r, a)
    if os.path.islink(link):
        os.remove(link)
    shutil.rmtree(state, ignore_errors=True); shutil.rmtree(state + ".partial", ignore_errors=True)
    print(f"{r['name']}: removed")


def main():
    p = argparse.ArgumentParser(description="Fetch FSL data sets on demand.", add_help=True)
    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="FILE")
    p.add_argument("--manifest", default=DEFAULT_MANIFEST)
    p.add_argument("--state", default=os.environ.get("FSL_DATA_STATE", "/var/lib/fsl-data"))
    p.add_argument("--fsldir", default=os.environ.get("FSLDIR", "/usr/lib/fsl"))
    a = p.parse_args()
    m = load_manifest(a.manifest)
    if a.list or (not a.names and not a.all and not a.status):
        print(f"{'NAME':<12} {'SIZE':>8}  {'STATE':<30} DESCRIPTION")
        for r in m.values():
            print(f"{r['name']:<12} {str(r['bytes'] >> 20) + ' MB':>8}  {status(r, a):<30} {r['description']}")
        return
    names = list(m) if (a.all or (a.status and not a.names)) else a.names
    for n in names:
        if n not in m:
            die(f"unknown data set '{n}' (see --list)", 64)
    if a.status:
        bad = [n for n in names if not status(m[n], a).startswith("installed")]
        for n in names:
            print(f"{n}: {status(m[n], a)}")
        sys.exit(1 if bad else 0)
    if a.local and len(names) != 1:
        die("--from needs exactly one NAME", 64)
    for n in names:
        if a.remove:
            remove(m[n], a)
        else:
            install(m[n], a, a.local)


if __name__ == "__main__":
    main()
