"""Post-reboot: start patched ComfyUI, submit, monitor. Single connection."""
import paramiko, time, json, sys

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)

def sh(cmd, t=30):
    stdin, stdout, stderr = c.exec_command(f"bash -lc '{cmd}'", timeout=t)
    return stdout.read().decode(errors='replace').strip()

# Verify patches survived reboot
p = sh("grep -c SHARED ~/ComfyUI/comfy/model_management.py")
print(f"SHARED refs in code: {p}")

# Start ComfyUI
sh("rm -f /tmp/comfyui.log")
sh("nohup /tmp/run_comfyui.sh > /tmp/comfyui.log 2>&1 &")
time.sleep(4)
pid = sh("pgrep -f main.py")
print(f"PID: {pid}")

# Wait for HTTP ready
print("Waiting for HTTP...", end='', flush=True)
for i in range(90):
    code = sh('curl -s -o /dev/null -w "%{http_code}" http://127.0.0.1:8188/ 2>/dev/null')
    if '200' in code:
        print(f" ready ({i*2}s)")
        break
    print('.', end='', flush=True)
    time.sleep(2)
else:
    print(" TIMEOUT")
    print(sh("tail -30 /tmp/comfyui.log"))
    sys.exit(1)

# Show startup state
log = sh("cat /tmp/comfyui.log")
for l in log.split('\n'):
    s = l.strip()
    if any(x in s for x in ['vram state', 'SHARED', 'Device:', 'Total VRAM']):
        print(f"  {s}")

# Submit workflow
wf = json.dumps({"prompt": {
    "1": {"class_type": "UnetLoaderGGUF", "inputs": {"unet_name": "z_image_turbo-Q5_K_S.gguf"}},
    "2": {"class_type": "CLIPLoaderGGUF", "inputs": {"clip_name": "Qwen3-4B.i1-Q5_K_S.gguf", "type": "qwen_image"}},
    "3": {"class_type": "VAELoader", "inputs": {"vae_name": "ae.safetensors"}},
    "4": {"class_type": "CLIPTextEncode", "inputs": {"text": "A red fox in a snowy forest, photorealistic", "clip": ["2", 0]}},
    "5": {"class_type": "EmptyLatentImage", "inputs": {"width": 512, "height": 512, "batch_size": 1}},
    "6": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "positive": ["4", 0], "negative": ["4", 0],
        "latent_image": ["5", 0], "seed": 99999, "steps": 8, "cfg": 1.0,
        "sampler_name": "euler", "scheduler": "simple", "denoise": 1.0}},
    "7": {"class_type": "VAEDecode", "inputs": {"samples": ["6", 0], "vae": ["3", 0]}},
    "8": {"class_type": "SaveImage", "inputs": {"images": ["7", 0], "filename_prefix": "ZImageTurbo_GPU"}}
}})
# Write workflow and submit
sh(f"echo '{wf}' > /tmp/wf.json")
resp = sh('curl -s -X POST http://127.0.0.1:8188/prompt -H "Content-Type: application/json" -d @/tmp/wf.json')
print(f"Submitted: {resp[:100]}")

# Monitor
t0 = time.time()
shown = set()
for _ in range(200):
    el = int(time.time() - t0)
    log = sh("cat /tmp/comfyui.log 2>/dev/null")
    
    for l in log.split('\n'):
        s = l.strip()
        if s not in shown and any(x in s for x in ['/8', 'loaded completely', 'load device', 'Requested to load', 'Prompt executed', 'Error', 'OOM']):
            if 'FETCH' not in s and 'audio' not in s:
                gpu = sh("cat /sys/class/drm/card0/device/gpu_busy_percent 2>/dev/null")
                print(f"  [{el:>3}s] GPU:{gpu}% | {s[-110:]}")
                shown.add(s)
    
    if 'Prompt executed' in log:
        print(f"\n*** DONE in {el}s! ***")
        imgs = sh("ls ~/ComfyUI/output/ZImageTurbo_GPU*.png 2>/dev/null")
        print(f"  Images: {imgs}")
        break
    
    alive = sh("pgrep -c -f main.py 2>/dev/null")
    if alive == '0':
        print(f"\nCRASHED at {el}s!")
        for l in log.split('\n')[-15:]:
            if l.strip(): print(f"  {l.strip()}")
        break
    
    time.sleep(5)

c.close()