import os
import sys
import subprocess
import json

SESSION_FILE = r"c:\Projects\Historical\runpod_session.json"
SSH_KEY = os.path.expanduser(r"~/.ssh/id_ed25519")

with open(SESSION_FILE, "r") as f:
    session = json.load(f)

HOST = session["host"]
PORT = str(session["port"])

def run(cmd):
    ssh_cmd = [
        "ssh", "-o", "StrictHostKeyChecking=no", "-o", "UserKnownHostsFile=/dev/null",
        "-i", SSH_KEY, "-p", PORT, f"root@{HOST}", cmd
    ]
    p = subprocess.run(ssh_cmd, capture_output=True, text=True, encoding="utf-8", errors="replace")
    print("STDOUT:", p.stdout)
    if p.stderr:
        print("STDERR:", p.stderr)
    return p.returncode

if __name__ == "__main__":
    cmd = " ".join(sys.argv[1:]) if len(sys.argv) > 1 else "nvidia-smi"
    run(cmd)