import argparse
import json
from pathlib import Path

import torch
from bit_transformer import BitTransformerLM


def list_candidates(path: Path):
    models = sorted(path.glob("*.pt"))
    for m in models:
        metrics_file = m.with_suffix(m.suffix + ".json")
        metrics = {}
        if metrics_file.exists():
            with open(metrics_file) as f:
                metrics = json.load(f)
        yield m, metrics


def main():
    parser = argparse.ArgumentParser(description="Review distilled submodels")
    parser.add_argument("directory", type=Path, help="Directory with candidate models")
    parser.add_argument("--approve-dir", type=Path, default=Path("approved"), help="Directory to store approved models")
    args = parser.parse_args()

    args.approve_dir.mkdir(exist_ok=True)
    log_file = args.approve_dir / "review_log.jsonl"

    for model_path, metrics in list_candidates(args.directory):
        print("Candidate:", model_path.name)
        for k, v in metrics.items():
            print(f"  {k}: {v}")
        ans = input("Approve this model? [y/N] ").strip().lower()
        if ans == "y":
            approved_path = args.approve_dir / model_path.name
            torch.save(torch.load(model_path), approved_path)
            entry = {"model": approved_path.name, "metrics": metrics, "approved": True}
            with open(log_file, "a") as lf:
                lf.write(json.dumps(entry) + "\n")
            print("Approved and saved to", approved_path)
        else:
            entry = {"model": model_path.name, "metrics": metrics, "approved": False}
            with open(log_file, "a") as lf:
                lf.write(json.dumps(entry) + "\n")
            print("Rejected", model_path.name)


if __name__ == "__main__":
    main()