from __future__ import annotations

import argparse
from datetime import datetime
from pathlib import Path
import subprocess
import sys


def sanitize_name(name: str) -> str:
    return "".join(ch if ch.isalnum() or ch in ("-", "_") else "_" for ch in name).strip("_")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Run a command and tee stdout/stderr into outputs/logs for submission evidence."
    )
    parser.add_argument("--name", required=True, help="Short log name, e.g. unit_tests or grpo_qwen3_0_6b")
    parser.add_argument("command", nargs=argparse.REMAINDER, help="Command to run after --")
    args = parser.parse_args()

    command = args.command
    if command and command[0] == "--":
        command = command[1:]
    if not command:
        parser.error("Provide a command after --")

    log_dir = Path("outputs/logs")
    log_dir.mkdir(parents=True, exist_ok=True)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    log_path = log_dir / f"{timestamp}_{sanitize_name(args.name)}.log"

    with log_path.open("w", encoding="utf-8") as log_file:
        header = [
            f"# Log: {args.name}",
            f"# Timestamp: {timestamp}",
            f"# Command: {' '.join(command)}",
            "",
        ]
        log_file.write("\n".join(header))
        log_file.flush()
        print(f"[log] writing to {log_path}")

        process = subprocess.Popen(
            command,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            encoding="utf-8",
            errors="replace",
        )
        assert process.stdout is not None
        for line in process.stdout:
            print(line, end="")
            log_file.write(line)
            log_file.flush()
        return_code = process.wait()
        log_file.write(f"\n# Exit code: {return_code}\n")

    print(f"[log] saved {log_path}")
    raise SystemExit(return_code)


if __name__ == "__main__":
    main()