"""Quick: check current state, patch for SHARED mode, restart."""
import paramiko, time, json, textwrap

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=60):
    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()

# Check current cmdline
print("=== Current process ===")
print(sh('ps aux | grep main.py | grep -v grep'))

# Check if patch exists
with sftp.open('/home/fabian/ComfyUI/comfy/model_management.py', 'r') as f:
    code = f.read().decode()
print(f"\n=== Patch status: {'APPLIED' if 'COMFYUI_SHARED_MEMORY' in code else 'NOT applied'} ===")

# Check current vram state in log
try:
    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 'vram state' in s or 'VAE load' in s or 'Device:' in s:
            print(f"  {s}")
except: pass

# =============================
# APPLY PATCH if not done
# =============================
if 'COMFYUI_SHARED_MEMORY' not in code:
    print("\nApplying SHARED memory patch...")
    # Backup
    sh('cp /home/fabian/ComfyUI/comfy/model_management.py /home/fabian/ComfyUI/comfy/model_management.py.bak')
    
    old = 'if cpu_state == CPUState.MPS:\n    vram_state = VRAMState.SHARED'
    new = '''if cpu_state == CPUState.MPS:
    vram_state = VRAMState.SHARED

# Shared memory APU: CPU+GPU share physical RAM (e.g. AMD BC-250)
import os as _os
if _os.environ.get("COMFYUI_SHARED_MEMORY") == "1" and cpu_state == CPUState.GPU:
    vram_state = VRAMState.SHARED
    logging.info("SHARED vram: APU shared memory mode enabled")'''
    
    if old in code:
        code = code.replace(old, new)
        with sftp.open('/home/fabian/ComfyUI/comfy/model_management.py', 'w') as f:
            f.write(code)
        print("  Patch applied!")
    else:
        print("  ERROR: Could not find patch target. Dumping area:")
        for i, line in enumerate(code.split('\n')):
            if 'MPS' in line and 'SHARED' in code.split('\n')[i+1] if i+1 < len(code.split('\n')) else '':
                for j in range(max(0,i-3), min(i+5, len(code.split('\n')))):
                    print(f"  {j+1}: {code.split(chr(10))[j]}")
else:
    print("  Patch already applied, good.")

# =============================
# KILL + RESTART with SHARED
# =============================
print("\nKilling ComfyUI...")
sh('pkill -9 -f "python3.*main.py" 2>/dev/null; sleep 2')

launcher = textwrap.dedent("""\
    #!/bin/bash
    export HSA_OVERRIDE_GFX_VERSION=10.1.0
    export HIP_VISIBLE_DEVICES=0
    export HSA_ENABLE_SDMA=0
    export HSA_TOOLS_LIB=""
    export HSA_TOOLS_REPORT_LOAD_FAILURE=0
    export PYTORCH_HIP_ALLOC_CONF=expandable_segments:False
    export OMP_NUM_THREADS=12
    export MKL_NUM_THREADS=12
    export OPENBLAS_NUM_THREADS=12
    export MIOPEN_FIND_MODE=3
    # APU shared memory mode: everything on GPU
    export COMFYUI_SHARED_MEMORY=1
    
    cd ~/ComfyUI
    source ~/comfyui-env/bin/activate
    
    # SHARED mode: models load on GPU directly (shared memory = zero copy)
    # --force-fp16: half precision for models
    # --fp16-vae: VAE in fp16 (160MB, fast on GPU)
    exec python3 main.py \\
        --listen 0.0.0.0 --port 8188 \\
        --force-fp16 \\
        --fp16-vae
""")
with sftp.open('/tmp/run_comfyui.sh', 'w') as f:
    f.write(launcher)
sh('chmod +x /tmp/run_comfyui.sh')

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"Started 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)

# Verify SHARED mode
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(k in s for k in ['vram state', 'SHARED', 'Device:', 'Total VRAM']):
        print(f"  {s}")

# Submit test
print("\nSubmitting test 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": 12345, "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[:120]}")

# Monitor
print("\nMonitoring (SHARED = VAE on GPU, everything on GPU)...")
t0 = time.time()
for i in range(200):
    el = int(time.time() - t0)
    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 '?'
    
    try:
        with sftp.open('/tmp/comfyui.log', 'r') as f:
            log = f.read().decode(errors='replace')
    except: log = ''
    
    samp = ''
    last = ''
    for line in log.split('\n'):
        s = line.strip()
        if '/8' in s and ('it/s' in s or 's/it' in s): samp = s
        if s and 'FETCH' not in s and 'startup' not in s and 'DEPRECATION' not in s: last = s
    
    print(f"  [{el:>4}s] {tc}C | {(samp or last)[-90:]}")
    
    imgs = sh('ls ~/ComfyUI/output/ZImageTurbo_GPU*.png 2>/dev/null', timeout=5)
    if imgs:
        et = ''
        for line in log.split('\n'):
            if 'Prompt executed' in line: et = line.strip()
        print(f"\n  *** DONE! ***  {imgs}")
        print(f"  {et}")
        print(f"  Wall: {el}s")
        for line in log.split('\n'):
            s = line.strip()
            if any(k in s for k in ['loaded', '/8', 'Prompt executed', 'VAE load', 'Requested']):
                if 'FETCH' not in s: print(f"  {s}")
        break
    
    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') and el > 30:
            time.sleep(2)
            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:")
                for line in log.split('\n')[-20:]:
                    s = line.strip()
                    if s and 'FETCH' not in s: print(f"  {s}")
            break
    except: pass
    
    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')[-30:]:
            if line.strip(): print(f"  {line.strip()}")
        break
    
    time.sleep(10)

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