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 = os.environ.get("VPS_IP", "85.155.190.233")
PORT = int(os.environ.get("VPS_PORT", "50022"))
VPS_IP = IP
VPS_PORT = PORT
USER = "root"
KEY_PATH = os.path.expanduser(os.environ.get("VPS_KEY_PATH", "~/.ssh/id_ed25519_new"))

def connect():
    safe_print(f"Connecting to {IP}:{PORT}...")
    client = paramiko.SSHClient()
    client.load_system_host_keys()
    client.set_missing_host_key_policy(paramiko.RejectPolicy())
    
    try:
        # Try connecting using the local ssh-agent (unlocked key is retrieved from memory)
        client.connect(IP, port=PORT, username=USER, key_filename=KEY_PATH, allow_agent=True, timeout=10)
        safe_print("Connected successfully (via SSH Agent)!")
        return client
    except Exception as agent_err:
        safe_print(f"SSH Agent connection failed or key not loaded. Trying direct file load...")
        try:
            key = paramiko.Ed25519Key.from_private_key_file(KEY_PATH)
            client.connect(IP, port=PORT, username=USER, pkey=key, timeout=10)
            safe_print("Connected successfully!")
            return client
        except paramiko.ssh_exception.PasswordRequiredException as pass_err:
            safe_print("Error: The SSH private key is encrypted with a passphrase.")
            safe_print("Please start your ssh-agent and load the key: ssh-add <path-to-key>")
            raise pass_err
        except Exception as e:
            safe_print(f"Connection failed: {e}")
            safe_print("Hint: Ensure the remote host key is added to ~/.ssh/known_hosts using 'ssh-keyscan -p <port> <ip> >> ~/.ssh/known_hosts'")
            raise e

def run_cmd(client, cmd):
    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')
    return exit_status, out, err