#!/usr/bin/env python3
"""Fix threading: set all 12 cores for PyTorch ops, patch ComfyUI startup, restart."""
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. Kill any existing ComfyUI
run("bash -c 'pkill -9 -f \"python main.py\" 2>/dev/null; sleep 1; echo killed'",
    desc="Kill existing ComfyUI")

# 2. Check the loader.py for the Dequantizing message source
run("bash -c 'grep -rn \"Dequantizing\" ~/ComfyUI/ --include=\"*.py\" 2>/dev/null | head -10'",
    desc="Find 'Dequantizing' message source")

# 3. Check torch thread defaults without any env vars
run("bash -c 'source ~/comfyui-env/bin/activate && python3 -c \"import torch; print(torch.get_num_threads(), torch.get_num_interop_threads())\"'",
    desc="Default torch thread count")

# 4. Verify it works with env vars
run("bash -c 'export OMP_NUM_THREADS=12; export MKL_NUM_THREADS=12; "
    "source ~/comfyui-env/bin/activate && python3 -c \"import torch; "
    "torch.set_num_threads(12); torch.set_num_interop_threads(4); "
    "print(torch.get_num_threads(), torch.get_num_interop_threads())\"'",
    desc="Torch threads with env vars set to 12")

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