"""Download the public PHerc0139 winding meshes (9.362 um tifxyz) and Vesuvius Challenge ink previews (ds8 JPGs).

  python tools/fetch.py        (from the Space's root; stdlib only; writes raw/, about 110 MB)

It lists the open data bucket itself, keeps the wNNN winding segments that have a 9.362 um mesh, and writes
raw/inv_PHerc0139.json plus raw/wNNN/{x,y,z.tif, meta.json, ink2um_ds8.jpg, ink1um_ds8.jpg, source.json}.
UNROLL_RAW overrides the output folder. tools/build.py reads the X-ray layer straight from the bucket.
"""
import json
import os
import re
import time
import urllib.parse
import urllib.request
from concurrent.futures import ThreadPoolExecutor

B = "https://vesuvius-challenge-open-data.s3.amazonaws.com"
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
RAW = os.environ.get("UNROLL_RAW", f"{ROOT}/raw")
SCROLL = os.environ.get("UNROLL_SCROLL", "PHerc0139")
MESH_RES = os.environ.get("UNROLL_MESH_RES", "9.362um")      # which tifxyz to take (the coarse scan keeps int16 coordinates in range)
MESH_VOL = os.environ.get("UNROLL_MESH_VOL", "")             # optional volume id, when a segment has two meshes at one resolution
INK_RES = os.environ.get("UNROLL_INK_RES", "")               # which ink preview to call ink2um_ds8.jpg ("" = the non-1.129 um one, newest model)
NO_INK = os.environ.get("UNROLL_NO_INK") == "1"              # X-ray only: fetch no ink at all


def get(url, tries=5):
    for i in range(tries):
        try:
            with urllib.request.urlopen(url, timeout=120) as r:
                data = r.read()
                n = r.headers.get("Content-Length")
                if n is None or len(data) == int(n):
                    return data
        except OSError:
            if i == tries - 1:
                raise
        time.sleep(1 + i)
    raise OSError(f"short reads from {url}")


def listing(prefix):
    """(sub-prefixes, [(key, size)]) under prefix, following continuation tokens."""
    dirs, keys, token = [], [], None
    while True:
        q = {"list-type": "2", "prefix": prefix, "delimiter": "/"}
        if token:
            q["continuation-token"] = token
        t = get(f"{B}/?{urllib.parse.urlencode(q)}").decode()
        dirs += [p for p in re.findall(r"<Prefix>(.*?)</Prefix>", t) if p != prefix]
        keys += [(k, int(s)) for k, s in re.findall(r"<Key>(.*?)</Key>.*?<Size>(\d+)</Size>", t)]
        m = re.search(r"<NextContinuationToken>(.*?)</NextContinuationToken>", t)
        if not m:
            return dirs, keys
        token = m.group(1)


def segment(seg_prefix):
    d = {"seg": seg_prefix.rstrip("/").split("/")[-1]}
    sub, _ = listing(seg_prefix)
    for x in sub:
        name = x.rstrip("/").split("/")[-1]
        if name == "mesh":
            d["mesh"] = [p.rstrip("/").split("/")[-1] for p in listing(x)[0]]
        elif name == "ink-detection":
            d["ink"] = [(k.split("/")[-1], s) for k, s in listing(x)[1]]
            d["ink_ds"] = [k.split("/")[-1] for k, _ in listing(x + "downsampled/")[1]]
        elif name == "surface-volumes":
            d["sv"] = [p.rstrip("/").split("/")[-1] for p in listing(x)[0]]
    return d


def main():
    os.makedirs(RAW, exist_ok=True)
    segs, _ = listing(f"{SCROLL}/segments/")
    with ThreadPoolExecutor(8) as ex:
        inv = list(ex.map(segment, segs))
    json.dump(inv, open(f"{RAW}/inv_{SCROLL}.json", "w"), indent=1)
    jobs = []
    named = [re.search(r"-(w\d{3})_", s["seg"]) for s in inv]
    nums = [m.group(1) for m in named if m]
    use_w = len(nums) >= 0.8 * len(inv) and len(set(nums)) == len(nums)
    for s, m in zip(inv, named):
        if use_w and not m:
            continue
        meshes = [x for x in s.get("mesh", []) if x.endswith(f"-{MESH_RES}.tifxyz") and MESH_VOL in x]
        if not meshes:
            continue
        inks = sorted(s.get("ink_ds", []))
        main_ink = [f for f in inks if (INK_RES in f if INK_RES else "1.129um" not in f)]
        if not NO_INK and not main_ink:
            print("skip", s["seg"], "(no ink preview)")
            continue
        # wNNN when the segment names carry unique winding numbers; otherwise the segment timestamp (build renumbers by radius)
        w = m.group(1) if use_w else "s" + s["seg"].split("-")[0]
        mesh = meshes[0]
        d = f"{RAW}/{w}"
        os.makedirs(d, exist_ok=True)
        base = f"{B}/{SCROLL}/segments/{s['seg']}"
        for f in ("x.tif", "y.tif", "z.tif", "meta.json"):
            jobs.append((f"{base}/mesh/{mesh}/{f}", f"{d}/{f}"))
        if not NO_INK:
            jobs.append((f"{base}/ink-detection/downsampled/{main_ink[-1]}", f"{d}/ink2um_ds8.jpg"))
            if True:
                for f in inks:
                    if "1.129um" in f:
                        jobs.append((f"{base}/ink-detection/downsampled/{f}", f"{d}/ink1um_ds8.jpg"))
        json.dump({"seg": s["seg"], "w": w, "mesh": mesh, "sv": s.get("sv"), "ink": s.get("ink"),
                   "ink_ds": None if NO_INK else main_ink[-1]},
                  open(f"{d}/source.json", "w"), indent=1)

    def fetch(job):
        url, path = job
        if os.path.exists(path) and os.path.getsize(path) > 0:
            return 0
        data = get(url)
        open(path, "wb").write(data)
        return len(data)

    with ThreadPoolExecutor(6) as ex:
        n = sum(ex.map(fetch, jobs))
    print(len(jobs), "files,", n // 1_000_000, "MB fetched into", RAW)


if __name__ == "__main__":
    main()