#!/usr/bin/env python3
"""Create ComfyUI startup script and launch it on BC-250."""
import paramiko
import time

ssh = paramiko.SSHClient()
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
ssh.connect('192.168.178.150', username='fabian', key_filename=r'C:\Users\fabia\.ssh\id_ed25519')

def run(cmd, timeout=120, desc=""):
    if desc:
        print(f"\n{'='*60}")
        print(f"  {desc}")
        print(f"{'='*60}")
    print(f"$ {cmd}")
    _, stdout, stderr = ssh.exec_command(cmd, timeout=timeout)
    out = stdout.read().decode()
    err = stderr.read().decode()
    rc = stdout.channel.recv_exit_status()
    if out.strip():
        lines = out.strip().split('\n')
        if len(lines) > 40:
            print(f"  ... ({len(lines)} lines, showing last 40)")
            print('\n'.join(lines[-40:]))
        else:
            print(out.strip())
    if err.strip():
        lines = err.strip().split('\n')
        show = lines[-20:] if len(lines) > 20 else lines
        print(f"STDERR: {chr(10).join(show)}")
    print(f"  Exit code: {rc}")
    return rc, out, err

# ──────────────────────────────────────────────────────────
# 1. Create the startup script
# ──────────────────────────────────────────────────────────
startup_script = r'''#!/bin/bash
# ComfyUI Startup Script for AsRock BC-250 (AMD Cyan Skillfish / ROCm 7.2)
# Usage: ~/start_comfyui.sh [--listen] [--port PORT]

set -euo pipefail

# ═══════════════════════════════════════════════════════════
#  BC-250 AMD GPU Environment Variables
# ═══════════════════════════════════════════════════════════
# Override gfx1013 → gfx1010 (RDNA 1.5 → RDNA 1 compat)
export HSA_OVERRIDE_GFX_VERSION=10.1.0

# Use device 0
export HIP_VISIBLE_DEVICES=0

# Disable SDMA (avoids queue errors on Cyan Skillfish)
export HSA_ENABLE_SDMA=0

# Suppress tool library warnings
export HSA_TOOLS_LIB=""
export HSA_TOOLS_REPORT_LOAD_FAILURE=0

# PyTorch / ROCm tuning
export PYTORCH_HIP_ALLOC_CONF="expandable_segments:True"
export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1

# Avoid OOM on 14GB shared VRAM — force float16 where possible
export COMFY_PRECISION=16

# ═══════════════════════════════════════════════════════════
#  Activate Virtual Environment
# ═══════════════════════════════════════════════════════════
source "$HOME/comfyui-env/bin/activate"

# ═══════════════════════════════════════════════════════════
#  Launch ComfyUI
# ═══════════════════════════════════════════════════════════
cd "$HOME/ComfyUI"

echo "=========================================="
echo " ComfyUI on BC-250 (ROCm 7.2)"
echo "=========================================="
echo " GPU:    AMD Cyan Skillfish (gfx1013→gfx1010)"
echo " PyTorch: $(python -c 'import torch; print(torch.__version__)')"
echo " HIP:    $(python -c 'import torch; print(torch.version.hip)')"
echo " CUDA:   $(python -c 'import torch; print(torch.cuda.is_available())')"
echo " Device: $(python -c 'import torch; print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else "N/A")')"
echo "=========================================="

# Default: listen on all interfaces for remote access
LISTEN_ARGS="--listen 0.0.0.0 --port 8188"

# Parse arguments (override defaults if provided)
if [ $# -gt 0 ]; then
    LISTEN_ARGS="$@"
fi

echo ""
echo "Starting ComfyUI with: $LISTEN_ARGS"
echo "Access at: http://$(hostname -I | awk '{print $1}'):8188"
echo ""

exec python main.py $LISTEN_ARGS
'''

# ──────────────────────────────────────────────────────────
# 2. Create fish shell wrapper too
# ──────────────────────────────────────────────────────────
fish_script = r'''#!/usr/bin/env fish
# ComfyUI launcher for fish shell on BC-250

# BC-250 GPU env vars
set -gx HSA_OVERRIDE_GFX_VERSION 10.1.0
set -gx HIP_VISIBLE_DEVICES 0
set -gx HSA_ENABLE_SDMA 0
set -gx HSA_TOOLS_LIB ""
set -gx HSA_TOOLS_REPORT_LOAD_FAILURE 0
set -gx PYTORCH_HIP_ALLOC_CONF "expandable_segments:True"
set -gx TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL 1

# Activate venv
source $HOME/comfyui-env/bin/activate.fish

# Launch
cd $HOME/ComfyUI
echo "Starting ComfyUI on BC-250..."
python main.py --listen 0.0.0.0 --port 8188 $argv
'''

# Upload scripts
sftp = ssh.open_sftp()
with sftp.open('/home/fabian/start_comfyui.sh', 'w') as f:
    f.write(startup_script)
with sftp.open('/home/fabian/start_comfyui.fish', 'w') as f:
    f.write(fish_script)
sftp.close()

run("chmod +x /home/fabian/start_comfyui.sh /home/fabian/start_comfyui.fish",
    desc="Make startup scripts executable")

# ──────────────────────────────────────────────────────────
# 3. Quick pre-flight check
# ──────────────────────────────────────────────────────────
run("bash -c 'source /home/fabian/comfyui-env/bin/activate && "
    "export HSA_OVERRIDE_GFX_VERSION=10.1.0 && "
    "export HIP_VISIBLE_DEVICES=0 && "
    "export HSA_ENABLE_SDMA=0 && "
    "cd /home/fabian/ComfyUI && "
    "python -c \""
    "import torch; "
    "print(f\\\"PyTorch {torch.__version__}, HIP {torch.version.hip}, CUDA {torch.cuda.is_available()}\\\"); "
    "print(f\\\"Device: {torch.cuda.get_device_name(0)}\\\"); "
    "import comfy.model_management; "
    "print(f\\\"ComfyUI model_management imported OK\\\"); "
    "\"'",
    desc="Pre-flight: PyTorch + ComfyUI import check")

# ──────────────────────────────────────────────────────────
# 4. Launch ComfyUI in background
# ──────────────────────────────────────────────────────────
run("bash -c 'pkill -f \"python main.py\" 2>/dev/null; echo killed || echo no_existing'",
    desc="Kill any existing ComfyUI process")

time.sleep(2)

run("bash -c 'nohup bash /home/fabian/start_comfyui.sh > /home/fabian/comfyui.log 2>&1 & echo PID=$!'",
    desc="Launch ComfyUI in background")

# Wait for startup
time.sleep(15)
run("tail -30 /home/fabian/comfyui.log 2>/dev/null", desc="ComfyUI startup log")

# Check if port is listening
time.sleep(10)
run("bash -c 'ss -tlnp | grep 8188 || echo PORT_NOT_LISTENING'",
    desc="Check if port 8188 is listening")

run("tail -50 /home/fabian/comfyui.log 2>/dev/null", desc="Full startup log")

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