"""Check if VAE is hanging or running, then fix and restart."""
import paramiko, time, json, textwrap

k = paramiko.Ed25519Key.from_private_key_file(r'C:\Users\fabia\.ssh\id_ed25519')
c = paramiko.SSHClient()
c.set_missing_host_key_policy(paramiko.AutoAddPolicy())
c.connect('192.168.178.150', username='fabian', pkey=k, timeout=10)
sftp = c.open_sftp()

def sh(cmd, timeout=30):
    chan = c.get_transport().open_session()
    chan.settimeout(timeout)
    chan.exec_command('/bin/bash -l -c ' + "'" + cmd.replace("'", "'\\''") + "'")
    out = b""
    while True:
        try:
            chunk = chan.recv(65536)
            if not chunk: break
            out += chunk
        except: break
    chan.close()
    return out.decode(errors='replace').strip()

# Check if process is truly hung or still computing
print("=== GPU MEMORY / ACTIVITY ===")
print(sh('rocm-smi --showmemuse --showuse 2>/dev/null | head -20'))
print()
print(sh('rocm-smi --showmeminfo vram 2>/dev/null'))
print()

# Check CPU usage of the process  
pid = sh('pgrep -f "python3.*main.py" | head -1')
print(f"PID: {pid}")
if pid:
    print(f"CPU%: {sh('ps -p ' + pid + ' -o %cpu,%mem,rss,vsz --no-headers')}")
    # Check /proc/pid/status for threads
    print(f"Threads: {sh('grep Threads /proc/' + pid + '/status 2>/dev/null')}")
    # strace peek - what syscall is it stuck on?
    print(f"\nStack peek (1s):")
    print(sh('timeout 2 strace -p ' + pid + ' -c 2>&1 | head -20', timeout=10))

# Check last log line timestamp
print("\n=== LOG LAST LINES ===")
with sftp.open('/tmp/comfyui.log', 'r') as f:
    log = f.read().decode(errors='replace')
lines = [l for l in log.split('\n') if l.strip()]
for l in lines[-10:]:
    print(f"  {l.strip()}")

# Check if output image exists
imgs = sh('ls -la ~/ComfyUI/output/ZImageTurbo_GPU*.png 2>/dev/null')
print(f"\nOutput images: {imgs or 'NONE'}")

# ============================================================
# THE FIX: VAE on GPU hangs on this APU. Force CPU VAE but 
# ensure ALL 12 threads are used via explicit torch patch.
# We'll also patch ComfyUI to call torch.set_num_threads(12)
# RIGHT BEFORE vae decode, in case something resets it.
# ============================================================

print("\n" + "="*60)
print("FIXING: Kill, patch VAE threading, restart with --cpu-vae")
print("="*60)

# Kill
sh('pkill -9 -f "python3.*main.py" 2>/dev/null; sleep 2')
print("Killed ComfyUI")

# Read current model_management.py
with sftp.open('/home/fabian/ComfyUI/comfy/model_management.py', 'r') as f:
    mm = f.read().decode()

# Find and patch vae_device to force CPU + set threads
# Current: def vae_device():  ... return vae_dev
# We need to find the vae decode path and add threading there
# But actually the simplest: patch the VAEDecode node itself

# Check nodes_latent.py for VAEDecode
vae_decode_path = sh('grep -rn "class VAEDecode" ~/ComfyUI/comfy_extras/ ~/ComfyUI/nodes.py 2>/dev/null')
print(f"\nVAEDecode location: {vae_decode_path}")

# Read the relevant file
vae_file = ''
vae_line = ''
for line in vae_decode_path.split('\n'):
    if 'class VAEDecode' in line and 'Tiled' not in line:
        parts = line.split(':')
        vae_file = parts[0]
        vae_line = parts[1]
        break
print(f"VAE file: {vae_file}, line: {vae_line}")

if vae_file:
    with sftp.open(vae_file, 'r') as f:
        vae_code = f.read().decode()
    # Show VAEDecode class
    vae_lines = vae_code.split('\n')
    start = int(vae_line) - 1
    print(f"\nVAEDecode class (from line {vae_line}):")
    for i in range(start, min(start+30, len(vae_lines))):
        print(f"  {i+1}: {vae_lines[i]}")

# Also check where vae.decode is called in the VAE wrapper
print("\n=== VAE decode method location ===")
vae_impl = sh('grep -rn "def decode" ~/ComfyUI/comfy/sd.py 2>/dev/null | head -5')
print(vae_impl)

# Read sd.py decode method
for line in vae_impl.split('\n'):
    if 'def decode' in line:
        parts = line.split(':')
        sd_file = parts[0]
        sd_line = int(parts[1])
        with sftp.open(sd_file, 'r') as f:
            sd_code = f.read().decode()
        sd_lines = sd_code.split('\n')
        print(f"\n{sd_file} decode method:")
        for i in range(sd_line-2, min(sd_line+40, len(sd_lines))):
            print(f"  {i+1}: {sd_lines[i]}")
        break

sftp.close()
c.close()
print("\nDiag complete. Next: apply fix.")