#!/usr/bin/env python3
"""Fix torchvision compatibility on BC-250.

The pip-installed torchvision conflicts with the system python-pytorch-rocm.
Need to use system torchvision-rocm or fix the version.
"""
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) > 40:
            print(f"  ... ({len(lines)} lines, showing last 40)")
            print('\n'.join(lines[-40:]))
        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 torchvision packages exist in repos
run("bash -c 'pacman -Ss torchvision 2>/dev/null'",
    desc="Search for torchvision packages in repos")

run("bash -c 'pacman -Ss torchaudio 2>/dev/null'",
    desc="Search for torchaudio packages")

# Check what's currently installed
run("bash -c 'pacman -Qs torch 2>/dev/null'",
    desc="Currently installed torch packages (system)")

run("bash -c 'source /home/fabian/comfyui-env/bin/activate && pip list 2>/dev/null | grep -i torch'",
    desc="torch packages in venv")

# Check the versions
run("bash -c 'source /home/fabian/comfyui-env/bin/activate && python -c \""
    "import torch; print(f\\\"torch: {torch.__version__} from {torch.__file__}\\\"); "
    "\"'",
    desc="Check torch location")

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