"""Monitor ComfyUI - just watch log + GPU until image appears."""
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 gpu_info():
    chan = c.get_transport().open_session()
    chan.settimeout(10)
    chan.exec_command('/bin/bash -c "cat /sys/class/drm/card0/device/hwmon/hwmon*/temp1_input 2>/dev/null; echo SEP; rocm-smi -P 2>&1 | grep Graphics; echo SEP; pgrep -f python3.*main.py >/dev/null && echo ALIVE || echo DEAD"')
    out = b""
    while True:
        try:
            chunk = chan.recv(65536)
            if not chunk: break
            out += chunk
        except: break
    chan.close()
    parts = out.decode(errors='replace').split('SEP')
    temp = parts[0].strip() if len(parts) > 0 else '?'
    temp_c = int(temp) // 1000 if temp.isdigit() else '?'
    power = parts[1].strip().split(':')[-1].strip() if len(parts) > 1 else '?'
    alive = 'ALIVE' in (parts[2] if len(parts) > 2 else '')
    return temp_c, power, alive

t0 = time.time()
print("Monitoring... GPU should be at high power (>100W) during sampling")

for i in range(120):
    elapsed = int(time.time() - t0)
    temp_c, power, alive = gpu_info()
    
    # Read log via SFTP
    try:
        with sftp.open('/tmp/comfyui.log', 'r') as f:
            log = f.read().decode(errors='replace')
    except:
        log = ''
    
    # Find last meaningful line
    last = ''
    for line in reversed(log.split('\n')):
        s = line.strip()
        if s and 'FETCH' not in s and 'startup tasks' not in s and 'DEPRECATION' not in s:
            last = s
            break
    
    # Check for sampling progress in log
    sampling = ''
    for line in log.split('\n'):
        if '/8' in line and ('it/s' in line or 's/it' in line):
            sampling = line.strip()
    
    # Check output
    chan = c.get_transport().open_session()
    chan.settimeout(5)
    chan.exec_command('/bin/bash -c "ls ~/ComfyUI/output/*.png 2>/dev/null"')
    imgs = b""
    while True:
        try:
            chunk = chan.recv(65536)
            if not chunk: break
            imgs += chunk
        except: break
    chan.close()
    imgs = imgs.decode().strip()
    
    line_out = f"[{elapsed:>4}s] {temp_c}C {power} | "
    if sampling:
        line_out += sampling[-80:]
    else:
        line_out += last[-80:]
    print(line_out)
    
    if imgs:
        print(f"\n*** IMAGE GENERATED! ***")
        print(f"File: {imgs}")
        print(f"Time: {elapsed}s")
        # Print last 15 log lines
        for line in log.split('\n')[-15:]:
            if line.strip():
                print(f"  {line.strip()}")
        break
    
    if not alive:
        print("\n*** PROCESS DEAD ***")
        for line in log.split('\n')[-20:]:
            if line.strip():
                print(f"  {line.strip()}")
        break
    
    # Check queue empty
    chan = c.get_transport().open_session()
    chan.settimeout(5)
    chan.exec_command('/bin/bash -c "curl -s http://127.0.0.1:8188/queue 2>/dev/null"')
    qraw = b""
    while True:
        try:
            chunk = chan.recv(65536)
            if not chunk: break
            qraw += chunk
        except: break
    chan.close()
    try:
        qd = json.loads(qraw.decode())
        if not qd.get('queue_running') and not qd.get('queue_pending') and elapsed > 30:
            time.sleep(2)
            # Final image check
            chan = c.get_transport().open_session()
            chan.settimeout(5)
            chan.exec_command('/bin/bash -c "ls ~/ComfyUI/output/*.png 2>/dev/null"')
            imgs2 = b""
            while True:
                try:
                    chunk = chan.recv(65536)
                    if not chunk: break
                    imgs2 += chunk
                except: break
            chan.close()
            if imgs2.decode().strip():
                print(f"\n*** IMAGE GENERATED! ***")
                print(f"File: {imgs2.decode().strip()}")
                print(f"Time: {elapsed}s")
            else:
                print(f"\nQueue empty, no image. Last log lines:")
                for line in log.split('\n')[-20:]:
                    if line.strip():
                        print(f"  {line.strip()}")
            break
    except:
        pass
    
    time.sleep(15)

sftp.close()
c.close()
print("\nDone.")