"""Pack the two benchmark sets (start, end and reference TS structures) into data/inputs.tar.gz.

    python reproduce/pack_inputs.py --sets <dir with the two set tables> \
                                    --structures <dir with kaist_docked/, geodesic_gxtb/>

Only the reactions listed in the two CSVs are packed. The archive is byte-reproducible
(sorted entries, fixed timestamps and owners). Unpacked inside rerun/, it puts the set
tables in sets/ and the structures in data/reaction_qm/, where the scripts look.
"""

from __future__ import annotations

import argparse
import csv
import gzip
import io
import tarfile
from pathlib import Path

OUT = Path(__file__).resolve().parent.parent / "data" / "inputs.tar.gz"
# source table -> name in the archive
SETS = {"kaist_docked_500.csv": "reaction_qm_500.csv", "geodesic_gxtb.csv": "geodesic_66.csv"}


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--sets", type=Path, required=True)
    ap.add_argument("--structures", type=Path, required=True)
    a = ap.parse_args()
    files: dict[str, bytes] = {}
    for name, packed in SETS.items():
        text = (a.sets / name).read_bytes()
        files[f"sets/{packed}"] = text
        for row in csv.DictReader(io.StringIO(text.decode())):
            for col in ("start_xyz", "end_xyz", "ts_reference_xyz"):
                files[f"data/reaction_qm/{row[col]}"] = (a.structures / row[col]).read_bytes()
            placement = (a.structures / row["start_xyz"]).parent / "placement.json"
            if placement.exists():
                files[f"data/reaction_qm/{Path(row['start_xyz']).parent}/placement.json"] = placement.read_bytes()
    buf = io.BytesIO()
    with tarfile.open(fileobj=buf, mode="w", format=tarfile.PAX_FORMAT) as tar:
        for name in sorted(files):
            info = tarfile.TarInfo(name)
            info.size, info.mtime, info.mode = len(files[name]), 0, 0o644
            tar.addfile(info, io.BytesIO(files[name]))
    with gzip.GzipFile(OUT, "wb", mtime=0) as gz:
        gz.write(buf.getvalue())
    print(f"{OUT.name}: {len(files)} files, {OUT.stat().st_size / 1e6:.1f} MB")


if __name__ == "__main__":
    main()