#!/usr/bin/env python3
"""FSI_FELON Training Monitor — watches over the forge while you sleep."""
import os, time, json, sys, subprocess
from datetime import datetime

LOG = "/tmp/opencode/snca/train_live.log"
SCRIPT = "/tmp/opencode/snca/train_live.py"
SCREEN = "deep_train"
CHECK_INTERVAL = 600  # 10 min (~100 steps)
REPORT = "/tmp/fsi_felon/training_report.json"
ALERT_FILE = "/tmp/fsi_felon/training_alert"

def tail_lines(path, n=5):
    try:
        with open(path) as f:
            lines = f.readlines()
        return lines[-n:]
    except:
        return []

def check_training():
    now = datetime.now().strftime("%H:%M:%S")
    lines = tail_lines(LOG, 3)
    if not lines:
        return {"status": "NO_LOG", "time": now, "msg": "No log file found"}

    last = lines[-1].strip()
    # Parse: [03:14:59] step=1,457 loss=7.3952 lr=3.00e-04
    try:
        parts = last.split()
        time_str = parts[0].strip("[]")
        step = int(parts[1].split("=")[1].replace(",", ""))
        loss = float(parts[2].split("=")[1])
        lr = float(parts[3].split("=")[1])
    except:
        return {"status": "PARSE_ERROR", "time": now, "msg": f"Could not parse: {last}"}

    report = {
        "status": "OK",
        "time": now,
        "step": step,
        "loss": loss,
        "lr": lr,
        "msg": f"step={step:,} loss={loss:.4f} lr={lr:.2e}"
    }

    # Check for anomalies
    if loss > 15.0:
        report["alert"] = f"Loss spike: {loss:.4f}"
        report["status"] = "ALERT"
        with open(ALERT_FILE, "w") as f:
            f.write(json.dumps(report))

    # Check if training stopped (log not updating)
    if os.path.exists(REPORT):
        try:
            with open(REPORT) as f:
                prev = json.load(f)
            if prev.get("step") == step and prev.get("status") != "RESTARTED":
                report["status"] = "STALLED"
                report["msg"] += " — LOG NOT UPDATING"
        except:
            pass

    with open(REPORT, "w") as f:
        json.dump(report, f, indent=2)
    return report

def restart_training():
    print(f"[{datetime.now().strftime('%H:%M:%S')}] Training stalled or crashed — restarting...")
    # Kill existing screen
    subprocess.run(["screen", "-S", SCREEN, "-X", "quit"], capture_output=True)
    time.sleep(2)
    # Start fresh
    cmd = f"cd /tmp/opencode/snca && screen -dmS {SCREEN} python3 train_live.py"
    subprocess.run(cmd, shell=True)
    time.sleep(5)
    report = check_training()
    report["status"] = "RESTARTED"
    with open(REPORT, "w") as f:
        json.dump(report, f, indent=2)
    print(f"  Restarted at step {report.get('step', '?')}")
    return report

def monitor():
    print("=" * 55)
    print("  FSI_FELON TRAINING MONITOR")
    print("  Checking every 10 minutes (~100 steps)")
    print(f"  Started: {datetime.now().strftime('%H:%M:%S')}")
    print("=" * 55)
    print()

    # Clear any previous alert
    if os.path.exists(ALERT_FILE):
        os.remove(ALERT_FILE)

    cycle = 0
    while True:
        cycle += 1
        report = check_training()
        ts = report["time"]
        status = report["status"]

        if status == "OK":
            print(f"[{ts}] ✓ step={report['step']:,} loss={report['loss']:.4f} lr={report['lr']:.2e}")
        elif status == "STALLED":
            print(f"[{ts}] ⚠ STALLED — restarting")
            report = restart_training()
        elif status == "ALERT":
            print(f"[{ts}] 🔴 ALERT: {report.get('alert')}")
        elif status == "NO_LOG":
            print(f"[{ts}] 🔴 No log — restarting")
            report = restart_training()
        else:
            print(f"[{ts}] ? {status}: {report.get('msg')}")

        # Check screen is alive
        r = subprocess.run(["screen", "-ls"], capture_output=True, text=True)
        if SCREEN not in r.stdout:
            print(f"[{datetime.now().strftime('%H:%M:%S')}] 🔴 Screen dead — restarting")
            restart_training()

        time.sleep(CHECK_INTERVAL)

if __name__ == "__main__":
    try:
        monitor()
    except KeyboardInterrupt:
        print("\nMonitor stopped.")