import argparse
import json
import os
import subprocess
import time
from datetime import datetime, timezone
from pathlib import Path

def get_utc_now_iso():
    return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")

def main():
    parser = argparse.ArgumentParser(description="Mission Control Terminal Tracker")
    parser.add_argument("cmd", help="The command to execute and track")
    parser.add_argument("--agent", default="user", help="The agent executing the command (e.g. codex, gemini, claude)")
    args = parser.parse_args()

    # Define paths
    repo_root = Path(__file__).resolve().parent.parent
    runs_dir = repo_root / "tools" / "story-runner" / "runs" / "terminal"
    
    # Create unique run ID
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    slug = "".join(c if c.isalnum() else "_" for c in args.cmd[:20]).strip("_")
    run_id = f"{timestamp}_{slug}"
    run_path = runs_dir / run_id
    run_path.mkdir(parents=True, exist_ok=True)
    
    status_file = run_path / "status.json"
    
    status_data = {
        "type": "terminal",
        "command": args.cmd,
        "agent": args.agent,
        "overall_status": "running",
        "start_time_utc": get_utc_now_iso(),
        "last_heartbeat_utc": get_utc_now_iso(),
        "exit_code": None
    }

    with open(status_file, "w", encoding="utf-8") as f:
        json.dump(status_data, f, indent=2)

    print(f"--- MC TRACK: Launching '{args.cmd}' (Agent: {args.agent}) ---")
    
    start_time = time.time()
    try:
        process = subprocess.Popen(args.cmd, shell=True)
        
        while process.poll() is None:
            time.sleep(1)
            status_data["last_heartbeat_utc"] = get_utc_now_iso()
            with open(status_file, "w", encoding="utf-8") as f:
                json.dump(status_data, f, indent=2)
        
        exit_code = process.returncode
    except Exception as e:
        print(f"--- MC TRACK: Execution Error: {e} ---")
        exit_code = -1

    end_time = time.time()
    duration = round(end_time - start_time, 2)

    status_data["overall_status"] = "done" if exit_code == 0 else "failed"
    status_data["exit_code"] = exit_code
    status_data["end_time_utc"] = get_utc_now_iso()
    status_data["duration_seconds"] = duration

    with open(status_file, "w", encoding="utf-8") as f:
        json.dump(status_data, f, indent=2)

    status_color = "\033[92m" if exit_code == 0 else "\033[91m"
    reset_color = "\033[0m"
    print(f"--- MC TRACK: Finished in {duration}s with code {exit_code} ({status_color}{status_data['overall_status']}{reset_color}) ---")

if __name__ == "__main__":
    main()