#!/usr/bin/env python3
"""Fix single-core bottleneck: set threading env vars and patch ComfyUI-GGUF for parallel dequant."""
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) > 50:
            print(f"  ... ({len(lines)} lines, showing last 50)")
            print('\n'.join(lines[-50:]))
        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 stuck ComfyUI
run("bash -c 'kill -9 484588 2>/dev/null; pkill -9 -f \"python main.py\" 2>/dev/null; sleep 2; echo done'",
    desc="Kill stuck ComfyUI process")

# 2. Check how many CPU cores
run("bash -c 'nproc'", desc="CPU core count")

# 3. Check the ComfyUI-GGUF dequant code to understand the bottleneck
run("bash -c 'grep -rn \"dequant\\|num_threads\\|torch.set_num_threads\\|ThreadPool\\|parallel\" ~/ComfyUI/custom_nodes/ComfyUI-GGUF/*.py 2>/dev/null | head -30'",
    desc="Search for threading in ComfyUI-GGUF")

run("bash -c 'cat ~/ComfyUI/custom_nodes/ComfyUI-GGUF/ops.py 2>/dev/null | head -80'",
    desc="ComfyUI-GGUF ops.py (dequant logic)")

# 4. Check the dequant function
run("bash -c 'grep -n \"def dequantize\\|class GGMLTensor\\|def forward\" ~/ComfyUI/custom_nodes/ComfyUI-GGUF/ops.py 2>/dev/null'",
    desc="Key functions in ops.py")

run("bash -c 'wc -l ~/ComfyUI/custom_nodes/ComfyUI-GGUF/ops.py ~/ComfyUI/custom_nodes/ComfyUI-GGUF/dequant.py 2>/dev/null'",
    desc="File sizes")

run("bash -c 'cat ~/ComfyUI/custom_nodes/ComfyUI-GGUF/dequant.py 2>/dev/null | head -60'",
    desc="dequant.py start")

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