#!/usr/bin/env python3
"""Try pre-built PyTorch ROCm from CachyOS repos, test on BC-250."""
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=300, desc=""):
    if desc:
        print(f"\n{'='*60}")
        print(f"  {desc}")
        print(f"{'='*60}")
    print(f"$ {cmd}")
    _, 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) > 80:
            print(f"  ... ({len(lines)} lines, showing last 80)")
            print('\n'.join(lines[-80:]))
        else:
            print(out.strip())
    if err.strip():
        lines = err.strip().split('\n')
        show = lines[-20:] if len(lines) > 20 else lines
        print(f"STDERR: {chr(10).join(show)}")
    print(f"  Exit code: {rc}")
    return rc, out, err

# Check what architectures the pre-built packages support
run("pacman -Si python-pytorch-rocm 2>&1 | head -20",
    desc="Pre-built PyTorch ROCm package info")

run("pacman -Si python-pytorch-opt-rocm 2>&1 | head -20",
    desc="Pre-built PyTorch Opt ROCm package info")

# Install the pre-built package (system-wide, venv will pick it up via --system-site-packages)
run("sudo pacman -S --needed --noconfirm python-pytorch-rocm 2>&1 | tail -30",
    desc="Install pre-built PyTorch ROCm",
    timeout=600)

# Test if it works in venv
run("""bash -c 'source ~/comfyui-env/bin/activate && \
    HSA_OVERRIDE_GFX_VERSION=10.1.0 \
    HIP_VISIBLE_DEVICES=0 \
    HSA_ENABLE_SDMA=0 \
    python3 -c "
import torch
print(f\\"PyTorch version: {torch.__version__}\\")
print(f\\"HIP version: {torch.version.hip}\\")
print(f\\"CUDA available (HIP): {torch.cuda.is_available()}\\")
if torch.cuda.is_available():
    print(f\\"Device count: {torch.cuda.device_count()}\\")
    print(f\\"Device name: {torch.cuda.get_device_name(0)}\\")
    print(f\\"Device arch: {torch.cuda.get_device_capability(0)}\\")
    # Try a simple tensor operation on GPU
    t = torch.randn(4, 4, device=\\"cuda\\")
    print(f\\"Tensor device: {t.device}\\")
    print(f\\"Tensor sum: {t.sum().item():.4f}\\")
    # Try matmul
    a = torch.randn(64, 64, device=\\"cuda\\")
    b = torch.randn(64, 64, device=\\"cuda\\")
    c = torch.matmul(a, b)
    print(f\\"Matmul result shape: {c.shape}\\")
    print(\\"GPU COMPUTE: WORKING\\")
else:
    print(\\"CUDA/HIP NOT AVAILABLE\\")
" 2>&1'""",
    desc="Test PyTorch on GPU",
    timeout=120)

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