#!/usr/bin/env python3
"""P8: fetch the GSE165845 IDATs (120 1000 Genomes individuals × 3 replicates, GSA-MD v1) and lay them out as
``<barcode>/<barcode>_<RxxCxx>_{Grn,Red}.idat.gz`` symlinks, reading barcode and position from each IDAT header
(the GEO file names are ``GSM…_<individual>_<replicate>_{Grn,Red}.idat.gz``).

Writes ``<out>/samples_map.tsv`` (sample_id, gsm, individual, replicate).

    python scripts/p8/prepare.py --tsv scripts/p8/gse165845_samples.tsv --dl DIR --out DIR [--only NA12878,HG00096]
"""

from __future__ import annotations

import argparse
import csv
import sys
import time
import urllib.request
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

from simurg.formats.idat import read_idat


def fetch(url: str, dest: Path, tries: int = 5) -> Path:
    if dest.exists() and dest.stat().st_size > 0:
        return dest
    tmp = dest.with_suffix(dest.suffix + ".part")
    for i in range(tries):
        try:
            with urllib.request.urlopen(url, timeout=120) as r, open(tmp, "wb") as fh:
                while chunk := r.read(1 << 20):
                    fh.write(chunk)
            tmp.rename(dest)
            return dest
        except OSError as e:
            print(f"retry {i + 1}/{tries} {url}: {e}", file=sys.stderr)
            time.sleep(5 * (i + 1))
    raise RuntimeError(f"download failed: {url}")


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--tsv", required=True, type=Path)
    ap.add_argument("--dl", required=True, type=Path, help="download folder (GEO names)")
    ap.add_argument("--out", required=True, type=Path, help="folder for the standard-name symlink tree")
    ap.add_argument("--only", default="", help="comma-separated individuals (test subsets)")
    ap.add_argument("--workers", type=int, default=16)
    a = ap.parse_args()

    rows = list(csv.DictReader(a.tsv.open(), delimiter="\t"))
    if a.only:
        keep = set(a.only.split(","))
        rows = [r for r in rows if r["individual"] in keep]
    a.dl.mkdir(parents=True, exist_ok=True)
    jobs = [(r[k], a.dl / r[k].rsplit("/", 1)[1]) for r in rows for k in ("grn_url", "red_url")]
    t0 = time.time()
    with ThreadPoolExecutor(a.workers) as ex:
        list(ex.map(lambda j: fetch(*j), jobs))
    size = sum(p.stat().st_size for _, p in jobs)
    print(f"{len(jobs)} IDATs, {size / 1e9:.1f} GB in {time.time() - t0:.0f} s")

    a.out.mkdir(parents=True, exist_ok=True)
    seen: dict[str, str] = {}
    with open(a.out / "samples_map.tsv", "w") as fh:
        fh.write("sample_id\tgsm\tindividual\treplicate\n")
        for r in rows:
            names = []
            for ch, key in (("Grn", "grn_url"), ("Red", "red_url")):
                src = a.dl / r[key].rsplit("/", 1)[1]
                idat = read_idat(src, load_arrays=False)
                names.append(f"{idat.barcode}_{idat.position}")
                d = a.out / idat.barcode
                d.mkdir(exist_ok=True)
                link = d / f"{idat.barcode}_{idat.position}_{ch}.idat.gz"
                if not link.exists():
                    link.symlink_to(src.resolve())
            if names[0] != names[1]:
                raise SystemExit(f"{r['gsm']}: Grn/Red headers disagree ({names})")
            sid = names[0]
            if sid in seen:
                raise SystemExit(f"{r['gsm']}: {sid} already used by {seen[sid]}")
            seen[sid] = r["gsm"]
            fh.write(f"{sid}\t{r['gsm']}\t{r['individual']}\t{r['replicate']}\n")
    print(f"{len(seen)} samples -> {a.out}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())