#!/usr/bin/env python3
"""
Check expected classifier result artifacts.

Usage:
    python tools/artifact_chk.py
    python tools/artifact_chk.py --output-root outputs
"""
import argparse
import json
from pathlib import Path


def parse_args():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--output-root", default="outputs", help="Root with logs/models folders")
    return p.parse_args()


def iter_config_paths(config_root: Path):
    for sub in ("phase1", "phase2"):
        yield from sorted((config_root / sub).glob("*.json"))


def main():
    args = parse_args()
    root = Path(__file__).resolve().parent.parent
    config_root = root / "configs"
    logs_dir = root / args.output_root / "logs"
    models_dir = root / args.output_root / "models"

    expected = []
    for cfg_path in iter_config_paths(config_root):
        with open(cfg_path) as f:
            cfg = json.load(f)
        run_name = cfg.get("run_name")
        if run_name:
            expected.append((run_name, cfg_path))

    missing_logs = []
    missing_models = []
    for run_name, cfg_path in expected:
        if not (logs_dir / f"{run_name}.json").exists():
            missing_logs.append((run_name, cfg_path))
        if not any(models_dir.glob(f"{run_name}_fold*_best.pt")):
            missing_models.append((run_name, cfg_path))

    print(f"Expected runs from configs: {len(expected)}")
    print(f"Missing logs: {len(missing_logs)}")
    for run_name, cfg_path in missing_logs:
        print(f"  - {run_name} ({cfg_path.relative_to(root)})")
    print(f"Missing checkpoints: {len(missing_models)}")
    for run_name, cfg_path in missing_models:
        print(f"  - {run_name} ({cfg_path.relative_to(root)})")

    if missing_logs or missing_models:
        raise SystemExit(1)
    print("All expected artifacts found.")


if __name__ == "__main__":
    main()