#!/usr/bin/env python3
"""Build and live-verify an inventory of public, complete human WES alignments.

The audit deliberately counts only BAM/CRAM objects with a matching BAI/CRAI.
It proves anonymous GetObject support using a one-byte HTTP Range request for
both the alignment and its index; it does not download multi-terabyte cohorts.
"""

from __future__ import annotations

import argparse
import base64
import csv
import hashlib
import json
import os
import re
import statistics
import sys
import threading
import time
import urllib.error
import urllib.parse
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Iterable


DATE = "2026-08-01"
USER_AGENT = "LGASieve-public-WES-audit/1.0"
PRINT_LOCK = threading.Lock()


def split_multi(value: str) -> list[str]:
    return [x.strip() for x in (value or "").split(";") if x.strip()]


def https_from_ftpish(value: str) -> str:
    value = value.strip()
    if value.startswith("ftp://"):
        value = "https://" + value[6:]
    elif not (value.startswith("http://") or value.startswith("https://")):
        value = "https://" + value

    # ENA object names can legally contain URI-reserved characters, notably
    # literal '#'.  With the default parser that character becomes a client-
    # side fragment and is never sent to the server.  Parse with fragments
    # disabled, then percent-encode the path while preserving existing escapes.
    parsed = urllib.parse.urlsplit(value, allow_fragments=False)
    encoded_path = urllib.parse.quote(
        parsed.path,
        safe="/%:@!$&'()*+,;=-._~",
    )
    encoded_query = urllib.parse.quote(
        parsed.query,
        safe="=&%/:?+,-._~",
    )
    return urllib.parse.urlunsplit(
        (parsed.scheme, parsed.netloc, encoded_path, encoded_query, "")
    )


def s3_https(uri: str) -> str:
    parsed = urllib.parse.urlparse(uri)
    if parsed.scheme != "s3":
        return uri
    return f"https://{parsed.netloc}.s3.amazonaws.com/{urllib.parse.quote(parsed.path.lstrip('/'), safe='/')}"


def gcs_https(bucket: str, name: str) -> str:
    return f"https://storage.googleapis.com/{bucket}/{urllib.parse.quote(name, safe='/')}"


def sample_from_name(name: str) -> str:
    m = re.search(r"\b(HG\d{3,}|NA\d{3,})\b", name)
    if m:
        return m.group(1)
    return Path(urllib.parse.urlparse(name).path).name.split(".")[0]


def load_ena(path: Path) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    pairs: list[dict[str, Any]] = []
    unindexed_bams = 0
    run_rows = 0
    studies: set[str] = set()
    samples: set[str] = set()
    with path.open("r", encoding="utf-8", newline="") as fh:
        for row in csv.DictReader(fh, delimiter="\t"):
            run_rows += 1
            studies.add(row["study_accession"])
            samples.add(row["sample_accession"])
            formats = split_multi(row.get("submitted_format", ""))
            urls = split_multi(row.get("submitted_ftp", ""))
            sizes = split_multi(row.get("submitted_bytes", ""))
            md5s = split_multi(row.get("submitted_md5", ""))
            objects: list[dict[str, Any]] = []
            for i, fmt in enumerate(formats):
                if i >= len(urls):
                    continue
                objects.append(
                    {
                        "format": fmt.upper(),
                        "url": https_from_ftpish(urls[i]),
                        "bytes": int(sizes[i]) if i < len(sizes) and sizes[i].isdigit() else None,
                        "md5": md5s[i] if i < len(md5s) else "",
                    }
                )
            bams = [o for o in objects if o["format"] == "BAM" or o["url"].lower().endswith(".bam")]
            bais = [o for o in objects if o["format"] == "BAI" or o["url"].lower().endswith(".bai")]
            used: set[int] = set()
            for bam in bams:
                bam_url_lower = bam["url"].lower()
                best_i = None
                for i, bai in enumerate(bais):
                    if i in used:
                        continue
                    u = bai["url"].lower()
                    if u == bam_url_lower + ".bai" or u == bam_url_lower[:-4] + ".bai":
                        best_i = i
                        break
                if best_i is None:
                    for i in range(len(bais)):
                        if i not in used:
                            best_i = i
                            break
                if best_i is None:
                    unindexed_bams += 1
                    continue
                used.add(best_i)
                bai = bais[best_i]
                pairs.append(
                    {
                        "repository": "ENA",
                        "collection": row["study_accession"],
                        "run_accession": row["run_accession"],
                        "sample_id": row["sample_accession"],
                        "reference_build": "submitted; inspect BAM header",
                        "platform_or_kit": "submitted WXS",
                        "alignment_format": "BAM",
                        "alignment_url": bam["url"],
                        "index_url": bai["url"],
                        "alignment_bytes": bam["bytes"],
                        "index_bytes": bai["bytes"],
                        "alignment_md5": bam["md5"],
                        "index_md5": bai["md5"],
                        "origin_note": "ENA read_run metadata: Homo sapiens; library_strategy=WXS",
                    }
                )
    return pairs, {
        "query_rows": run_rows,
        "query_studies": len(studies),
        "query_samples": len(samples),
        "indexed_alignment_objects": len(pairs),
        "unindexed_alignment_objects": unindexed_bams,
        "metadata_query": 'tax_eq(9606) AND library_strategy="WXS" AND submitted_format="BAM"',
    }


def load_depmap(path: Path) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    pairs: list[dict[str, Any]] = []
    with path.open("r", encoding="utf-8-sig", newline="") as fh:
        for row in csv.DictReader(fh):
            if row.get("datatype") != "wes":
                continue
            bam_uri = row["cram_bam_url"]
            bai_uri = row["crai_bai_url"]
            pairs.append(
                {
                    "repository": "DepMap_CCLE_AWS",
                    "collection": "DepMap/CCLE",
                    "run_accession": row["sample_id"],
                    "sample_id": row.get("model_id") or row.get("sample_id"),
                    "reference_build": row.get("reference_genome", ""),
                    "platform_or_kit": row.get("baits", "") or "unspecified WES bait kit",
                    "alignment_format": Path(bam_uri).suffix.lstrip(".").upper(),
                    "alignment_url": s3_https(bam_uri),
                    "index_url": s3_https(bai_uri),
                    "alignment_bytes": int(row["cram_bam_size"]),
                    "index_bytes": None,
                    "alignment_md5": row.get("cram_bam_checksum_value", ""),
                    "index_md5": "",
                    "origin_note": f"Human cancer cell line {row.get('cell_line_name','')}; public AWS object; DepMap terms apply",
                }
            )
    return pairs, {"inventory_rows_wes": len(pairs)}


def load_giab(index_root: Path) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    paths = [
        p
        for p in index_root.rglob("alignment.*")
        if "exome" in p.name.lower() and not re.search(r"\.HG00[2345]$", p.name)
    ]
    pairs: list[dict[str, Any]] = []
    for path in sorted(paths):
        technology = path.name.split("alignment.index.", 1)[-1]
        with path.open("r", encoding="utf-8") as fh:
            for line in fh:
                if not line.startswith("ftp://"):
                    continue
                fields = line.rstrip("\n").split()
                if len(fields) < 4:
                    continue
                bam, bam_md5, bai, bai_md5 = fields[:4]
                sample = sample_from_name(bam)
                pairs.append(
                    {
                        "repository": "NIST_GIAB",
                        "collection": path.parent.name,
                        "run_accession": Path(urllib.parse.urlparse(bam).path).name,
                        "sample_id": sample,
                        "reference_build": "GRCh37",
                        "platform_or_kit": technology,
                        "alignment_format": "BAM",
                        "alignment_url": https_from_ftpish(bam),
                        "index_url": https_from_ftpish(bai),
                        "alignment_bytes": None,
                        "index_bytes": None,
                        "alignment_md5": bam_md5,
                        "index_md5": bai_md5,
                        "origin_note": "NIST Genome in a Bottle official exome alignment index",
                    }
                )
    unique: dict[str, dict[str, Any]] = {p["alignment_url"]: p for p in pairs}
    return list(unique.values()), {"index_files": len(paths), "unique_alignment_objects": len(unique)}


def gcs_list(prefix: str) -> list[dict[str, Any]]:
    base = "https://storage.googleapis.com/storage/v1/b/brain-genomics-public/o"
    token = ""
    items: list[dict[str, Any]] = []
    while True:
        params = {"prefix": prefix, "maxResults": "1000"}
        if token:
            params["pageToken"] = token
        url = base + "?" + urllib.parse.urlencode(params)
        req = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
        with urllib.request.urlopen(req, timeout=60) as resp:
            payload = json.load(resp)
        items.extend(payload.get("items", []))
        token = payload.get("nextPageToken", "")
        if not token:
            return items


def load_google() -> tuple[list[dict[str, Any]], dict[str, Any]]:
    prefixes: list[str] = []
    for build in ("grch37", "grch38"):
        for platform in ("hiseq4000", "novaseq"):
            for kit in ("wes_agilent", "wes_idt", "wes_truseq"):
                prefixes.append(f"research/sequencing/{build}/bam/{platform}/{kit}/")
    objects: dict[str, dict[str, Any]] = {}
    for prefix in prefixes:
        for item in gcs_list(prefix):
            objects[item["name"]] = item
    pairs: list[dict[str, Any]] = []
    missing_index: list[str] = []
    for name, item in sorted(objects.items()):
        if not name.endswith(".bam"):
            continue
        bai_name = name + ".bai"
        if bai_name not in objects:
            alt = name[:-4] + ".bai"
            bai_name = alt if alt in objects else ""
        if not bai_name:
            missing_index.append(name)
            continue
        bai = objects[bai_name]
        parts = name.split("/")
        build, platform, kit = parts[2], parts[4], parts[5]
        pairs.append(
            {
                "repository": "Google_Brain_GIAB_benchmark",
                "collection": "Baid_et_al_GIAB_benchmark",
                "run_accession": Path(name).name,
                "sample_id": sample_from_name(name),
                "reference_build": build.upper(),
                "platform_or_kit": f"{platform}/{kit}",
                "alignment_format": "BAM",
                "alignment_url": gcs_https("brain-genomics-public", name),
                "index_url": gcs_https("brain-genomics-public", bai_name),
                "alignment_bytes": int(item["size"]),
                "index_bytes": int(bai["size"]),
                "alignment_md5": base64.b64decode(item.get("md5Hash", "")).hex() if item.get("md5Hash") else "",
                "index_md5": base64.b64decode(bai.get("md5Hash", "")).hex() if bai.get("md5Hash") else "",
                "origin_note": "Google public GIAB benchmark WES; same libraries aligned independently to two builds",
            }
        )
    return pairs, {
        "prefixes_queried": prefixes,
        "objects_in_prefixes": len(objects),
        "indexed_alignment_objects": len(pairs),
        "missing_index": missing_index,
    }


def load_gatk() -> tuple[list[dict[str, Any]], dict[str, Any]]:
    bam_key = "exome_bam/1000G_wex_hg38/HG00133.alt_bwamem_GRCh38DH.20150826.GBR.exome.bam"
    bai_key = bam_key[:-1] + "i"
    return [
        {
            "repository": "GATK_test_data_AWS",
            "collection": "1000G_wex_hg38",
            "run_accession": "HG00133",
            "sample_id": "HG00133",
            "reference_build": "GRCh38DH",
            "platform_or_kit": "1000G exome re-alignment",
            "alignment_format": "BAM",
            "alignment_url": f"https://gatk-test-data.s3.amazonaws.com/{bam_key}",
            "index_url": f"https://gatk-test-data.s3.amazonaws.com/{bai_key}",
            "alignment_bytes": 13_105_459_770,
            "index_bytes": 4_379_640,
            "alignment_md5": "",
            "index_md5": "",
            "origin_note": "Broad GATK public test-data; underlying sample is from 1000 Genomes",
        }
    ], {"indexed_alignment_objects": 1}


def range_probe(url: str, expected: int | None, attempts: int = 4) -> dict[str, Any]:
    last_error = ""
    for attempt in range(attempts):
        try:
            req = urllib.request.Request(
                url,
                headers={"Range": "bytes=0-0", "User-Agent": USER_AGENT, "Accept": "*/*"},
            )
            with urllib.request.urlopen(req, timeout=45) as resp:
                payload = resp.read(2)
                status = getattr(resp, "status", None) or resp.getcode()
                content_range = resp.headers.get("Content-Range", "")
                content_length = resp.headers.get("Content-Length", "")
                total = None
                m = re.search(r"/(\d+)$", content_range)
                if m:
                    total = int(m.group(1))
                elif status == 200 and content_length.isdigit():
                    total = int(content_length)
                size_match = expected is None or total is None or total == expected
                return {
                    "ok": status in (200, 206) and bool(payload) and size_match,
                    "status": status,
                    "content_range": content_range,
                    "observed_total_bytes": total,
                    "expected_bytes": expected,
                    "size_match": size_match,
                    "bytes_received": len(payload),
                    "etag": resp.headers.get("ETag", "").strip('"'),
                    "last_modified": resp.headers.get("Last-Modified", ""),
                    "attempts": attempt + 1,
                }
        except Exception as exc:  # urllib exposes several transport exception types
            last_error = f"{type(exc).__name__}: {exc}"
            if attempt + 1 < attempts:
                time.sleep(0.5 * (2**attempt))
    return {
        "ok": False,
        "status": None,
        "content_range": "",
        "observed_total_bytes": None,
        "expected_bytes": expected,
        "size_match": False,
        "bytes_received": 0,
        "etag": "",
        "last_modified": "",
        "attempts": attempts,
        "error": last_error,
    }


def verify_pairs(pairs: list[dict[str, Any]], workers: int, progress_path: Path) -> None:
    tasks: list[tuple[int, str, str, int | None]] = []
    for i, pair in enumerate(pairs):
        tasks.append((i, "alignment", pair["alignment_url"], pair.get("alignment_bytes")))
        tasks.append((i, "index", pair["index_url"], pair.get("index_bytes")))
    started = time.time()
    ok = 0
    done = 0
    with ThreadPoolExecutor(max_workers=workers) as pool:
        futures = {pool.submit(range_probe, url, expected): (i, kind) for i, kind, url, expected in tasks}
        for fut in as_completed(futures):
            i, kind = futures[fut]
            result = fut.result()
            pairs[i][f"{kind}_probe"] = result
            ok += int(result["ok"])
            done += 1
            if done % 250 == 0 or done == len(tasks):
                progress = {
                    "status": "running" if done < len(tasks) else "complete",
                    "tasks_done": done,
                    "tasks_total": len(tasks),
                    "tasks_ok": ok,
                    "elapsed_seconds": round(time.time() - started, 1),
                    "updated_at": datetime.now(timezone.utc).isoformat(),
                }
                progress_path.write_text(json.dumps(progress, indent=2), encoding="utf-8")
                with PRINT_LOCK:
                    print(json.dumps(progress), flush=True)


def write_outputs(
    pairs: list[dict[str, Any]],
    repo_metadata: dict[str, Any],
    out_csv: Path,
    out_json: Path,
    exclusions: list[dict[str, Any]],
) -> None:
    verified = [
        p
        for p in pairs
        if p.get("alignment_probe", {}).get("ok") and p.get("index_probe", {}).get("ok")
    ]
    columns = [
        "repository",
        "collection",
        "run_accession",
        "sample_id",
        "reference_build",
        "platform_or_kit",
        "alignment_format",
        "alignment_url",
        "index_url",
        "alignment_bytes",
        "index_bytes",
        "alignment_md5",
        "index_md5",
        "alignment_http_status",
        "index_http_status",
        "alignment_observed_bytes",
        "index_observed_bytes",
        "verified_downloadable",
        "origin_note",
    ]
    with out_csv.open("w", encoding="utf-8", newline="") as fh:
        writer = csv.DictWriter(fh, fieldnames=columns)
        writer.writeheader()
        for p in pairs:
            writer.writerow(
                {
                    **{k: p.get(k, "") for k in columns},
                    "alignment_http_status": p.get("alignment_probe", {}).get("status", ""),
                    "index_http_status": p.get("index_probe", {}).get("status", ""),
                    "alignment_observed_bytes": p.get("alignment_probe", {}).get("observed_total_bytes", ""),
                    "index_observed_bytes": p.get("index_probe", {}).get("observed_total_bytes", ""),
                    "verified_downloadable": bool(
                        p.get("alignment_probe", {}).get("ok") and p.get("index_probe", {}).get("ok")
                    ),
                }
            )
    repo_summary: dict[str, dict[str, Any]] = {}
    for repo in sorted({p["repository"] for p in pairs}):
        all_repo = [p for p in pairs if p["repository"] == repo]
        good_repo = [p for p in verified if p["repository"] == repo]
        sizes = [p["alignment_bytes"] for p in all_repo if p.get("alignment_bytes") is not None]
        repo_summary[repo] = {
            "candidate_alignment_objects": len(all_repo),
            "verified_alignment_and_index_pairs": len(good_repo),
            "unique_sample_identifiers": len({p["sample_id"] for p in all_repo}),
            "known_alignment_bytes_total": sum(sizes),
            "known_alignment_bytes_min": min(sizes) if sizes else None,
            "known_alignment_bytes_median": int(statistics.median(sizes)) if sizes else None,
            "known_alignment_bytes_max": max(sizes) if sizes else None,
            "metadata": repo_metadata.get(repo, {}),
        }
    failures = [
        {
            "repository": p["repository"],
            "alignment_url": p["alignment_url"],
            "index_url": p["index_url"],
            "alignment_probe": p.get("alignment_probe"),
            "index_probe": p.get("index_probe"),
        }
        for p in pairs
        if p not in verified
    ]
    payload = {
        "audit_date": DATE,
        "generated_at_utc": datetime.now(timezone.utc).isoformat(),
        "scope": "Public complete human WES BAM/CRAM alignment objects with matching BAI/CRAI",
        "verification_method": "Anonymous HTTP GET Range bytes=0-0 on every alignment and every index; expected sizes compared when supplied by the official inventory",
        "interpretation": {
            "file_not_person": "Counts are downloadable alignment objects, not necessarily unique people. ENA may have multiple runs per sample; Google provides the same libraries aligned to GRCh37 and GRCh38; GIAB includes technical replicates.",
            "not_clinical_truth": "Public WES availability is not equivalent to curated BRCA exon-CNV truth. Most files are useful for technical robustness, reference normalization, or specificity rather than sensitivity.",
        },
        "totals": {
            "candidate_alignment_objects": len(pairs),
            "verified_alignment_and_index_pairs": len(verified),
            "failed_pairs": len(failures),
            "unique_repository_labels": len(repo_summary),
            "known_alignment_bytes_total": sum(
                p["alignment_bytes"] for p in pairs if p.get("alignment_bytes") is not None
            ),
        },
        "repositories": repo_summary,
        "failures": failures,
        "excluded_or_not_counted_major_resources": exclusions,
        "artifacts": {"csv": out_csv.name},
    }
    out_json.write_text(json.dumps(payload, indent=2, sort_keys=False), encoding="utf-8")


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--analysis-dir", type=Path, required=True)
    parser.add_argument("--giab-index-root", type=Path, required=True)
    parser.add_argument("--workers", type=int, default=32)
    args = parser.parse_args()
    analysis = args.analysis_dir
    analysis.mkdir(parents=True, exist_ok=True)
    ena_path = analysis / f"ena_human_wxs_submitted_bam_inventory_{DATE}.tsv"
    depmap_path = analysis / f"depmap_omics_ccle_inventory_{DATE}.csv"
    if not ena_path.exists() or not depmap_path.exists():
        raise FileNotFoundError("Required official ENA and DepMap metadata inventories are missing")

    pairs: list[dict[str, Any]] = []
    metadata: dict[str, Any] = {}
    for label, loader in (
        ("ENA", lambda: load_ena(ena_path)),
        ("DepMap_CCLE_AWS", lambda: load_depmap(depmap_path)),
        ("NIST_GIAB", lambda: load_giab(args.giab_index_root)),
        ("Google_Brain_GIAB_benchmark", load_google),
        ("GATK_test_data_AWS", load_gatk),
    ):
        loaded, info = loader()
        pairs.extend(loaded)
        metadata[label] = info
        print(json.dumps({"loaded": label, "pairs": len(loaded), "metadata": info}), flush=True)

    # Deduplicate exact objects only; biological duplicates remain visible by design.
    unique = {(p["alignment_url"], p["index_url"]): p for p in pairs}
    pairs = list(unique.values())
    progress = analysis / f"public_complete_wes_alignment_probe_progress_{DATE}.json"
    verify_pairs(pairs, max(1, args.workers), progress)

    exclusions = [
        {"repository": "GDC/TCGA and BeatAML", "anonymous_indexed_WES": 0, "reason": "Individual-level BAM/FASTQ aligned reads are controlled access through dbGaP."},
        {"repository": "ICGC/PCAWG", "anonymous_indexed_WES": 0, "reason": "Primary BAMs are credentialed/controlled; public bucket contains metadata or aggregate results."},
        {"repository": "gnomAD", "anonymous_indexed_WES": 0, "reason": "Public release is aggregate summary data, not individual exome alignments."},
        {"repository": "UK Biobank", "anonymous_indexed_WES": 0, "reason": "Approved-project access in UKB-RAP; WES/WGS download is not permitted under current policy."},
        {"repository": "All of Us", "anonymous_indexed_WES": 0, "reason": "Genomics is Controlled Tier and current primary sequencing is WGS, not a public WES BAM collection."},
        {"repository": "NHLBI ESP/dbGaP", "anonymous_indexed_WES": 0, "reason": "Individual sequence data are controlled; public Exome Variant Server products are aggregate variants."},
        {"repository": "GenomeAsia 100K", "anonymous_indexed_WES": 0, "reason": "Individual data require a data-access agreement/EGA credentials; public browser provides variants."},
        {"repository": "HGDP/SGDP", "anonymous_indexed_WES": 0, "reason": "Modern public aligned-read releases are WGS rather than WES; restricted samples require approval."},
        {"repository": "ENCODE", "anonymous_indexed_WES": 0, "reason": "Functional-genomics portal; no released whole-exome alignment collection identified."},
        {"repository": "SRA normalized archive", "anonymous_indexed_WES": 0, "reason": "Public WES runs exist as SRA/raw data, but normalized .sra objects are not ready indexed BAM/CRAM pairs; ENA submitted BAM/BAI pairs are counted separately."},
    ]
    out_csv = analysis / f"public_complete_wes_alignment_inventory_verified_{DATE}.csv"
    out_json = analysis / f"public_complete_wes_repository_audit_{DATE}.json"
    write_outputs(pairs, metadata, out_csv, out_json, exclusions)
    progress.write_text(
        json.dumps(
            {
                "status": "complete",
                "candidate_pairs": len(pairs),
                "csv": str(out_csv),
                "json": str(out_json),
                "completed_at": datetime.now(timezone.utc).isoformat(),
            },
            indent=2,
        ),
        encoding="utf-8",
    )
    print(json.dumps({"status": "complete", "pairs": len(pairs), "csv": str(out_csv), "json": str(out_json)}))
    return 0


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