"""Parse training log file and write metrics to TensorBoard event files.

Usage:
    # One-shot conversion
    python scripts/log_to_tensorboard.py --log /tmp/streampi_t3_robodojo_30k_v1_20260909.log

    # Watch mode (continuously monitor log for new entries)
    python scripts/log_to_tensorboard.py --log /tmp/streampi_t3_robodojo_30k_v1_20260909.log --watch

    # Custom output directory
    python scripts/log_to_tensorboard.py --log /tmp/run.log --output-dir runs/my_run
"""

import argparse
import os
import re
import time
from pathlib import Path

try:
  from torch.utils.tensorboard import SummaryWriter
except ImportError:
  from tensorboardX import SummaryWriter


def parse_step_line(line: str) -> dict | None:
    """Parse a log line like 'Step 0: grad_norm=4.0662, loss=0.4126, param_norm=1802.3864'.

    Returns a dict with 'step' (int) and metric key-value pairs (float), or None.
    """
    match = re.match(r"Step\s+(\d+):\s*(.+)", line.strip())
    if not match:
        return None

    step = int(match.group(1))
    metrics_str = match.group(2)

    result: dict = {"step": step}
    for pair in metrics_str.split(","):
        pair = pair.strip()
        if "=" not in pair:
            continue
        key, value = pair.split("=", 1)
        key = key.strip()
        value = value.strip()
        try:
            result[key] = float(value)
        except ValueError:
            continue

    return result


def parse_log_file(log_path: str, last_step: int = -1) -> list[dict]:
    """Return deduplicated entries with step > last_step.

    If a resumed run repeats a step, the last occurrence in the log is kept,
    because it belongs to the trajectory resumed from the latest checkpoint.
    """
    entries_by_step: dict[int, dict] = {}
    with open(log_path, "r", encoding="utf-8", errors="ignore") as f:
        for line in f:
            parsed = parse_step_line(line)
            if parsed is None:
                continue

            step = parsed["step"]
            if step <= last_step:
                continue

            # Later occurrences replace pre-resume entries at the same step.
            entries_by_step[step] = parsed

    return [entries_by_step[step] for step in sorted(entries_by_step)]


def write_entries(writer: SummaryWriter, entries: list[dict]) -> int:
    """Write parsed entries to TensorBoard. Returns the max step written."""
    max_step = -1
    for entry in entries:
        step = entry["step"]
        for key, value in entry.items():
            if key == "step":
                continue
            writer.add_scalar(key, value, step)
        max_step = max(max_step, step)

    writer.flush()
    return max_step


def main():
    parser = argparse.ArgumentParser(description="Convert training log to TensorBoard events.")
    parser.add_argument("--log", type=str, required=True, help="Path to the training log file.")
    parser.add_argument(
        "--output-dir",
        type=str,
        default=None,
        help="TensorBoard output directory. Defaults to runs/<log_filename_stem>.",
    )
    parser.add_argument(
        "--watch",
        action="store_true",
        help="Continuously monitor the log file for new entries.",
    )
    parser.add_argument(
        "--interval",
        type=float,
        default=10.0,
        help="Polling interval in seconds for --watch mode (default: 10).",
    )
    args = parser.parse_args()

    log_path = args.log
    if not os.path.exists(log_path):
        raise FileNotFoundError(f"Log file not found: {log_path}")

    log_stem = Path(log_path).stem
    output_dir = args.output_dir or os.path.join("runs", log_stem)
    os.makedirs(output_dir, exist_ok=True)

    writer = SummaryWriter(log_dir=output_dir)
    print(f"TensorBoard log dir: {os.path.abspath(output_dir)}")
    print(f"Monitoring log: {log_path}")

    # One-shot parse
    entries = parse_log_file(log_path)
    max_step = write_entries(writer, entries)
    print(f"Wrote {len(entries)} entries (max step: {max_step})")

    if not args.watch:
        writer.close()
        print(f"\nDone. Run: tensorboard --logdir {os.path.abspath(output_dir)}")
        return

    # Watch mode
    print(f"Watching for new entries (interval={args.interval}s). Press Ctrl+C to stop.")
    try:
        while True:
            time.sleep(args.interval)
            new_entries = parse_log_file(log_path, last_step=max_step)
            if new_entries:
                max_step = write_entries(writer, new_entries)
                print(f"  Updated: +{len(new_entries)} entries (max step: {max_step})")
    except KeyboardInterrupt:
        print("\nStopping watch mode.")
    finally:
        writer.close()
        print(f"TensorBoard events written to: {os.path.abspath(output_dir)}")
        print(f"Run: tensorboard --logdir {os.path.abspath(output_dir)}")


if __name__ == "__main__":
    main()