#!/usr/bin/env python3
"""Diagnose ComfyUI state on BC-250 — check if stuck or OOM."""
import paramiko

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, desc=""):
    if desc:
        print(f"\n{'='*60}")
        print(f"  {desc}")
        print(f"{'='*60}")
    _, stdout, stderr = ssh.exec_command(cmd, timeout=timeout)
    out = stdout.read().decode()
    err = stderr.read().decode()
    rc = stdout.channel.recv_exit_status()
    if out.strip():
        print(out.strip())
    if err.strip():
        lines = err.strip().split('\n')[-10:]
        print(f"STDERR: {chr(10).join(lines)}")
    print(f"  Exit: {rc}")
    return rc, out, err

# Check if ComfyUI process is alive
run("bash -c 'ps aux | grep \"python main.py\" | grep -v grep'",
    desc="ComfyUI process status")

# Memory state
run("bash -c 'free -h'", desc="RAM/Swap usage")

# GPU VRAM
run("bash -c 'cat /sys/class/drm/card1/device/mem_info_vram_used 2>/dev/null; "
    "echo \"---\"; cat /sys/class/drm/card1/device/mem_info_vram_total 2>/dev/null'",
    desc="GPU VRAM usage")

# dmesg for OOM
run("bash -c 'dmesg | tail -20'", desc="Recent kernel messages")

# Last 80 lines of comfyui log
run("bash -c 'tail -80 /home/fabian/comfyui.log 2>/dev/null'", desc="ComfyUI log (last 80)")

# Check if port still listening
run("bash -c 'ss -tlnp | grep 8188 || echo PORT_GONE'", desc="Port 8188")

ssh.close()
print("\nDone.")