import argparse
import csv
import re
from pathlib import Path


METRIC_PATTERN = re.compile(
    r"^\s*step\s+(?P<step>[\d,]+)\s+"
    r"tokens\s+(?P<tokens>[\d,]+)\s+"
    r"loss\s+(?P<loss>[-+]?\d+(?:\.\d+)?)\s+"
    r".*?\bwall\s+(?P<wall>[-+]?\d+(?:\.\d+)?)s\b"
)


def parse_log(path: Path) -> list[dict[str, object]]:
    rows = []
    with path.open(encoding="utf-8", errors="replace") as log_file:
        for line in log_file:
            match = METRIC_PATTERN.search(line)
            if match is None:
                continue
            rows.append(
                {
                    "run": path.parent.name,
                    "step": int(match.group("step").replace(",", "")),
                    "tokens": int(match.group("tokens").replace(",", "")),
                    "loss": float(match.group("loss")),
                    "wall_time_seconds": float(match.group("wall")),
                }
            )
    if not rows:
        raise ValueError(f"No metric rows found in {path}")
    return rows


def main() -> None:
    argument_parser = argparse.ArgumentParser()
    argument_parser.add_argument("logs", nargs="+", type=Path)
    argument_parser.add_argument("--output", "-o", required=True, type=Path)
    args = argument_parser.parse_args()

    rows = []
    for log_path in args.logs:
        rows.extend(parse_log(log_path))

    fieldnames = ["run", "step", "tokens", "loss", "wall_time_seconds"]
    with args.output.open("w", newline="", encoding="utf-8") as output_file:
        writer = csv.DictWriter(output_file, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)

    for log_path in args.logs:
        log_rows = [row for row in rows if row["run"] == log_path.parent.name]
        print(f"{log_path}: {len(log_rows)} metric rows")
    print(f"Wrote {len(rows)} rows to {args.output}")


if __name__ == "__main__":
    main()