#!/usr/bin/env python3
"""Generate standard Qwen SFT summary plots from captured VERL logs."""

from __future__ import annotations

import argparse
from pathlib import Path

from plot import make_plot, parse_log, write_csv


DEFAULT_RUNS = [
    (
        "qwen2_5_7b_lora16_sft",
        "Qwen2.5 7B LoRA-16 SFT",
        Path("plotting/logs/qwen2_5_7b-lora16-hep-sft.log"),
    ),
]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Plot the standard Qwen 7B LoRA VERL SFT training log.")
    parser.add_argument(
        "--out-dir",
        type=Path,
        default=Path("plotting"),
        help="Directory for generated PNG and CSV files.",
    )
    parser.add_argument(
        "--logs-dir",
        type=Path,
        default=Path("plotting/logs"),
        help="Directory containing qwen2_5_7b-lora16-hep-sft.log.",
    )
    parser.add_argument(
        "--csv",
        action="store_true",
        help="Also write parsed metrics CSV files.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    made_any = False

    for slug, title, default_log in DEFAULT_RUNS:
        log_path = args.logs_dir / default_log.name
        if not log_path.exists():
            print(f"Skipping missing log: {log_path}")
            continue

        rows = parse_log(log_path)
        if not rows:
            print(f"Skipping {log_path}: no `step:` metric lines found")
            continue

        png_path = args.out_dir / f"{slug}_summary.png"
        make_plot(rows, log_path, png_path, title)
        print(f"Wrote {png_path}")
        made_any = True

        if args.csv:
            csv_path = args.out_dir / f"{slug}_metrics.csv"
            write_csv(rows, csv_path)
            print(f"Wrote {csv_path}")

    if not made_any:
        raise SystemExit(
            "No plots were generated. Put captured VERL console logs under plotting/logs/ first."
        )


if __name__ == "__main__":
    main()