"""Diagnose and fix VAE hang. Check log, fix threading, restart."""
import paramiko, time, json

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()

# 1. What's running?
print("=== CURRENT STATE ===")
ps = sh('ps aux | grep main.py | grep -v grep')
print(f"Process: {ps or 'NONE'}")

# Check env of running process
env = sh(r'cat /proc/$(pgrep -f "python3.*main.py" | head -1)/environ 2>/dev/null | tr "\0" "\n" | grep -E "OMP|MKL|COMFYUI|THREAD|OPENBLAS"')
print(f"Env:\n{env}")

# 2. Log tail
print("\n=== LOG TAIL ===")
try:
    with sftp.open('/tmp/comfyui.log', 'r') as f:
        log = f.read().decode(errors='replace')
    for line in log.split('\n')[-40:]:
        s = line.strip()
        if s: print(f"  {s}")
except Exception as e:
    log = ''
    print(f"  No log: {e}")

# 3. Check launcher
print("\n=== LAUNCHER ===")
try:
    with sftp.open('/tmp/run_comfyui.sh', 'r') as f:
        print(f.read().decode())
except: print("  No launcher")

# 4. model_management.py patch?
with sftp.open('/home/fabian/ComfyUI/comfy/model_management.py', 'r') as f:
    mm = f.read().decode()
print(f"SHARED patch: {'YES' if 'COMFYUI_SHARED_MEMORY' in mm else 'NO'}")

# 5. Check torch threads in same env
print("\n=== TORCH THREADS ===")
tcheck = sh('''source ~/comfyui-env/bin/activate
export HSA_OVERRIDE_GFX_VERSION=10.1.0
python3 -c "
import torch, os
print(f'torch.get_num_threads() = {torch.get_num_threads()}')
print(f'OMP_NUM_THREADS = {os.environ.get(chr(34)+'OMP_NUM_THREADS'+chr(34), chr(34)+'NOT SET'+chr(34))}')
"''', timeout=30)
print(tcheck)

# 6. Check sitecustomize
print("\n=== SITECUSTOMIZE ===")
sc = sh('cat ~/comfyui-env/lib/python*/site-packages/sitecustomize.py 2>/dev/null || echo MISSING')
print(sc[:500])

# 7. Key log lines
print("\n=== KEY LOG ENTRIES ===")
for line in log.split('\n'):
    s = line.strip()
    if any(x in s.lower() for x in ['vram state', 'shared', 'loaded', 'offloaded', 'device:', 'total vram', 'vae', 'thread']):
        print(f"  {s}")

sftp.close()
c.close()
print("\nDiag done.")