#!/usr/bin/env python3
"""Fetch official ENA titles for every study in the WXS BAM/CRAM inventory."""

from __future__ import annotations

import csv
import io
import json
import sys
import time
import urllib.parse
import urllib.request
from collections import defaultdict
from datetime import datetime, timezone
from pathlib import Path


DATE = "2026-08-01"
API = "https://www.ebi.ac.uk/ena/portal/api/search"


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


def main() -> int:
    if len(sys.argv) != 2:
        raise SystemExit("usage: fetch_ena_wxs_study_titles.py ANALYSIS_DIR")
    analysis = Path(sys.argv[1])
    sources = {
        "BAM": analysis / f"ena_human_wxs_submitted_bam_inventory_{DATE}.tsv",
        "CRAM": analysis / f"ena_human_wxs_submitted_cram_inventory_{DATE}.tsv",
    }

    stats: dict[str, dict[str, object]] = defaultdict(
        lambda: {
            "bam_run_rows": 0,
            "bam_indexed_pairs": 0,
            "bam_without_submitted_index": 0,
            "cram_indexed_pairs": 0,
            "sample_accessions": set(),
        }
    )
    for kind, path in sources.items():
        with path.open("r", encoding="utf-8", newline="") as fh:
            for row in csv.DictReader(fh, delimiter="\t"):
                study = row["study_accession"]
                record = stats[study]
                record["sample_accessions"].add(row["sample_accession"])  # type: ignore[union-attr]
                formats = [part.upper() for part in semis(row["submitted_format"])]
                if kind == "BAM":
                    record["bam_run_rows"] = int(record["bam_run_rows"]) + 1
                    bams = formats.count("BAM")
                    bais = formats.count("BAI")
                    indexed = min(bams, bais)
                    record["bam_indexed_pairs"] = int(record["bam_indexed_pairs"]) + indexed
                    record["bam_without_submitted_index"] = int(record["bam_without_submitted_index"]) + (bams - indexed)
                else:
                    record["cram_indexed_pairs"] = int(record["cram_indexed_pairs"]) + min(
                        formats.count("CRAM"), formats.count("CRAI")
                    )

    accessions = sorted(stats)
    official: dict[str, dict[str, str]] = {}
    request_urls: list[str] = []
    for offset in range(0, len(accessions), 40):
        batch = accessions[offset : offset + 40]
        query = " OR ".join(f'study_accession="{accession}"' for accession in batch)
        params = {
            "result": "study",
            "query": query,
            "fields": "study_accession,secondary_study_accession,study_title,study_description",
            "format": "tsv",
            "limit": "0",
        }
        url = API + "?" + urllib.parse.urlencode(params)
        request_urls.append(url)
        last_error = ""
        for attempt in range(5):
            try:
                request = urllib.request.Request(url, headers={"User-Agent": "LGASieve-ENA-study-audit/1.0"})
                with urllib.request.urlopen(request, timeout=90) as response:
                    body = response.read().decode("utf-8")
                for row in csv.DictReader(io.StringIO(body), delimiter="\t"):
                    official[row["study_accession"]] = row
                break
            except Exception as exc:
                last_error = f"{type(exc).__name__}: {exc}"
                if attempt == 4:
                    raise RuntimeError(f"ENA batch failed after five attempts: {last_error}") from exc
                time.sleep(1.5 * (2**attempt))

    fields = [
        "study_accession", "secondary_study_accession", "study_title", "study_description",
        "bam_run_rows", "bam_indexed_pairs", "bam_without_submitted_index",
        "cram_indexed_pairs", "indexed_alignment_pairs_total", "unique_sample_accessions",
        "official_study_url",
    ]
    output = analysis / f"ena_human_wxs_study_summary_{DATE}.csv"
    with output.open("w", encoding="utf-8", newline="") as fh:
        writer = csv.DictWriter(fh, fieldnames=fields)
        writer.writeheader()
        for accession in accessions:
            record = stats[accession]
            metadata = official.get(accession, {})
            writer.writerow(
                {
                    "study_accession": accession,
                    "secondary_study_accession": metadata.get("secondary_study_accession", ""),
                    "study_title": metadata.get("study_title", ""),
                    "study_description": metadata.get("study_description", ""),
                    "bam_run_rows": record["bam_run_rows"],
                    "bam_indexed_pairs": record["bam_indexed_pairs"],
                    "bam_without_submitted_index": record["bam_without_submitted_index"],
                    "cram_indexed_pairs": record["cram_indexed_pairs"],
                    "indexed_alignment_pairs_total": int(record["bam_indexed_pairs"]) + int(record["cram_indexed_pairs"]),
                    "unique_sample_accessions": len(record["sample_accessions"]),  # type: ignore[arg-type]
                    "official_study_url": f"https://www.ebi.ac.uk/ena/browser/view/{accession}",
                }
            )

    report = {
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "source": "official ENA Portal API result=study",
        "studies_requested": len(accessions),
        "studies_returned": len(official),
        "studies_missing_metadata": sorted(set(accessions) - set(official)),
        "api_request_count": len(request_urls),
        "output_csv": output.name,
    }
    report_path = output.with_suffix(".json")
    report_path.write_text(json.dumps(report, indent=2), encoding="utf-8")
    print(json.dumps(report), flush=True)
    return 0


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