"""Stream one MR-RATE atlas batch tar -> one uint8 memmap shard + index CSV.

One shard per batch rather than one file per study keeps the file count small.

Written to --out:
    <batch>.u8    raw uint8, shape (n_studies, 4, *geometry output shape), C order
    <batch>.csv   row,study_uid,has_t1w,has_t2w,has_flair,has_swi,n_series
"""

from __future__ import annotations

import argparse
import csv
import io
import os
import sys
import tarfile
import time
from multiprocessing import Pool
from threading import Semaphore

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from preprocess import MODALITIES, process_zip  # noqa: E402

GEOMETRY = "d15"


def _init(geometry):
    global GEOMETRY
    GEOMETRY = geometry


def _work(job):
    uid, raw = job
    try:
        per, n = process_zip(io.BytesIO(raw), (GEOMETRY,))
        vol, presence = per[GEOMETRY]
        return uid, vol, presence, n, ""
    except Exception as exc:  # a bad study must not kill the shard
        return uid, None, None, 0, f"{type(exc).__name__}: {exc}"


def members(tar_path: str, wanted: set[str] | None):
    """Yield (uid, zip_bytes) streaming forward through the tar exactly once."""
    with tarfile.open(tar_path, "r|") as tf:
        for m in tf:
            if not m.isfile() or not m.name.endswith(".zip"):
                continue
            uid = os.path.basename(m.name)[: -len(".zip")]
            if uid.endswith("_atlas"):
                uid = uid[: -len("_atlas")]
            if wanted is not None and uid not in wanted:
                continue
            if m.size == 0:
                continue
            data = tf.extractfile(m).read()
            yield uid, data


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--tar", required=True)
    ap.add_argument("--batch", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--geometry", default="d15", choices=("d15", "d2"))
    ap.add_argument("--workers", type=int, default=16)
    ap.add_argument("--uid-list", default=None, help="optional file of study_uids to keep")
    ap.add_argument("--limit", type=int, default=0)
    args = ap.parse_args()

    wanted = None
    if args.uid_list:
        with open(args.uid_list) as f:
            wanted = {ln.strip() for ln in f if ln.strip()}

    os.makedirs(args.out, exist_ok=True)
    bin_path = os.path.join(args.out, f"{args.batch}.u8")
    csv_path = os.path.join(args.out, f"{args.batch}.csv")
    tmp_bin = bin_path + ".part"
    tmp_csv = csv_path + ".part"

    t0 = time.time()
    written = 0
    failed = 0
    src = members(args.tar, wanted)
    if args.limit:
        src = (x for i, x in enumerate(src) if i < args.limit)

    # Pool.imap's feeder queues without bound and each item is a ~150 MB zip,
    # so gate it on completed results to keep the tar out of RAM.
    inflight = Semaphore(args.workers * 3)

    def throttled(it):
        for item in it:
            inflight.acquire()
            yield item

    src = throttled(src)

    with open(tmp_bin, "wb", buffering=1024 * 1024) as fbin, open(tmp_csv, "w", newline="") as fcsv:
        w = csv.writer(fcsv)
        w.writerow(["row", "study_uid", *[f"has_{m}" for m in MODALITIES], "n_series"])
        with Pool(args.workers, maxtasksperchild=64, initializer=_init, initargs=(args.geometry,)) as pool:
            for uid, vol, presence, n, err in pool.imap(_work, src, chunksize=1):
                inflight.release()
                if vol is None or n == 0:
                    failed += 1
                    if failed <= 20:
                        print(f"[skip] {uid} {err or 'no usable series'}", flush=True)
                    continue
                fbin.write(vol.tobytes())
                w.writerow([written, uid, *presence.tolist(), n])
                written += 1
                if written % 200 == 0:
                    el = time.time() - t0
                    print(f"[{args.batch}] {written} studies  {el/60:.1f} min "
                          f"({written/max(el,1e-9):.2f}/s, {failed} skipped)", flush=True)

    os.replace(tmp_bin, bin_path)
    os.replace(tmp_csv, csv_path)
    el = time.time() - t0
    print(f"[{args.batch}] DONE {written} studies, {failed} skipped, "
          f"{el/60:.1f} min, {os.path.getsize(bin_path)/2**30:.1f} GiB", flush=True)
    return 0


if __name__ == "__main__":
    sys.exit(main())