import paramiko
import sys

c = paramiko.SSHClient()
c.set_missing_host_key_policy(paramiko.AutoAddPolicy())
c.connect('192.168.178.150', username='fabian', key_filename=r'C:\Users\fabia\.ssh\id_ed25519')

def run(cmd, timeout=30):
    stdin, stdout, stderr = c.exec_command(f"bash -c '{cmd}'", timeout=timeout)
    out = stdout.read().decode()
    err = stderr.read().decode()
    rc = stdout.channel.recv_exit_status()
    return out.strip(), err.strip(), rc

def section(title):
    print(f"\n{'='*50}")
    print(f" {title}")
    print(f"{'='*50}")

# Test 1: Module status
section("Test 1: amdgpu Module")
out, _, _ = run("lsmod | grep amdgpu | head -3")
print(out)

# Test 2: v3 Module messages
section("Test 2: v3 Module (GFXOFF, BC-250)")
out, _, _ = run('sudo dmesg | grep -E "BC-250|GFXOFF|out-of-tree" | head -5')
print(out)

# Test 3: KIQ errors
section("Test 3: KIQ Fence Timeouts")
out, _, _ = run('sudo dmesg | grep -ci "timeout waiting for kiq fence" 2>/dev/null || echo 0')
print(f"KIQ timeout count: {out}")

# Test 4: GPU dead events
section("Test 4: GPU Unreachable/Dead Events")
out, _, _ = run('sudo dmesg | grep -ci "GPU unreachable\\|GPU died" 2>/dev/null || echo 0')
print(f"GPU dead count: {out}")

# Test 5: ppfeaturemask
section("Test 5: ppfeaturemask")
out, _, _ = run("cat /sys/module/amdgpu/parameters/ppfeaturemask")
print(f"ppfeaturemask: {out}")

# Test 6: Devices
section("Test 6: DRM/KFD Devices")
out, _, _ = run("ls -la /dev/dri/ 2>&1; echo; ls -la /dev/kfd 2>&1")
print(out)

# Test 7: GPU clocks
section("Test 7: GPU Clock Levels")
out, _, _ = run("cat /sys/class/drm/card0/device/pp_dpm_sclk 2>&1")
print(out)

# Test 8: Governor
section("Test 8: Cyan Skillfish Governor")
out, _, _ = run("systemctl is-active cyan-skillfish-governor.service 2>&1")
print(f"Governor: {out}")

# Test 9: rocminfo
section("Test 9: rocminfo")
out, _, _ = run('export HSA_OVERRIDE_GFX_VERSION=10.1.0 HIP_VISIBLE_DEVICES=0 HSA_ENABLE_SDMA=0 HSA_TOOLS_LIB= HSA_TOOLS_REPORT_LOAD_FAILURE=0; rocminfo 2>&1 | grep -E "Name:|Marketing Name:|Compute Unit:|gfx|Done"', timeout=60)
print(out)

# Test 10: hip_probe
section("Test 10: hip_probe (Dani's diagnostic)")
out, _, rc = run('export HSA_OVERRIDE_GFX_VERSION=10.1.0 HIP_VISIBLE_DEVICES=0 HSA_ENABLE_SDMA=0 HSA_TOOLS_LIB= HSA_TOOLS_REPORT_LOAD_FAILURE=0; /tmp/hip_probe 2>&1', timeout=60)
print(out)
print(f"Exit code: {rc}")

# Check KIQ after hip_probe
import time
time.sleep(3)
out, _, _ = run('sudo dmesg | tail -20 | grep -ci "timeout waiting for kiq fence" 2>/dev/null || echo 0')
print(f"\nKIQ errors after hip_probe: {out}")

# Test 11: hip_minimal_test
section("Test 11: hip_minimal_test (kernel compute)")
out, _, rc = run('export HSA_OVERRIDE_GFX_VERSION=10.1.0 HIP_VISIBLE_DEVICES=0 HSA_ENABLE_SDMA=0 HSA_TOOLS_LIB= HSA_TOOLS_REPORT_LOAD_FAILURE=0; /tmp/hip_minimal_test 2>&1', timeout=60)
print(out)
print(f"Exit code: {rc}")

time.sleep(3)
out, _, _ = run('sudo dmesg | tail -20 | grep -ci "timeout waiting for kiq fence" 2>/dev/null || echo 0')
print(f"\nKIQ errors after hip_minimal_test: {out}")

# Test 12: hip_vector_add (managed memory + timing)
section("Test 12: hip_vector_add (managed memory, sin²+cos²)")
out, _, rc = run('export HSA_OVERRIDE_GFX_VERSION=10.1.0 HIP_VISIBLE_DEVICES=0 HSA_ENABLE_SDMA=0 HSA_TOOLS_LIB= HSA_TOOLS_REPORT_LOAD_FAILURE=0; /tmp/hip_vector_add 2>&1', timeout=60)
print(out)
print(f"Exit code: {rc}")

time.sleep(3)
out, _, _ = run('sudo dmesg | tail -20 | grep -ci "timeout waiting for kiq fence" 2>/dev/null || echo 0')
print(f"\nKIQ errors after hip_vector_add: {out}")

# Test 13: Sequential stress (3 rounds)
section("Test 13: Sequential GPU Stress (3 rounds)")
for i in range(1, 4):
    out, _, rc = run('export HSA_OVERRIDE_GFX_VERSION=10.1.0 HIP_VISIBLE_DEVICES=0 HSA_ENABLE_SDMA=0 HSA_TOOLS_LIB= HSA_TOOLS_REPORT_LOAD_FAILURE=0; /tmp/hip_minimal_test 2>&1 | tail -3', timeout=60)
    status = "PASS" if rc == 0 else "FAIL"
    print(f"Round {i}: {status} (rc={rc}) — {out.split(chr(10))[-1]}")
    time.sleep(2)

# Final KIQ check
time.sleep(5)
section("Final: Post-Stress KIQ Check")
out, _, _ = run('sudo dmesg | grep -ci "timeout waiting for kiq fence" 2>/dev/null || echo 0')
print(f"Total KIQ timeout count: {out}")

out, _, _ = run('sudo dmesg | grep -ci "GPU unreachable\\|GPU died" 2>/dev/null || echo 0')
print(f"Total GPU dead count: {out}")

# Test 14: rocminfo AFTER all GPU tests
section("Test 14: rocminfo After Stress (previously would hang)")
out, _, rc = run('export HSA_OVERRIDE_GFX_VERSION=10.1.0 HIP_VISIBLE_DEVICES=0 HSA_ENABLE_SDMA=0 HSA_TOOLS_LIB= HSA_TOOLS_REPORT_LOAD_FAILURE=0; timeout 30 rocminfo 2>&1 | grep "Done"', timeout=60)
print(f"rocminfo: {out} (rc={rc})")

print("\n" + "="*50)
print(" ALL VERIFICATION COMPLETE")
print("="*50)
c.close()