from __future__ import annotations

import argparse
import json
import resource
import time
from pathlib import Path

import numpy as np
import skops.io as sio
from threadpoolctl import threadpool_limits

from banking_intent_error_predictor.training import (
    SOURCE_REVISION,
    download,
    fit_once,
    load_rows,
    sha256_file,
)


def main() -> None:
    parser = argparse.ArgumentParser(description="Reproduce the reviewed release")
    parser.add_argument("--cache-dir", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--reference", type=Path, required=True)
    args = parser.parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=False)
    train_rows = load_rows(download(args.cache_dir, "train.csv"))
    test_rows = load_rows(download(args.cache_dir, "test.csv"))
    download(args.cache_dir, "categories.json")
    wall_start = time.perf_counter()
    cpu_start = time.process_time()
    with threadpool_limits(limits=8):
        result = fit_once(train_rows, test_rows)
    reference = result["reproduction_reference"]
    with np.load(args.reference, allow_pickle=False) as expected:
        exact = all(
            np.array_equal(reference[key], expected[key]) for key in expected.files
        )
    if not exact:
        raise RuntimeError("reproduced decision scores differ from v1.0.0")
    primary_path = args.output_dir / "primary_baseline.skops"
    candidate_path = args.output_dir / "model.skops"
    sio.dump(
        {
            "features": result["models"]["feature_extractor"],
            "classifier": result["models"]["primary"],
            "labels": result["labels"],
        },
        primary_path,
    )
    sio.dump(
        {
            "classifier": result["models"]["candidate"],
            "labels": result["labels"],
            "review_rate": 0.20,
        },
        candidate_path,
    )
    for path in (primary_path, candidate_path):
        if sio.get_untrusted_types(file=path):
            raise RuntimeError(f"reproduced artifact requires untrusted types: {path}")
    report = {
        "dataset_revision": SOURCE_REVISION,
        "reference_outputs_exact": exact,
        "primary_artifact_sha256": sha256_file(primary_path),
        "candidate_artifact_sha256": sha256_file(candidate_path),
        "wall_seconds": time.perf_counter() - wall_start,
        "cpu_seconds": time.process_time() - cpu_start,
        "thread_limit": 8,
        "peak_rss_bytes": int(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss),
    }
    (args.output_dir / "reproduction.json").write_text(
        json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    print(json.dumps(report, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()