#!/usr/bin/env python3
"""Install missing ROCm math libraries for PyTorch build 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) > 60:
            print(f"  ... ({len(lines)} lines, showing last 60)")
            print('\n'.join(lines[-60:]))
        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

# Find all available ROCm packages
run("pacman -Ss rocm | grep -E '^cachyos|^extra|^core' | head -40",
    desc="Available ROCm packages")

# Install all ROCm math/compute libraries needed by PyTorch
run("sudo pacman -S --needed --noconfirm "
    "hiprand rocrand "
    "hipblas rocblas "
    "hipfft rocfft "
    "hipsparse rocsparse "
    "hipsolver rocsolver "
    "miopen-hip "
    "rocprim hipcub "
    "rocthrust "
    "rccl "
    "hipblaslt "
    "roctracer "
    "2>&1 | tail -40",
    desc="Install ROCm math libraries",
    timeout=600)

# Verify hiprand is now available
run("find /opt/rocm -name 'hiprandConfig.cmake' -o -name 'hiprand-config.cmake' 2>/dev/null | head -5",
    desc="Verify hiprand cmake config")

# Check all libraries
run("ls /opt/rocm/lib/libhiprand.so /opt/rocm/lib/librocblas.so /opt/rocm/lib/libhipblas.so /opt/rocm/lib/librocfft.so /opt/rocm/lib/libMIOpen.so 2>&1",
    desc="Verify key libraries exist")

# Clean the failed build and restart
run("bash -c 'source ~/comfyui-env/bin/activate && cd ~/pytorch && python3 setup.py clean 2>&1 | tail -5'",
    desc="Clean failed build")

# Restart build
run("rm -f /home/fabian/pytorch_build.log", desc="Clean old log")
run("bash -c 'nohup bash /home/fabian/build_pytorch.sh </dev/null >/dev/null 2>&1 & echo PID=$!'",
    desc="Restart PyTorch build")

time.sleep(15)
run("pgrep -fa 'setup.py\\|cmake\\|ninja' | head -10",
    desc="Verify build restarted")

time.sleep(45)
run("tail -40 /home/fabian/pytorch_build.log 2>/dev/null",
    desc="Build progress")

ssh.close()
print("\nDone — ROCm libs installed and build restarted.")