#!/usr/bin/env python3
"""Wait for ComfyUI generation to complete on BC-250."""
import paramiko
import 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, 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():
        lines = out.strip().split('\n')
        if len(lines) > 30:
            print(f"  ... ({len(lines)} lines, showing last 30)")
            print('\n'.join(lines[-30:]))
        else:
            print(out.strip())
    if err.strip():
        print(f"STDERR: {err.strip()[-500:]}")
    print(f"  Exit: {rc}")
    return rc, out, err

# Monitor every 30s for up to 20 minutes
for i in range(40):
    time.sleep(30)
    rc, out, _ = run(f"bash -c 'wc -l /home/fabian/comfyui.log; echo \"---\"; tail -5 /home/fabian/comfyui.log'",
        desc=f"Check {i+1} ({(i+1)*30}s)")
    
    if 'Prompt executed in' in out:
        print("\n  GENERATION COMPLETE!")
        run("bash -c 'tail -30 /home/fabian/comfyui.log'", desc="Final log")
        break
    if 'Traceback' in out or 'Exception' in out:
        print("\n  ERROR!")
        run("bash -c 'tail -50 /home/fabian/comfyui.log'", desc="Error log")
        break
    
    # Also check process CPU
    _, pout, _ = run("bash -c 'ps -p 484588 -o %cpu,%mem,vsz,rss --no-headers 2>/dev/null || echo DEAD'")
    if 'DEAD' in pout:
        print("\n  PROCESS DIED!")
        run("bash -c 'tail -50 /home/fabian/comfyui.log'", desc="Death log")
        break

# Final output check
run("bash -c 'ls -la ~/ComfyUI/output/ 2>/dev/null'", desc="Output directory")
run("bash -c 'tail -30 /home/fabian/comfyui.log 2>/dev/null'", desc="Final log state")

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