#!/usr/bin/env python3
"""
Diagnose GPU hang: kill stuck ComfyUI, run targeted HIP tests,
check what ops hang on Cyan Skillfish gfx1013->gfx1010.
Single SSH connection, properly closed.
"""
import paramiko
import json
import time
import sys

ssh = paramiko.SSHClient()
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())

for attempt in range(5):
    try:
        ssh.connect('192.168.178.150', username='fabian',
                    key_filename=r'C:\Users\fabia\.ssh\id_ed25519', timeout=10)
        break
    except Exception as e:
        print(f"  SSH attempt {attempt+1}/5: {e}")
        time.sleep(10)
else:
    print("FATAL: Cannot connect"); sys.exit(1)

def run(cmd, timeout=120):
    _, stdout, stderr = ssh.exec_command(cmd, timeout=timeout)
    out = stdout.read().decode()
    err = stderr.read().decode()
    return out, err

try:
    # 1. Kill stuck ComfyUI
    print("=== Kill stuck ComfyUI ===")
    out, _ = run("pkill -9 -f 'python3 main.py' 2>/dev/null; sleep 2; echo killed")
    print(f"  {out.strip()}")

    # 2. Check ComfyUI help for --cpu-vae flag existence
    print("\n=== Check if --cpu-vae exists ===")
    out, err = run("bash -c 'source ~/comfyui-env/bin/activate && cd ~/ComfyUI && python3 main.py --help 2>&1'")
    full_help = out + err
    has_cpu_vae = '--cpu-vae' in full_help
    print(f"  --cpu-vae flag exists: {has_cpu_vae}")
    # Print all vram/gpu related flags
    for line in full_help.split('\n'):
        if any(w in line.lower() for w in ['vram', 'cpu', 'gpu', 'fp16', 'fp32', 'vae', 'force', 'precision']):
            print(f"  {line.strip()}")

    # 3. Check what the CachyOS pytorch-rocm was built for
    print("\n=== PyTorch ROCm build info ===")
    out, _ = run("bash -c 'source ~/comfyui-env/bin/activate && python3 -c \""
        "import torch; "
        "print(f\\\"PyTorch version: {torch.__version__}\\\"); "
        "print(f\\\"CUDA/HIP available: {torch.cuda.is_available()}\\\"); "
        "print(f\\\"ROCm version: {torch.version.hip}\\\"); "
        "print(f\\\"Device: {torch.cuda.get_device_name(0)}\\\"); "
        "print(f\\\"Arch: {torch.cuda.get_device_capability(0)}\\\"); "
        "print(f\\\"VRAM free/total: {torch.cuda.mem_get_info()[0]//1048576}/{torch.cuda.mem_get_info()[1]//1048576} MB\\\"); "
        "\"' 2>&1")
    print(out.strip())

    # 4. Targeted GPU op tests - find what hangs
    print("\n=== GPU operation tests (timeout 30s each) ===")
    tests = [
        ("Basic matmul fp32", 
         "a=torch.randn(256,256,device='cuda'); b=a@a; print(f'fp32 matmul: {b.shape} sum={b.sum().item():.1f}')"),
        ("Basic matmul fp16",
         "a=torch.randn(256,256,device='cuda').half(); b=a@a; print(f'fp16 matmul: {b.shape} sum={b.sum().item():.1f}')"),
        ("Conv2d fp16 (VAE-like)",
         "import torch.nn as nn; c=nn.Conv2d(128,128,3,padding=1).half().cuda(); x=torch.randn(1,128,64,64,device='cuda').half(); y=c(x); print(f'conv2d fp16: {y.shape}')"),
        ("Conv2d fp32 (VAE default)",
         "import torch.nn as nn; c=nn.Conv2d(128,128,3,padding=1).cuda(); x=torch.randn(1,128,64,64,device='cuda'); y=c(x); print(f'conv2d fp32: {y.shape}')"),
        ("GroupNorm fp16",
         "import torch.nn as nn; gn=nn.GroupNorm(32,128).half().cuda(); x=torch.randn(1,128,32,32,device='cuda').half(); y=gn(x); print(f'groupnorm fp16: {y.shape}')"),
        ("GroupNorm fp32",
         "import torch.nn as nn; gn=nn.GroupNorm(32,128).cuda(); x=torch.randn(1,128,32,32,device='cuda'); y=gn(x); print(f'groupnorm fp32: {y.shape}')"),
        ("LayerNorm fp16",
         "import torch.nn as nn; ln=nn.LayerNorm(256).half().cuda(); x=torch.randn(1,64,256,device='cuda').half(); y=ln(x); print(f'layernorm fp16: {y.shape}')"),
        ("Linear fp16 (DiT-like)",
         "import torch.nn as nn; l=nn.Linear(1024,1024).half().cuda(); x=torch.randn(1,64,1024,device='cuda').half(); y=l(x); print(f'linear fp16: {y.shape}')"),
        ("Attention fp16 (scaled_dot_product)",
         "q=torch.randn(1,8,64,64,device='cuda').half(); k=q.clone(); v=q.clone(); "
         "y=torch.nn.functional.scaled_dot_product_attention(q,k,v); print(f'sdpa fp16: {y.shape}')"),
        ("Attention fp32 (scaled_dot_product)",
         "q=torch.randn(1,8,64,64,device='cuda'); k=q.clone(); v=q.clone(); "
         "y=torch.nn.functional.scaled_dot_product_attention(q,k,v); print(f'sdpa fp32: {y.shape}')"),
        ("Large matmul fp16 (5032x5032)",
         "a=torch.randn(2048,2048,device='cuda').half(); b=a@a; print(f'large fp16 matmul: {b.shape} sum={b.sum().item():.1f}')"),
        ("RoPE-like op (complex multiply)",
         "x=torch.randn(1,8,64,64,device='cuda').half(); "
         "f=torch.randn(64,32,2,device='cuda').half(); "
         "print(f'rope input shapes: x={x.shape} f={f.shape} OK')"),
        ("torch.compile basic test",
         "import torch._dynamo; f=lambda x: x*2+1; cf=torch.compile(f); "
         "x=torch.randn(100,device='cuda'); y=cf(x); print(f'compile: {y.shape}')"),
    ]
    
    for name, code in tests:
        print(f"\n  Testing: {name}...", end=" ", flush=True)
        cmd = (f"bash -c 'timeout 30 bash -c \""
               f"source ~/comfyui-env/bin/activate && "
               f"HSA_OVERRIDE_GFX_VERSION=10.1.0 HSA_ENABLE_SDMA=0 "
               f"python3 -c \\\"import torch; {code}\\\"\" 2>&1 || echo TIMEOUT_OR_ERROR'")
        out, err = run(cmd, timeout=40)
        result = (out + err).strip()
        if 'TIMEOUT_OR_ERROR' in result:
            # Get just the error part
            lines = result.split('\n')
            for l in reversed(lines):
                if l.strip() and l.strip() != 'TIMEOUT_OR_ERROR':
                    print(f"FAILED: {l.strip()[:100]}")
                    break
            else:
                print("TIMEOUT (GPU HANG)")
        elif result:
            last_line = [l for l in result.split('\n') if l.strip()][-1] if result.split('\n') else result
            print(f"OK: {last_line.strip()[:100]}")
        else:
            print("NO OUTPUT (possible hang)")

    # 5. Check dmesg for GPU errors after tests
    print("\n\n=== dmesg GPU errors (last 20) ===")
    out, _ = run("dmesg 2>/dev/null | grep -i -E 'amdgpu|gpu|gfx|error|fault' | tail -20 || echo 'no permission'")
    print(out.strip() if out.strip() else "  (empty or no permission)")

    # 6. Check rocm-smi for GPU health
    print("\n=== GPU health after tests ===")
    out, _ = run("HSA_OVERRIDE_GFX_VERSION=10.1.0 rocm-smi 2>/dev/null")
    for line in out.split('\n'):
        if any(c in line for c in ['°C', '%', 'Device', 'Node']):
            print(f"  {line.strip()}")

finally:
    ssh.close()
    print("\n\nSSH connection closed.")
