#!/usr/bin/env python3
"""Fix torchvision — use system package instead of pip 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

# 1. Uninstall pip torchvision and torchaudio from venv
run("bash -c 'source /home/fabian/comfyui-env/bin/activate && pip uninstall -y torchvision torchaudio 2>&1'",
    desc="Uninstall pip torchvision and torchaudio from venv")

# 2. Install system python-torchvision via pacman
run("sudo pacman -S --noconfirm python-torchvision",
    desc="Install system python-torchvision (matches system pytorch)")

# 3. Verify torchvision now works
run("bash -c 'source /home/fabian/comfyui-env/bin/activate && "
    "export HSA_OVERRIDE_GFX_VERSION=10.1.0 && "
    "export HIP_VISIBLE_DEVICES=0 && "
    "export HSA_ENABLE_SDMA=0 && "
    "python -c \""
    "import torch; print(f\\\"torch {torch.__version__} from {torch.__file__}\\\"); "
    "import torchvision; print(f\\\"torchvision {torchvision.__version__} from {torchvision.__file__}\\\"); "
    "print(\\\"torchvision ops OK\\\"); "
    "\"'",
    desc="Verify torchvision import works")

# 4. Kill old ComfyUI and relaunch
run("bash -c 'pkill -f \"python main.py\" 2>/dev/null; sleep 2; echo done'",
    desc="Kill old ComfyUI process")

run("bash -c 'nohup bash /home/fabian/start_comfyui.sh > /home/fabian/comfyui.log 2>&1 & echo PID=$!'",
    desc="Relaunch ComfyUI")

time.sleep(20)
run("tail -40 /home/fabian/comfyui.log 2>/dev/null",
    desc="ComfyUI startup log")

time.sleep(10)
run("bash -c 'ss -tlnp | grep 8188 || echo PORT_NOT_LISTENING'",
    desc="Check if port 8188 is listening")

run("tail -60 /home/fabian/comfyui.log 2>/dev/null",
    desc="Full ComfyUI log")

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