"""Print training status as KEY=VALUE lines. Runs on the GCP VM."""
import json
import glob
import os
import re
import subprocess

# Process detection - check for both AIDE and direct training
aide_procs = int(os.popen('ps aux | grep -E "run_aide" | grep -v grep | wc -l').read().strip())
train_procs = int(os.popen('ps aux | grep -E "train_100ep|train_optimal|train_fulldata" | grep -v grep | wc -l').read().strip())
print(f"AIDE_RUNNING={aide_procs + train_procs}")

# Determine mode: direct training or AIDE
mode = "unknown"
log_file = None

# Priority 1: Full-data training (train_full.log)
if os.path.exists(os.path.expanduser("~/train_full.log")) and os.path.getsize(os.path.expanduser("~/train_full.log")) > 0:
    log_file = os.path.expanduser("~/train_full.log")
    mode = "direct"
# Priority 2: Direct training (train_100ep.log)
elif os.path.exists(os.path.expanduser("~/train_100ep.log")) and os.path.getsize(os.path.expanduser("~/train_100ep.log")) > 0:
    log_file = os.path.expanduser("~/train_100ep.log")
    mode = "direct"
# Priority 2: Optimal training
elif os.path.exists(os.path.expanduser("~/optimal_train.log")) and os.path.getsize(os.path.expanduser("~/optimal_train.log")) > 0:
    log_file = os.path.expanduser("~/optimal_train.log")
    mode = "direct"
elif os.path.exists(os.path.expanduser("~/optimal_train2.log")) and os.path.getsize(os.path.expanduser("~/optimal_train2.log")) > 0:
    log_file = os.path.expanduser("~/optimal_train2.log")
    mode = "direct"
# Priority 3: AIDE
elif os.path.exists(os.path.expanduser("~/aide-run.log")) and os.path.getsize(os.path.expanduser("~/aide-run.log")) > 0:
    log_file = os.path.expanduser("~/aide-run.log")
    mode = "aide"

print(f"MODE={mode}")

if log_file:
    content = open(log_file).read()
    lines = content.strip().split("\n")
    print(f"LOG_LINES={len(lines)}")
    print(f"LOG_TAIL={lines[-1].strip()[:80] if lines else 'empty'}")

    if mode == "direct":
        # Parse epoch progress from ultralytics output: "31/100  8.1G ..."
        # Match lines like "     31/100      8.1G" - epoch/total at start of training line
        epoch_pattern = re.findall(r'^\s*(\d+)/(\d+)\s+[\d.]+G', content, re.MULTILINE)

        if epoch_pattern:
            cur, total = int(epoch_pattern[-1][0]), int(epoch_pattern[-1][1])
            print(f"EPOCH={cur}/{total}")
            print(f"STEP={cur}")
            print(f"BEST_SCORE={cur/total:.4f}")  # Progress as proxy
            print(f"BUGGY=0")

            # Check if training completed
            if "Done!" in content or cur >= total:
                print(f"TRAINING_DONE=1")
            else:
                print(f"TRAINING_DONE=0")
        else:
            print(f"EPOCH=0/100")
            print(f"STEP=0")
            print(f"BEST_SCORE=0.0000")
            print(f"BUGGY=0")
            print(f"TRAINING_DONE=0")

        # Show last few log lines
        for line in lines[-3:]:
            print(f"LOG_LINE={line.strip()[:90]}")

        # Find trained model weights
        for wdir in glob.glob(os.path.expanduser("~/aide/train_*/run/weights/best.pt")):
            print(f"TRAINED_MODEL={wdir}")
        for wdir in glob.glob(os.path.expanduser("~/aide/working/*/weights/best.pt")):
            print(f"TRAINED_MODEL={wdir}")

    elif mode == "aide":
        # AIDE mode - check journal
        files = glob.glob(os.path.expanduser("~/aide/logs/*/journal.json"))
        if files:
            j = json.load(open(files[0]))
            nodes = j.get("nodes", [])
            step = len(nodes)
            best = 0
            buggy = sum(1 for n in nodes if n.get("is_buggy", False))
            for n in nodes:
                m = n.get("metric", {})
                v = m.get("value") if isinstance(m, dict) else m
                if v is not None and isinstance(v, (int, float)) and v > best:
                    best = v
            print(f"STEP={step}")
            print(f"BEST_SCORE={best:.4f}")
            print(f"BUGGY={buggy}")
        else:
            print("STEP=0")
            print("BEST_SCORE=0.0000")
            print("BUGGY=0")
else:
    print("LOG_LINES=0")
    print("LOG_TAIL=empty")
    print("STEP=0")
    print("BEST_SCORE=0.0000")
    print("BUGGY=0")

# GPU
try:
    gpu_out = subprocess.check_output(
        ["nvidia-smi", "--query-gpu=utilization.gpu,memory.used,memory.total",
         "--format=csv,noheader,nounits"], text=True
    ).strip().split(", ")
    print(f"GPU_UTIL={gpu_out[0]}")
    print(f"GPU_MEM_USED={gpu_out[1]}")
    print(f"GPU_MEM_TOTAL={gpu_out[2]}")
except Exception:
    pass