#!/usr/bin/env python3
"""Fix startup: use proper subprocess instead of exec, set threads via sitecustomize."""
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[: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. Create a sitecustomize.py in the venv to set threads on import
sitecustomize = '''# Auto-set PyTorch threading to use all 12 CPU cores on BC-250
import os
os.environ.setdefault("OMP_NUM_THREADS", "12")
os.environ.setdefault("MKL_NUM_THREADS", "12")
os.environ.setdefault("OPENBLAS_NUM_THREADS", "12")

try:
    import torch
    torch.set_num_threads(12)
    torch.set_num_interop_threads(12)
except Exception:
    pass
'''

# Find the venv site-packages path
rc, out, _ = run("bash -c 'source ~/comfyui-env/bin/activate && python3 -c \"import site; print(site.getsitepackages()[0])\"'",
    desc="Find venv site-packages")
site_packages = out.strip()
print(f"  Site-packages: {site_packages}")

# Write sitecustomize.py
sftp = ssh.open_sftp()
sitecust_path = f"{site_packages}/sitecustomize.py"
# Check if it exists first
try:
    sftp.stat(sitecust_path)
    print(f"  sitecustomize.py already exists, backing up")
    sftp.rename(sitecust_path, f"{sitecust_path}.bak")
except FileNotFoundError:
    pass
with sftp.open(sitecust_path, 'w') as f:
    f.write(sitecustomize)
sftp.close()
print(f"  Written: {sitecust_path}")

# 2. Update startup script — simple, using exec python main.py directly
startup_script = r'''#!/bin/bash
# ComfyUI Startup Script for AsRock BC-250 (AMD Cyan Skillfish / ROCm 7.2)
# All 12 CPU cores + lowvram for 7.6GB shared VRAM
set -euo pipefail

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

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

# ═══════════════ Memory tuning ═══════════════
export PYTORCH_HIP_ALLOC_CONF="expandable_segments:False"
export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1

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

echo "=========================================="
echo " ComfyUI on BC-250 (ROCm 7.2)"
echo " GPU:    AMD Cyan Skillfish (gfx1010)"
echo " PyTorch: $(python3 -c 'import torch; print(torch.__version__)')"
echo " Threads: $(python3 -c 'import torch; print(f"intra={torch.get_num_threads()} inter={torch.get_num_interop_threads()}")')"
echo " CPU:    $(nproc) cores"
echo " VRAM:   7.6GB shared — lowvram mode"
echo "=========================================="

# Default args: listen on all, lowvram for tight VRAM
ARGS="--listen 0.0.0.0 --port 8188 --lowvram"
if [ $# -gt 0 ]; then
    ARGS="$@"
fi

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

exec python3 main.py $ARGS
'''

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 executable")

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

time.sleep(20)
run("bash -c 'tail -40 /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")

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

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