"""Download and sample the two requested noncommercial research datasets."""

import argparse
import csv
import hashlib
import io
import json
import random
import urllib.request
import zipfile
from datetime import UTC, datetime
from pathlib import Path

from dataset import write_jsonl

SOURCES = {
    "saml_d": {
        "ref": "berkanoztas/synthetic-transaction-monitoring-dataset-aml",
        "member": "SAML-D.csv",
        "label": "Is_laundering",
        "authors": "Berkan Oztas et al.",
        "description": "Synthetic AML transactions; laundering labels are not private-share fraud labels.",
    },
    "baf": {
        "ref": "sgpjesus/bank-account-fraud-dataset-neurips-2022",
        "member": "Base.csv",
        "label": "fraud_bool",
        "authors": "Sergio Jesus, Jose Pombal, Duarte Alves, Andre Cruz, Pedro Saleiro, Rita P. Ribeiro, Joao Gama, Pedro Bizarro; Feedzai",
        "description": "Synthetic bank account applications; application fraud is not share ownership fraud.",
    },
}


def sample_archive(path, member, label_column, per_class, seed):
    """Stratified reservoir sample across the whole file, not its first rows."""
    rng = random.Random(seed)
    samples, counts = {"0": [], "1": []}, {"0": 0, "1": 0}
    with zipfile.ZipFile(path) as archive, archive.open(member) as stream:
        reader = csv.DictReader(io.TextIOWrapper(stream))
        for index, row in enumerate(reader, start=1):
            label = row[label_column]
            if label not in samples:
                raise ValueError(f"Unexpected label {label!r} at CSV data row {index}")
            counts[label] += 1
            candidate = {"csv_data_row": index, "original_label": label, "fields": row}
            if len(samples[label]) < per_class:
                samples[label].append(candidate)
            else:
                slot = rng.randrange(counts[label])
                if slot < per_class:
                    samples[label][slot] = candidate
    if any(len(rows) != per_class for rows in samples.values()):
        raise ValueError("Source does not contain enough records in both classes")
    return samples["0"] + samples["1"], counts


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--per-class", type=int, default=100)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--output", default="data/source-samples-001")
    args = parser.parse_args()
    if args.per_class <= 0:
        raise ValueError("per-class must be positive")
    output = Path(args.output)
    if output.exists():
        raise FileExistsError("Use a new output directory")
    raw = Path("data/raw")
    raw.mkdir(parents=True, exist_ok=True)
    records, manifest = (
        [],
        {
            "scope": "noncommercial research",
            "seed": args.seed,
            "sample_per_class": args.per_class,
            "created_at": datetime.now(UTC).isoformat(),
            "sources": {},
        },
    )
    for source, spec in SOURCES.items():
        archive = raw / f"{source}.zip"
        url = "https://www.kaggle.com/api/v1/datasets/download/" + spec["ref"]
        if not archive.exists():
            partial = archive.with_suffix(".partial")
            urllib.request.urlretrieve(url, partial)
            partial.rename(archive)
        with archive.open("rb") as stream:
            checksum = hashlib.file_digest(stream, "sha256").hexdigest()
        rows, counts = sample_archive(
            archive, spec["member"], spec["label"], args.per_class, args.seed
        )
        for row in rows:
            row.update(
                source=source,
                original_id=f"{spec['member']}:{row['csv_data_row']}",
                archive_sha256=checksum,
                dataset_license="CC-BY-NC-SA-4.0",
                label_semantics=spec["description"],
                adaptation_status="unadapted; original source label only",
            )
        records.extend(rows)
        manifest["sources"][source] = {
            **spec,
            "download_url": url,
            "source_url": "https://www.kaggle.com/datasets/" + spec["ref"],
            "license": "CC-BY-NC-SA-4.0",
            "license_url": "https://creativecommons.org/licenses/by-nc-sa/4.0/",
            "archive_sha256": checksum,
            "source_label_counts": counts,
            "sample_label_counts": {"0": args.per_class, "1": args.per_class},
            "caveat": "Balanced sample does not represent source prevalence; not yet mapped to share-review cases.",
        }
        print(
            source, "full-file label counts", counts, "sampled", len(rows), flush=True
        )
    output.mkdir(parents=True)
    write_jsonl(output / "source_records.jsonl", records)
    manifest["source_records_sha256"] = hashlib.sha256(
        (output / "source_records.jsonl").read_bytes()
    ).hexdigest()
    (output / "manifest.json").write_text(json.dumps(manifest, indent=2))
    print("Saved", len(records), "source records to", output)


if __name__ == "__main__":
    main()