#!/usr/bin/env python3
"""Step 1: Install build dependencies on BC-250 for PyTorch build."""
import paramiko
import sys
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():
        print(out.strip()[-2000:])  # Last 2000 chars
    if err.strip():
        # Filter out common noise
        lines = [l for l in err.strip().split('\n') if not l.startswith('warning:')]
        if lines:
            print(f"STDERR: {chr(10).join(lines[-20:])}")
    print(f"  Exit code: {rc}")
    return rc, out, err

# Install build tools
run("sudo pacman -S --needed --noconfirm python-pip ninja ccache",
    desc="Install pip, ninja, ccache")

# Install PyTorch build dependencies
run("sudo pacman -S --needed --noconfirm cmake blas lapack openblas "
    "python-numpy python-pyyaml python-typing_extensions "
    "intel-oneapi-mkl 2>/dev/null; echo done",
    desc="Install build dependencies (cmake, blas, numpy, etc.)")

# Install additional deps that PyTorch needs
run("sudo pacman -S --needed --noconfirm python-cffi python-setuptools "
    "python-wheel python-filelock python-sympy python-networkx",
    desc="Install Python dependencies")

# Verify installs
run("pip --version && ninja --version && ccache --version | head -1 && cmake --version | head -1",
    desc="Verify installations")

# Check pip can install packages  
run("pip install --user --upgrade pip setuptools wheel 2>&1 | tail -5",
    desc="Upgrade pip/setuptools")

ssh.close()
print("\n\nDone — build dependencies installed.")