#!/usr/bin/env python3
"""Deep diagnosis: WHY only 1 core? Check OpenMP, threading, GGUF code path."""
import paramiko, json, textwrap
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=60):
    _, stdout, stderr = ssh.exec_command(cmd, timeout=timeout)
    out = stdout.read().decode()
    err = stderr.read().decode()
    rc = stdout.channel.recv_exit_status()
    return rc, out, err

def show(label, cmd, timeout=60):
    print(f"\n{'='*60}")
    print(f"  {label}")
    print(f"{'='*60}")
    rc, out, err = run(cmd, timeout)
    if out.strip():
        print(out.strip())
    if err.strip():
        for line in err.strip().split('\n')[-10:]:
            print(f"STDERR: {line}")
    return out

# 1. Check ComfyUI queue/history - did generation succeed or fail?
show("Queue status", "curl -s http://localhost:8188/queue")
hist_out = show("History", "curl -s http://localhost:8188/history")
try:
    h = json.loads(hist_out.strip())
    for pid, info in h.items():
        status = info.get('status', {})
        print(f"\n  Prompt {pid}: status={status}")
        outputs = info.get('outputs', {})
        if outputs:
            print(f"  Outputs: {json.dumps(outputs, indent=2)[:500]}")
        else:
            print("  NO OUTPUTS")
except:
    pass

# 2. Check if PyTorch has OpenMP 
show("PyTorch OpenMP & threading", 
    "bash -c 'source /home/fabian/comfyui-env/bin/activate && "
    "OMP_NUM_THREADS=12 python3 -c \""
    "import torch; "
    "print(f\\\"OpenMP available: {torch.backends.openmp.is_available()}\\\"); "
    "print(f\\\"MKL available: {torch.backends.mkl.is_available()}\\\"); "
    "print(f\\\"Num threads: {torch.get_num_threads()}\\\"); "
    "print(f\\\"Num interop threads: {torch.get_num_interop_threads()}\\\"); "
    "print(f\\\"torch.__config__.show():  \\\"); "
    "print(torch.__config__.show()); "
    "\"'")

# 3. Check if libomp/libgomp is available
show("OpenMP libraries", 
    "bash -c 'ldconfig -p 2>/dev/null | grep -i omp; "
    "echo ---; "
    "pacman -Qs openmp 2>/dev/null; "
    "echo ---; "
    "pacman -Qs libgomp 2>/dev/null; "
    "echo ---; "
    "ls -la /usr/lib/libomp* /usr/lib/libgomp* 2>/dev/null || echo none'")

# 4. Check pytorch shared lib dependencies for OpenMP
show("PyTorch .so OpenMP deps",
    "bash -c 'ldd /usr/lib/python3.14/site-packages/torch/lib/libtorch_cpu.so 2>/dev/null | grep -i omp'")

# 5. Actual thread test - does a matrix multiply use multiple cores?
show("Matrix multiply CPU benchmark (should use all cores)",
    "bash -c 'source /home/fabian/comfyui-env/bin/activate && "
    "OMP_NUM_THREADS=12 python3 -c \""
    "import torch, time, os; "
    "print(f\\\"PID: {os.getpid()}\\\"); "
    "torch.set_num_threads(12); "
    "print(f\\\"Threads set to: {torch.get_num_threads()}\\\"); "
    "a = torch.randn(4096, 4096); "
    "b = torch.randn(4096, 4096); "
    "# warmup; "
    "c = torch.mm(a, b); "
    "import subprocess; "
    "# Start monitoring in background; "
    "start = time.time(); "
    "for i in range(5): c = torch.mm(a, b); "
    "elapsed = time.time() - start; "
    "print(f\\\"5x matmul 4096x4096: {elapsed:.2f}s\\\"); "
    "\"'")

# 6. Check what the GGUF dequant code actually does (single-threaded python loop?)
show("GGUF dequant code - is there a Python for-loop?",
    "bash -c 'grep -n \"for \" /home/fabian/ComfyUI/custom_nodes/ComfyUI-GGUF/dequant.py | head -20; "
    "echo \"---\"; "
    "grep -n \"for \" /home/fabian/ComfyUI/custom_nodes/ComfyUI-GGUF/ops.py | head -20; "
    "echo \"---\"; "
    "grep -n \"def dequantize\" /home/fabian/ComfyUI/custom_nodes/ComfyUI-GGUF/dequant.py'")

# 7. Check if lowvram is causing sequential layer-by-layer processing
show("ComfyUI lowvram model loading code",
    "bash -c 'grep -rn \"lowvram\\|low_vram\\|offload\" /home/fabian/ComfyUI/comfy/model_management.py 2>/dev/null | head -30'")

# 8. Output directory
show("Output files", "ls -lah /home/fabian/ComfyUI/output/ 2>/dev/null")

ssh.close()
print("\n\nDONE.")