"""BC-250: Start ComfyUI on GPU and generate an image. Single SSH connection. No shell escaping issues."""
import paramiko
import time
import json
import sys
import textwrap

# ==== CONFIG ====
SSH_HOST = '192.168.178.150'
SSH_USER = 'fabian'
SSH_KEY  = r'C:\Users\fabia\.ssh\id_ed25519'

def connect():
    k = paramiko.Ed25519Key.from_private_key_file(SSH_KEY)
    c = paramiko.SSHClient()
    c.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    c.connect(SSH_HOST, username=SSH_USER, pkey=k, timeout=15)
    return c

def sh(c, cmd, timeout=60):
    """Run a bash command. All commands go through bash explicitly."""
    chan = c.get_transport().open_session()
    chan.settimeout(timeout)
    chan.exec_command(f'/bin/bash -l -c {_quote(cmd)}')
    out = b""
    while True:
        try:
            chunk = chan.recv(65536)
            if not chunk:
                break
            out += chunk
        except Exception:
            break
    chan.close()
    return out.decode(errors='replace').strip()

def _quote(s):
    """Shell-quote a string using single quotes."""
    return "'" + s.replace("'", "'\\''") + "'"

def write_remote_file(c, path, content):
    """Write a file on the remote via SFTP. No shell escaping needed."""
    sftp = c.open_sftp()
    with sftp.open(path, 'w') as f:
        f.write(content)
    sftp.close()

# ================================================================
print("="*60)
print("STEP 1: Connect + kill old ComfyUI")
print("="*60)
c = connect()
sh(c, 'pkill -9 -f "python3.*main.py" 2>/dev/null || true')
time.sleep(2)
alive = sh(c, 'pgrep -af "python3.*main.py" 2>/dev/null || echo NONE')
print(f"  Old processes: {alive}")

# ================================================================
print("\n" + "="*60)
print("STEP 2: Write launcher script on BC-250")
print("="*60)

# Write a bash launcher script directly via SFTP - avoids ALL shell escaping issues
launcher = textwrap.dedent("""\
    #!/bin/bash
    # GPU environment
    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
    # Threading
    export OMP_NUM_THREADS=12
    export MKL_NUM_THREADS=12
    export OPENBLAS_NUM_THREADS=12
    # MIOpen
    export MIOPEN_FIND_MODE=1
    
    cd ~/ComfyUI
    source ~/comfyui-env/bin/activate
    
    exec python3 main.py \\
        --listen 0.0.0.0 --port 8188 \\
        --lowvram \\
        --force-fp16 \\
        --cpu-vae \\
        --disable-smart-memory
""")

write_remote_file(c, '/tmp/run_comfyui.sh', launcher)
sh(c, 'chmod +x /tmp/run_comfyui.sh')
print("  Launcher script written to /tmp/run_comfyui.sh")
print("  Flags: --lowvram --force-fp16 --cpu-vae --disable-smart-memory")

# ================================================================
print("\n" + "="*60)
print("STEP 3: Verify GPU works with PyTorch")
print("="*60)

gpu_script = textwrap.dedent("""\
    #!/bin/bash
    export HSA_OVERRIDE_GFX_VERSION=10.1.0
    export HIP_VISIBLE_DEVICES=0
    export HSA_ENABLE_SDMA=0
    source ~/comfyui-env/bin/activate
    python3 -c "
import torch
print('PyTorch:', torch.__version__)
print('CUDA/ROCm available:', torch.cuda.is_available())
if torch.cuda.is_available():
    print('Device:', torch.cuda.get_device_name(0))
    f,t = torch.cuda.mem_get_info(0)
    print(f'VRAM: {f//1048576}MB free / {t//1048576}MB total')
    x = torch.randn(512,512,device='cuda',dtype=torch.float16)
    y = x @ x
    print('GPU compute test: PASS')
else:
    print('FATAL: NO GPU')
    exit(1)
"
""")
write_remote_file(c, '/tmp/gpu_test.sh', gpu_script)
sh(c, 'chmod +x /tmp/gpu_test.sh')
out = sh(c, '/tmp/gpu_test.sh', timeout=30)
print(f"  {out}")
if 'FATAL' in out or 'False' in out:
    print("  *** GPU not working! Aborting. ***")
    c.close()
    sys.exit(1)
print("  GPU OK!")

# ================================================================
print("\n" + "="*60)
print("STEP 4: Start ComfyUI")
print("="*60)

sh(c, 'rm -f /tmp/comfyui.log; touch /tmp/comfyui.log')
sh(c, 'nohup /tmp/run_comfyui.sh > /tmp/comfyui.log 2>&1 &')
time.sleep(3)

pid = sh(c, 'pgrep -f "python3.*main.py" 2>/dev/null || echo DEAD')
if pid == 'DEAD':
    print("  FAILED to start! Log:")
    print(sh(c, 'cat /tmp/comfyui.log'))
    c.close()
    sys.exit(1)
print(f"  PID: {pid}")

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

# Show startup flags from log
log_head = sh(c, 'head -10 /tmp/comfyui.log')
print(f"\n  Startup log:\n  {log_head[:300]}")

# ================================================================
print("\n" + "="*60)
print("STEP 5: Submit workflow")
print("="*60)

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"}
        }
    }
}

write_remote_file(c, '/tmp/wf.json', json.dumps(workflow))
# Verify it's valid JSON with correct nodes
verify = sh(c, 'python3 -c "import json; d=json.load(open(\'/tmp/wf.json\')); p=d[\'prompt\']; print(len(p), \'nodes:\', sorted(p.keys()))"')
print(f"  Workflow: {verify}")

# Submit
resp = sh(c, '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("  *** SUBMIT FAILED! ***")
    print(f"  Full response: {resp}")
    print(f"  Log: {sh(c, 'tail -10 /tmp/comfyui.log')}")
    c.close()
    sys.exit(1)

prompt_id = json.loads(resp).get('prompt_id', '?')
print(f"  Prompt ID: {prompt_id}")

# ================================================================
print("\n" + "="*60)
print("STEP 6: Monitor generation (checking GPU usage)")
print("="*60)

t0 = time.time()
for i in range(200):  # up to ~50 min
    elapsed = int(time.time() - t0)
    
    gpu_pct  = sh(c, 'cat /sys/class/drm/card0/device/gpu_busy_percent 2>/dev/null || echo ?', timeout=5)
    gpu_temp = sh(c, 'cat /sys/class/drm/card0/device/hwmon/hwmon*/temp1_input 2>/dev/null || echo 0', timeout=5)
    temp_c   = int(gpu_temp) // 1000 if gpu_temp.isdigit() else '?'
    
    log_tail = sh(c, 'tail -3 /tmp/comfyui.log 2>/dev/null', timeout=5)
    last_line = log_tail.strip().split('\n')[-1] if log_tail else ''
    
    # Check for output image
    imgs = sh(c, 'ls ~/ComfyUI/output/*.png 2>/dev/null || echo NONE', timeout=5)
    
    print(f"  [{elapsed:>4}s] GPU:{gpu_pct:>3}% {temp_c}C | {last_line[-90:]}")
    
    if imgs != 'NONE':
        print(f"\n  >>> IMAGE GENERATED! <<<")
        print(f"  Files: {imgs}")
        print(f"  Time: {elapsed}s")
        final = sh(c, 'tail -20 /tmp/comfyui.log')
        print(f"\n  Final log:\n{final}")
        break
    
    # Check queue empty (= done or error)
    q = sh(c, 'curl -s http://127.0.0.1:8188/queue 2>/dev/null || echo {}', timeout=5)
    try:
        qd = json.loads(q)
        if not qd.get('queue_running') and not qd.get('queue_pending') and elapsed > 20:
            time.sleep(3)
            imgs2 = sh(c, 'ls ~/ComfyUI/output/*.png 2>/dev/null || echo NONE', timeout=5)
            if imgs2 != 'NONE':
                print(f"\n  >>> IMAGE GENERATED! <<<")
                print(f"  Files: {imgs2}")
                print(f"  Time: {elapsed}s")
            else:
                print(f"\n  Queue empty, no image. Checking log for errors...")
                print(sh(c, 'tail -30 /tmp/comfyui.log'))
            break
    except json.JSONDecodeError:
        pass
    
    # Check process still alive
    alive = sh(c, 'pgrep -f "python3.*main.py" >/dev/null 2>&1 && echo YES || echo NO', timeout=5)
    if alive == 'NO':
        print(f"\n  *** ComfyUI CRASHED! ***")
        print(sh(c, 'tail -40 /tmp/comfyui.log'))
        break
    
    time.sleep(15)

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