import paramiko
import os
import sys

# Reconfigure stdout/stderr to utf-8 if possible
if hasattr(sys.stdout, 'reconfigure'):
    try:
        sys.stdout.reconfigure(encoding='utf-8')
    except Exception:
        pass
if hasattr(sys.stderr, 'reconfigure'):
    try:
        sys.stderr.reconfigure(encoding='utf-8')
    except Exception:
        pass

def safe_print(msg):
    try:
        print(msg)
    except UnicodeEncodeError:
        enc = sys.stdout.encoding or 'utf-8'
        print(msg.encode(enc, errors='replace').decode(enc))

IP = "85.155.190.233"
PORT = 50022
USER = "root"
KEY_PATH = os.path.expanduser("~/.ssh/id_ed25519")

def run(cmd):
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    
    try:
        key = paramiko.Ed25519Key.from_private_key_file(KEY_PATH)
        client.connect(IP, port=PORT, username=USER, pkey=key, timeout=10)
    except Exception as e:
        safe_print(f"Connection failed: {e}")
        return

    safe_print(f"Executing: {cmd}")
    stdin, stdout, stderr = client.exec_command(cmd)
    exit_status = stdout.channel.recv_exit_status()
    out = stdout.read().decode('utf-8', errors='replace')
    err = stderr.read().decode('utf-8', errors='replace')
    
    if exit_status != 0:
        safe_print(f"Command failed with exit code: {exit_status}")
    else:
        safe_print("Command finished successfully.")
        
    if out.strip():
        safe_print(f"STDOUT:\n{out}")
    if err.strip():
        safe_print(f"STDERR:\n{err}")
        
    client.close()

if __name__ == "__main__":
    if len(sys.argv) > 1:
        run(" ".join(sys.argv[1:]))
    else:
        run("/root/Gemini_Swarm_Core/venv/bin/python /root/Gemini_Swarm_Core/scripts/test_gemini.py")