#!/usr/bin/env python3
"""
Fix: Add MIOpen fast-find, pre-warm GPU, restart ComfyUI, generate.
All GPU ops confirmed working. The hang is likely cold MIOpen kernel cache.
Single SSH connection.
"""
import paramiko, json, time, sys

ssh = paramiko.SSHClient()
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
for attempt in range(5):
    try:
        ssh.connect('192.168.178.150', username='fabian',
                    key_filename=r'C:\Users\fabia\.ssh\id_ed25519', timeout=10)
        break
    except Exception as e:
        print(f"  SSH attempt {attempt+1}/5: {e}")
        time.sleep(10)
else:
    print("FATAL: Cannot connect"); sys.exit(1)

def run(cmd, timeout=300):
    _, stdout, stderr = ssh.exec_command(cmd, timeout=timeout)
    out = stdout.read().decode()
    err = stderr.read().decode()
    return out, err

try:
    # 1. Kill any lingering ComfyUI
    print("=== Kill any ComfyUI ===")
    run("pkill -9 -f 'python3 main.py' 2>/dev/null; sleep 1")
    print("  Done")

    # 2. Pre-warm MIOpen kernel cache with DiT-like operations
    print("\n=== Pre-warming MIOpen kernel cache ===")
    print("  This compiles HIP kernels that Z-Image-Turbo will need.")
    print("  First run after reboot is slow (kernel compilation)...")
    
    warmup_code = r"""
import torch, torch.nn as nn, time

# Simulate Z-Image-Turbo DiT operations
device = 'cuda'
dtype = torch.float16

print("Warming up HIP kernels for DiT inference...")
t0 = time.time()

# 1. Linear layers (DiT blocks)
print("  Linear layers...", end=" ", flush=True)
for size in [(1024,1024), (4096,1024), (1024,4096)]:
    l = nn.Linear(*size).to(device, dtype)
    x = torch.randn(1, 64, size[0], device=device, dtype=dtype)
    y = l(x)
    del l, x, y
torch.cuda.synchronize()
print(f"{time.time()-t0:.1f}s")

# 2. Attention (SDPA - the core of DiT)
print("  Scaled dot-product attention...", end=" ", flush=True)
t1 = time.time()
for heads in [8, 16, 24]:
    q = torch.randn(1, heads, 64, 64, device=device, dtype=dtype)
    k = torch.randn(1, heads, 64, 64, device=device, dtype=dtype)    
    v = torch.randn(1, heads, 64, 64, device=device, dtype=dtype)
    y = torch.nn.functional.scaled_dot_product_attention(q, k, v)
    del q, k, v, y
torch.cuda.synchronize()
print(f"{time.time()-t1:.1f}s")

# 3. LayerNorm / RMSNorm
print("  Normalization layers...", end=" ", flush=True)
t1 = time.time()
for dim in [1024, 2048, 4096]:
    ln = nn.LayerNorm(dim).to(device, dtype)
    x = torch.randn(1, 64, dim, device=device, dtype=dtype)
    y = ln(x)
    del ln, x, y
torch.cuda.synchronize()
print(f"{time.time()-t1:.1f}s")

# 4. Conv2d (VAE-like, but we'll run VAE on CPU)
print("  Conv2d layers...", end=" ", flush=True)
t1 = time.time()
for ch in [64, 128, 256]:
    c = nn.Conv2d(ch, ch, 3, padding=1).to(device, dtype)
    x = torch.randn(1, ch, 32, 32, device=device, dtype=dtype)
    y = c(x)
    del c, x, y
torch.cuda.synchronize()
print(f"{time.time()-t1:.1f}s")

# 5. Full mini-DiT forward pass simulation  
print("  Mini-DiT forward pass simulation...", end=" ", flush=True)
t1 = time.time()
hidden = 1024
seq_len = 256
heads = 16
head_dim = hidden // heads
# Simulate a DiT block
x = torch.randn(1, seq_len, hidden, device=device, dtype=dtype)
norm = nn.LayerNorm(hidden).to(device, dtype)
qkv = nn.Linear(hidden, hidden*3).to(device, dtype)
proj = nn.Linear(hidden, hidden).to(device, dtype)
ff1 = nn.Linear(hidden, hidden*4).to(device, dtype)
ff2 = nn.Linear(hidden*4, hidden).to(device, dtype)
for step in range(3):
    h = norm(x)
    q, k, v = qkv(h).chunk(3, dim=-1)
    q = q.view(1, seq_len, heads, head_dim).transpose(1,2)
    k = k.view(1, seq_len, heads, head_dim).transpose(1,2)
    v = v.view(1, seq_len, heads, head_dim).transpose(1,2)
    attn = torch.nn.functional.scaled_dot_product_attention(q, k, v)
    attn = attn.transpose(1,2).contiguous().view(1, seq_len, hidden)
    x = x + proj(attn)
    x = x + ff2(torch.nn.functional.gelu(ff1(norm(x))))
torch.cuda.synchronize()
print(f"{time.time()-t1:.1f}s")

total = time.time() - t0
print(f"\nGPU kernel warmup complete in {total:.1f}s")
print(f"VRAM used: {torch.cuda.memory_allocated()//1048576} MB")
torch.cuda.empty_cache()
print(f"VRAM after cleanup: {torch.cuda.memory_allocated()//1048576} MB")
print("WARMUP_DONE")
"""
    # Write warmup script
    run(f"cat > /tmp/gpu_warmup.py << 'PYEOF'\n{warmup_code}\nPYEOF")
    
    out, err = run("bash -c 'source ~/comfyui-env/bin/activate && "
                   "HSA_OVERRIDE_GFX_VERSION=10.1.0 HSA_ENABLE_SDMA=0 "
                   "MIOPEN_FIND_MODE=3 MIOPEN_FIND_ENFORCE=3 "
                   "python3 /tmp/gpu_warmup.py' 2>&1", timeout=300)
    print(out.strip())
    if 'WARMUP_DONE' not in out:
        print(f"  WARNING: Warmup may have failed")
        print(f"  STDERR: {err.strip()[:500]}")

    # 3. Update startup script with MIOpen settings
    print("\n=== Update startup script ===")
    script = """#!/bin/bash
# BC-250 ComfyUI Launcher — GPU (ROCm) + CPU VAE

# GPU identity (Cyan Skillfish gfx1013 -> gfx1010)
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

# MIOpen: fast kernel selection (avoid long auto-tune on first run)
export MIOPEN_FIND_MODE=3
export MIOPEN_FIND_ENFORCE=3

# Use all 12 CPU cores for CPU-side work (dequant, text encoding)
export OMP_NUM_THREADS=12
export MKL_NUM_THREADS=12
export OPENBLAS_NUM_THREADS=12

# Activate venv
source /home/fabian/comfyui-env/bin/activate
cd /home/fabian/ComfyUI

# --novram: offload models to RAM, send layers to GPU one at a time
# --force-fp16: fp16 diffusion to halve VRAM usage
# --cpu-vae: VAE decode on CPU (GPU hangs on full VAE forward pass)
exec python3 main.py \\
    --listen 0.0.0.0 --port 8188 \\
    --novram \\
    --force-fp16 \\
    --cpu-vae
"""
    run(f"cat > /home/fabian/start_comfyui.sh << 'HEREDOC_END'\n{script}HEREDOC_END\n"
        f"chmod +x /home/fabian/start_comfyui.sh")
    print("  Written with MIOpen fast-find + --novram --force-fp16 --cpu-vae")

    # 4. Launch ComfyUI
    print("\n=== Launch ComfyUI ===")
    run("nohup /home/fabian/start_comfyui.sh > /home/fabian/comfyui.log 2>&1 &")
    time.sleep(2)
    
    print("  Waiting for server...")
    for i in range(50):
        time.sleep(3)
        out, _ = run("curl -s -o /dev/null -w '%{http_code}' http://localhost:8188/ 2>/dev/null || echo 0")
        if out.strip() == '200':
            print(f"  Server ready ({(i+1)*3}s)")
            break
        if i % 5 == 4:
            log, _ = run("tail -2 /home/fabian/comfyui.log 2>/dev/null")
            last = [l.strip() for l in log.strip().split('\n') if l.strip()]
            print(f"  [{(i+1)*3}s] ... {last[-1][:80] if last else ''}")
    else:
        print("  Timeout!")
        out, _ = run("tail -40 /home/fabian/comfyui.log")
        print(out)
        sys.exit(1)

    # 5. Submit workflow
    print("\n=== Submit workflow (512x512, 8 steps) ===")
    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": "CLIPTextEncode",
                  "inputs": {"text": "", "clip": ["2", 0]}},
            "6": {"class_type": "EmptyLatentImage",
                  "inputs": {"width": 512, "height": 512, "batch_size": 1}},
            "7": {"class_type": "KSampler",
                  "inputs": {"model": ["1", 0], "seed": 42, "steps": 8, "cfg": 1.0,
                             "sampler_name": "euler", "scheduler": "simple",
                             "positive": ["4", 0], "negative": ["5", 0],
                             "latent_image": ["6", 0], "denoise": 1.0}},
            "8": {"class_type": "VAEDecode",
                  "inputs": {"samples": ["7", 0], "vae": ["3", 0]}},
            "9": {"class_type": "SaveImage",
                  "inputs": {"images": ["8", 0], "filename_prefix": "ZImageTurbo_test"}}
        }
    }
    
    wf_json = json.dumps(workflow)
    # Write workflow to file to avoid shell escaping issues
    run(f"cat > /tmp/zimage_wf.json << 'JSONEOF'\n{wf_json}\nJSONEOF")
    out, _ = run("curl -s -X POST http://localhost:8188/prompt "
                 "-H 'Content-Type: application/json' "
                 "-d @/tmp/zimage_wf.json")
    try:
        resp = json.loads(out.strip())
        if 'error' in resp:
            print(f"  API ERROR: {resp['error']}")
            sys.exit(1)
        print(f"  Prompt ID: {resp.get('prompt_id')}")
    except:
        print(f"  Response: {out.strip()[:500]}")

    # 6. Monitor — wait up to 15 minutes (first run can be slow due to kernel cache)
    print("\n=== Monitoring generation (GPU kernels may compile on first step) ===")
    last_log = ""
    for i in range(60):  # up to 15 minutes
        time.sleep(15)
        
        stats, _ = run("bash -c '"
            "PID=$(pgrep -f \"python3 main.py\" | head -1); "
            "if [ -n \"$PID\" ]; then "
            "  CPU=$(ps -p $PID -o %cpu --no-headers); "
            "  RSS=$(ps -p $PID -o rss --no-headers); "
            "  LOAD=$(cat /proc/loadavg | cut -d\" \" -f1); "
            "  GPU_T=$(cat /sys/class/drm/card1/device/hwmon/hwmon*/temp1_input 2>/dev/null || echo 0); "
            "  echo \"CPU:${CPU}% RSS:$((RSS/1024))M LOAD:${LOAD} GPU:$((GPU_T/1000))C\"; "
            "else echo DEAD; fi'")
        
        log, _ = run("tail -12 /home/fabian/comfyui.log 2>/dev/null")
        log_s = log.strip()
        
        elapsed = (i+1) * 15
        m, s = divmod(elapsed, 60)
        stats_s = stats.strip()
        print(f"  [{m}m{s:02d}s] {stats_s}")
        
        # Show new log content
        if log_s != last_log:
            for line in reversed(log_s.split('\n')):
                l = line.strip()
                if l and not l.startswith('FETCH') and not l.startswith('[DEPRECATION') and not l.startswith('[ComfyUI-Manager]'):
                    print(f"    LOG: {l[:120]}")
                    break
            last_log = log_s
        
        if 'DEAD' in stats_s:
            print("\n  PROCESS DIED!")
            out, _ = run("tail -60 /home/fabian/comfyui.log")
            print(out)
            break
        
        if 'Prompt executed in' in log_s:
            print(f"\n  SUCCESS! Image generated!")
            out, _ = run("tail -25 /home/fabian/comfyui.log")
            print(out)
            break
        
        if 'Traceback' in log_s or 'RuntimeError' in log_s:
            print("\n  ERROR detected!")
            out, _ = run("tail -60 /home/fabian/comfyui.log")
            print(out)
            break

    # 7. Check output
    print("\n=== Output files ===")
    out, _ = run("ls -lah ~/ComfyUI/output/ 2>/dev/null")
    print(out.strip())
    
    # 8. Check history
    out, _ = run("curl -s http://localhost:8188/history 2>/dev/null")
    try:
        h = json.loads(out)
        for pid, info in h.items():
            status = info.get('status', {})
            outputs = info.get('outputs', {})
            print(f"\n  Prompt {pid[:12]}...: status={status}")
            if outputs:
                for nid, nout in outputs.items():
                    if isinstance(nout, dict) and 'images' in nout:
                        for img in nout['images']:
                            print(f"    Image: {img.get('filename', 'unknown')}")
    except:
        pass

finally:
    ssh.close()
    print("\nSSH connection closed.")