#!/usr/bin/env python3
"""Read full ops.py and check torch thread defaults, then fix threading."""
import paramiko

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=30, desc=""):
    if desc:
        print(f"\n{'='*60}")
        print(f"  {desc}")
        print(f"{'='*60}")
    _, 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():
        print(out.strip())
    if err.strip():
        lines = err.strip().split('\n')[-10:]
        print(f"STDERR: {chr(10).join(lines)}")
    print(f"  Exit: {rc}")
    return rc, out, err

# Check current torch thread defaults
run("bash -c 'source ~/comfyui-env/bin/activate && "
    "export HSA_OVERRIDE_GFX_VERSION=10.1.0 && "
    "python -c \""
    "import torch; "
    "print(f\\\"num_threads={torch.get_num_threads()}\\\"); "
    "print(f\\\"num_interop_threads={torch.get_num_interop_threads()}\\\"); "
    "import os; "
    "print(f\\\"OMP_NUM_THREADS={os.environ.get(\\\\\\\"OMP_NUM_THREADS\\\\\\\", \\\\\\\"not set\\\\\\\")}\\\"); "
    "print(f\\\"MKL_NUM_THREADS={os.environ.get(\\\\\\\"MKL_NUM_THREADS\\\\\\\", \\\\\\\"not set\\\\\\\")}\\\"); "
    "\"'",
    desc="Check default torch thread settings")

# Read forward_ggml_cast_weights (where dequant happens during inference)
run("bash -c 'sed -n \"200,281p\" ~/ComfyUI/custom_nodes/ComfyUI-GGUF/ops.py'",
    desc="ops.py lines 200-281 (forward functions)")

# Check __init__.py for any loading/patching
run("bash -c 'cat ~/ComfyUI/custom_nodes/ComfyUI-GGUF/__init__.py 2>/dev/null | head -40'",
    desc="ComfyUI-GGUF __init__.py")

# Check nodes.py for model loading
run("bash -c 'cat ~/ComfyUI/custom_nodes/ComfyUI-GGUF/nodes.py 2>/dev/null'",
    desc="ComfyUI-GGUF nodes.py (model loader)")

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