#!/usr/bin/env python3
"""Fix: maximize threading to 12 cores, add --lowvram for 7.6GB VRAM, restart ComfyUI."""
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=300, desc=""):
    if desc:
        print(f"\n{'='*60}")
        print(f"  {desc}")
        print(f"{'='*60}")
    print(f"$ {cmd[:300]}")
    _, 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[-15:] if len(lines) > 15 else lines
        print(f"STDERR: {chr(10).join(show)}")
    print(f"  Exit code: {rc}")
    return rc, out, err

# 1. Kill any existing ComfyUI
run("bash -c 'pkill -9 -f \"python main.py\" 2>/dev/null; sleep 2; echo ok'",
    desc="Kill existing ComfyUI")

# 2. Write updated start_comfyui.sh with full threading + lowvram
startup_script = r'''#!/bin/bash
# ComfyUI Startup Script for AsRock BC-250 (AMD Cyan Skillfish / ROCm 7.2)
# Optimized for all 12 CPU cores + 7.6GB shared VRAM
set -euo pipefail

# ═══════════════ BC-250 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

# ═══════════════ THREADING — ALL 12 CORES ═══════════════
export OMP_NUM_THREADS=12
export MKL_NUM_THREADS=12
export OPENBLAS_NUM_THREADS=12
export VECLIB_MAXIMUM_THREADS=12
export NUMEXPR_NUM_THREADS=12
# PyTorch intra-op (tensor math) and inter-op (parallel node execution) threads
export TORCH_NUM_THREADS=12

# ═══════════════ Memory / ROCm tuning ═══════════════
export PYTORCH_HIP_ALLOC_CONF="expandable_segments:False"
export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1
# Disable HIP memory caching to avoid fragmentation on shared VRAM
export PYTORCH_NO_HIP_MEMORY_CACHING=0

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

cd "$HOME/ComfyUI"

# Force PyTorch to use all 12 cores
python3 -c "import torch; torch.set_num_threads(12); torch.set_num_interop_threads(12); print(f'Threads: intra={torch.get_num_threads()} inter={torch.get_num_interop_threads()}')"

echo "=========================================="
echo " ComfyUI on BC-250 (ROCm 7.2)"
echo " GPU:    AMD Cyan Skillfish (gfx1013→gfx1010)"
echo " PyTorch: $(python3 -c 'import torch; print(torch.__version__)')"
echo " HIP:    $(python3 -c 'import torch; print(torch.version.hip)')"
echo " CUDA:   $(python3 -c 'import torch; print(torch.cuda.is_available())')"
echo " CPU:    $(nproc) cores (all used)"
echo " VRAM:   7.6GB shared — using --lowvram mode"
echo "=========================================="

# Default: listen on all interfaces, --lowvram for 7.6GB shared VRAM
LISTEN_ARGS="--listen 0.0.0.0 --port 8188 --lowvram"

if [ $# -gt 0 ]; then
    LISTEN_ARGS="$@"
fi

echo "Starting: python main.py $LISTEN_ARGS"
echo "Access at: http://192.168.178.150:8188"
echo ""

# Set threads inside the actual process too
exec python3 -c "
import torch, sys, os
torch.set_num_threads(12)
torch.set_num_interop_threads(12)
# Now exec ComfyUI main
sys.argv = ['main.py'] + '$LISTEN_ARGS'.split()
exec(open('main.py').read())
"
'''

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

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

# 3. Launch ComfyUI with new settings
run("bash -c 'rm -f /home/fabian/comfyui.log'", desc="Clean old log")
run("bash -c 'nohup bash /home/fabian/start_comfyui.sh > /home/fabian/comfyui.log 2>&1 & echo PID=$!'",
    desc="Launch ComfyUI with 12-core threading + lowvram")

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

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

# Verify threads are set
run("bash -c 'tail -40 /home/fabian/comfyui.log 2>/dev/null'",
    desc="Full startup log")

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