"""Fix: Back to --novram (proven working for GPU sampling) + --cpu-vae."""
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=15)

def sh(cmd, timeout=60):
    chan = c.get_transport().open_session()
    chan.settimeout(timeout)
    # Use bash array to avoid quoting issues
    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()

def sftp_write(path, content):
    sftp = c.open_sftp()
    with sftp.open(path, 'w') as f:
        f.write(content)
    sftp.close()

def sftp_read(path):
    sftp = c.open_sftp()
    with sftp.open(path, 'r') as f:
        data = f.read().decode(errors='replace')
    sftp.close()
    return data

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

# ---- STEP 2: Write launcher ----
print("\nSTEP 2: Write launcher with --novram (PROVEN to work on this APU)")
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=1
    
    cd ~/ComfyUI
    source ~/comfyui-env/bin/activate
    
    # --novram = model weights on CPU, GPU only for compute (correct for shared-memory APU)
    # --cpu-vae = VAE decode on CPU (fixes known hang on this GPU)
    # --force-fp16 = half precision to save memory
    exec python3 main.py \\
        --listen 0.0.0.0 --port 8188 \\
        --novram \\
        --force-fp16 \\
        --cpu-vae \\
        --disable-smart-memory
""")
sftp_write('/tmp/run_comfyui.sh', launcher)
sh('chmod +x /tmp/run_comfyui.sh')
print("  Written: --novram --force-fp16 --cpu-vae --disable-smart-memory")

# ---- STEP 3: Start ----
print("\nSTEP 3: Start ComfyUI")
sh('rm -f /tmp/comfyui.log; touch /tmp/comfyui.log')
sh('nohup /tmp/run_comfyui.sh > /tmp/comfyui.log 2>&1 &')
time.sleep(3)
pid = sh('pgrep -f "python3.*main.py"')
if not pid:
    print("  FAILED!")
    print(sftp_read('/tmp/comfyui.log'))
    c.close()
    exit(1)
print(f"  PID: {pid}")

# ---- STEP 4: Wait for HTTP 200 ----
print("\nSTEP 4: Wait for server ready", end='', flush=True)
for i in range(120):
    code = sh('curl -s -o /dev/null -w "%{http_code}" http://127.0.0.1:8188/ 2>/dev/null', timeout=5)
    if '200' in code:
        print(f" READY ({i*2}s)")
        break
    # Check if process died
    alive = sh('pgrep -f "python3.*main.py" >/dev/null && echo Y || echo N', timeout=5)
    if alive == 'N':
        print("\n  Process died!")
        print(sftp_read('/tmp/comfyui.log'))
        c.close()
        exit(1)
    if i % 10 == 0 and i > 0:
        log = sftp_read('/tmp/comfyui.log')
        lines = [l for l in log.split('\n') if l.strip()]
        print(f"\n  [{i*2}s] {lines[-1][:80] if lines else '...'}", end='', flush=True)
    else:
        print('.', end='', flush=True)
    time.sleep(2)
else:
    print("\n  TIMEOUT!")
    print(sftp_read('/tmp/comfyui.log')[-1000:])
    c.close()
    exit(1)

# Verify startup flags
log = sftp_read('/tmp/comfyui.log')
if 'NO_VRAM' in log or 'NOVRAM' in log.upper():
    print("  Confirmed: NOVRAM mode (GPU compute only, model on CPU)")
for line in log.split('\n'):
    if 'vram state' in line.lower():
        print(f"  {line.strip()}")
    if 'Device:' in line:
        print(f"  {line.strip()}")

# ---- STEP 5: Submit workflow ----
print("\nSTEP 5: Submit workflow")
workflow = {
    "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": 42, "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"}}
    }
}
sftp_write('/tmp/wf.json', json.dumps(workflow))
resp = sh('curl -s -X POST http://127.0.0.1:8188/prompt -H "Content-Type: application/json" -d @/tmp/wf.json', timeout=10)
print(f"  Response: {resp[:200]}")
if 'prompt_id' not in resp:
    print("  FAILED!")
    c.close()
    exit(1)
prompt_id = json.loads(resp).get('prompt_id', '?')
print(f"  Prompt ID: {prompt_id}")

# ---- STEP 6: Monitor ----
print("\nSTEP 6: Monitor (expect GPU power >100W during sampling)")
t0 = time.time()
sampling_seen = False

for i in range(200):
    elapsed = int(time.time() - t0)
    
    gpu_temp = sh('cat /sys/class/drm/card0/device/hwmon/hwmon*/temp1_input 2>/dev/null', timeout=5)
    temp_c = int(gpu_temp) // 1000 if gpu_temp.isdigit() else '?'
    
    # Read log via SFTP to avoid shell issues
    try:
        log = sftp_read('/tmp/comfyui.log')
    except:
        log = ''
    
    lines = log.strip().split('\n')
    # Find last meaningful line (skip manager spam)
    last = ''
    for line in reversed(lines):
        if 'FETCH ComfyRegistry' not in line and 'All startup tasks' not in line and 'FETCH DATA' not in line and line.strip():
            last = line.strip()
            break
    
    # Detect sampling progress
    for line in lines:
        if '/8' in line and 'it/s' in line:
            sampling_seen = True
    
    print(f"  [{elapsed:>4}s] {temp_c}C | {last[-100:]}")
    
    # Check for output image
    imgs = sh('ls ~/ComfyUI/output/*.png 2>/dev/null', timeout=5)
    if imgs:
        print(f"\n  *** IMAGE GENERATED! ***")
        print(f"  File: {imgs}")
        print(f"  Total: {elapsed}s")
        # Show last 15 lines
        for line in lines[-15:]:
            if line.strip():
                print(f"  {line.strip()}")
        break
    
    # Check queue
    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 elapsed > 30:
            time.sleep(3)
            imgs = sh('ls ~/ComfyUI/output/*.png 2>/dev/null', timeout=5)
            if imgs:
                print(f"\n  *** IMAGE GENERATED! ***")
                print(f"  File: {imgs}")
            else:
                print(f"\n  Queue empty, no image. Error in log:")
                for line in lines[-20:]:
                    if line.strip():
                        print(f"  {line}")
            break
    except:
        pass
    
    # Check process alive
    alive = sh('pgrep -f "python3.*main.py" >/dev/null && echo Y || echo N', timeout=5)
    if alive == 'N':
        print(f"\n  *** CRASHED ***")
        for line in lines[-30:]:
            if line.strip():
                print(f"  {line}")
        break
    
    time.sleep(15)

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