"""Fix: keep UNet+VAE both on GPU (shared memory). No offloading."""
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()

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

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

lines = code.split('\n')

# Show current offload functions to understand exact code
print("\n=== Finding offload functions ===")
for i, line in enumerate(lines):
    if 'def unet_offload_device' in line or 'def vae_offload_device' in line:
        print(f"\n--- {line.strip()} at line {i+1} ---")
        for j in range(i, min(i+10, len(lines))):
            print(f"  {j+1}: {lines[j]}")

# ============ PATCH unet_offload_device ============
# Current: returns CPU unless HIGH_VRAM
# Fix: also return GPU for SHARED (APU shared memory = no point offloading)
old_unet = None
new_unet = None

for i, line in enumerate(lines):
    if 'def unet_offload_device' in line:
        # Grab the function body (next ~6 lines)
        chunk = '\n'.join(lines[i:i+8])
        print(f"\n=== unet_offload_device chunk ===\n{chunk}")
        
        # The function checks HIGH_VRAM only. Add SHARED.
        if 'HIGH_VRAM' in chunk and 'SHARED' not in chunk:
            old_unet = chunk
            new_unet = chunk.replace(
                'vram_state == VRAMState.HIGH_VRAM',
                'vram_state == VRAMState.HIGH_VRAM or vram_state == VRAMState.SHARED'
            )
            print(f"\n  -> Will patch to include SHARED")
        elif 'SHARED' in chunk:
            print(f"\n  -> Already patched for SHARED")
        break

# ============ PATCH vae_offload_device ============
# Current: returns CPU unless --gpu-only
# Fix: also return GPU for SHARED
old_vae = None
new_vae = None

for i, line in enumerate(lines):
    if 'def vae_offload_device' in line:
        chunk = '\n'.join(lines[i:i+8])
        print(f"\n=== vae_offload_device chunk ===\n{chunk}")
        
        if 'args.gpu_only' in chunk and 'SHARED' not in chunk:
            old_vae = chunk
            new_vae = chunk.replace(
                'args.gpu_only',
                'args.gpu_only or vram_state == VRAMState.SHARED'
            )
            print(f"\n  -> Will patch to include SHARED")
        elif 'SHARED' in chunk:
            print(f"\n  -> Already patched for SHARED")
        break

# Also check text_encoder_offload_device
for i, line in enumerate(lines):
    if 'def text_encoder_offload_device' in line:
        chunk = '\n'.join(lines[i:i+8])
        print(f"\n=== text_encoder_offload_device chunk ===\n{chunk}")
        break

# Apply patches
patched = False
if old_unet and new_unet:
    code = code.replace(old_unet, new_unet)
    patched = True
    print("\n[OK] Patched unet_offload_device")

if old_vae and new_vae:
    code = code.replace(old_vae, new_vae)
    patched = True
    print("[OK] Patched vae_offload_device")

if patched:
    # Backup and write
    sh('cp /home/fabian/ComfyUI/comfy/model_management.py /home/fabian/ComfyUI/comfy/model_management.py.bak2')
    with sftp.open('/home/fabian/ComfyUI/comfy/model_management.py', 'w') as f:
        f.write(code)
    print("[OK] Written to disk")
else:
    print("[INFO] No patches needed (already applied or code changed)")

# Verify
with sftp.open('/home/fabian/ComfyUI/comfy/model_management.py', 'r') as f:
    verify = f.read().decode()
for i, line in enumerate(verify.split('\n')):
    if 'def unet_offload_device' in line or 'def vae_offload_device' in line:
        print(f"\n--- VERIFY {line.strip()} ---")
        for j in range(i, min(i+8, len(verify.split(chr(10))))):
            print(f"  {j+1}: {verify.split(chr(10))[j]}")

# ============ RESTART ============
print("\n=== RESTARTING ===")
sh('rm -f /tmp/comfyui.log; rm -f ~/ComfyUI/output/ZImageTurbo_GPU*.png')
sh('nohup /tmp/run_comfyui.sh > /tmp/comfyui.log 2>&1 &')
time.sleep(3)
pid = sh('pgrep -f "python3.*main.py"')
print(f"PID: {pid}")

# Wait for ready
print("Waiting for HTTP", end='', flush=True)
for i in range(90):
    r = sh('curl -s -o /dev/null -w "%{http_code}" http://127.0.0.1:8188/ 2>/dev/null', timeout=5)
    if '200' in r:
        print(f" OK ({i*2}s)")
        break
    print('.', end='', flush=True)
    time.sleep(2)
else:
    print(" TIMEOUT")

# Check SHARED mode active
with sftp.open('/tmp/comfyui.log', 'r') as f:
    log = f.read().decode(errors='replace')
for line in log.split('\n'):
    s = line.strip()
    if any(x in s.lower() for x in ['vram state', 'shared', 'device:', 'total vram']):
        print(f"  {s}")

# Submit workflow
print("\nSubmitting workflow...")
wf = {"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"}}
}}
with sftp.open('/tmp/wf.json', 'w') as f:
    f.write(json.dumps(wf))
resp = sh('curl -s -X POST http://127.0.0.1:8188/prompt -H "Content-Type: application/json" -d @/tmp/wf.json')
print(f"  {resp[:150]}")

# Monitor - watch for UNet+VAE both on GPU, fast VAE
print("\nMonitoring...")
t0 = time.time()
for i in range(120):
    el = int(time.time() - t0)
    
    try:
        with sftp.open('/tmp/comfyui.log', 'r') as f:
            log = f.read().decode(errors='replace')
    except: log = ''
    
    # GPU usage
    gpu = sh('cat /sys/class/drm/card0/device/gpu_busy_percent 2>/dev/null || echo ?', timeout=5)
    temp = sh('cat /sys/class/drm/card0/device/hwmon/hwmon*/temp1_input 2>/dev/null', timeout=5)
    tc = int(temp)//1000 if temp.isdigit() else '?'
    
    # Get latest progress line
    last_progress = ''
    last_line = ''
    for line in log.split('\n'):
        s = line.strip()
        if '/8' in s and ('it/s' in s or 's/it' in s): last_progress = s
        if 'loaded' in s.lower() or 'VAE' in s or 'Requested' in s or 'Prompt executed' in s:
            last_line = s
        if s and 'FETCH' not in s: last_line = s
    
    status = last_progress or last_line
    print(f"  [{el:>3}s] GPU:{gpu}% {tc}C | {status[-100:]}")
    
    # Check for output
    imgs = sh('ls ~/ComfyUI/output/ZImageTurbo_GPU*.png 2>/dev/null', timeout=5)
    if imgs:
        print(f"\n  *** IMAGE DONE! *** {imgs}")
        print(f"  Wall time: {el}s")
        # Print key log lines
        for line in log.split('\n'):
            s = line.strip()
            if any(x in s for x in ['loaded', 'load device', 'offload device', 'Prompt executed', '/8', 'Requested', 'VAE']):
                if 'FETCH' not in s:
                    print(f"    {s}")
        break
    
    # Check queue empty
    if el > 30:
        q = sh('curl -s http://127.0.0.1:8188/queue 2>/dev/null', timeout=5)
        try:
            qd = json.loads(q)
            if not qd.get('queue_running') and not qd.get('queue_pending'):
                time.sleep(3)
                imgs = sh('ls ~/ComfyUI/output/ZImageTurbo_GPU*.png 2>/dev/null', timeout=5)
                if imgs:
                    print(f"\n  *** DONE: {imgs} ***")
                else:
                    print(f"\n  Queue empty, no image. Error?")
                    for line in log.split('\n')[-20:]:
                        if line.strip() and 'FETCH' not in line: print(f"    {line.strip()}")
                break
        except: pass
    
    # Check alive
    if sh('pgrep -f "python3.*main.py" >/dev/null && echo Y || echo N', timeout=5) == 'N':
        print("\n  CRASHED!")
        for line in log.split('\n')[-20:]:
            if line.strip(): print(f"    {line.strip()}")
        break
    
    time.sleep(5)

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