#!/usr/bin/env python3
"""Monitor ComfyUI generation progress - poll every 20s."""
import paramiko, json, time

ssh = paramiko.SSHClient()
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
ssh.connect('192.168.178.150', username='fabian', key_filename=r'C:\Users\fabia\.ssh\id_ed25519')

def run(cmd, timeout=30):
    _, stdout, stderr = ssh.exec_command(cmd, timeout=timeout)
    return stdout.read().decode()

last_log_hash = ""
for i in range(90):  # up to 30 minutes
    time.sleep(20)
    
    # Process stats
    stats = run("bash -c 'PID=$(pgrep -f \"python3 main.py\" | head -1); "
        "if [ -n \"$PID\" ]; then "
        "  CPU=$(ps -p $PID -o %cpu --no-headers); "
        "  MEM=$(ps -p $PID -o rss --no-headers); "
        "  echo \"CPU:${CPU}% RSS:$((MEM/1024))MB LOAD:$(cat /proc/loadavg | cut -d\" \" -f1-3)\"; "
        "else echo DEAD; fi'").strip()
    
    # GPU
    gpu = run("bash -c 'rocm-smi --showuse --showmemuse 2>/dev/null | grep -E \"GPU|%\" | head -5 || echo no-gpu'").strip()
    
    # Log tail
    log = run("tail -10 /home/fabian/comfyui.log 2>/dev/null").strip()
    log_hash = hash(log)
    
    elapsed = (i+1) * 20
    mins = elapsed // 60
    secs = elapsed % 60
    print(f"[{mins}m{secs:02d}s] {stats}")
    
    # Show GPU line
    for line in gpu.split('\n'):
        if '%' in line or 'GPU' in line:
            print(f"  GPU: {line.strip()}")
            break
    
    # Show last meaningful log line
    if log_hash != last_log_hash:
        for line in reversed(log.split('\n')):
            l = line.strip()
            if l and not l.startswith('FETCH'):
                print(f"  LOG: {l}")
                break
        last_log_hash = log_hash
    
    if 'DEAD' in stats:
        print("\nPROCESS DIED!")
        print(run("tail -60 /home/fabian/comfyui.log 2>/dev/null"))
        break
    
    if 'Prompt executed in' in log:
        print(f"\nSUCCESS! Image generated!")
        print(run("tail -30 /home/fabian/comfyui.log 2>/dev/null"))
        print("\n=== OUTPUT FILES ===")
        print(run("ls -lah /home/fabian/ComfyUI/output/ 2>/dev/null"))
        break
    
    if 'Traceback' in log or 'CUDA out of memory' in log or 'RuntimeError' in log:
        print(f"\nERROR!")
        print(run("tail -60 /home/fabian/comfyui.log 2>/dev/null"))
        break

else:
    print("\nTimed out after 30 minutes")
    print(run("tail -40 /home/fabian/comfyui.log 2>/dev/null"))

ssh.close()