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")
LOCAL_ADC_PATH = os.path.expandvars(r"%APPDATA%\gcloud\application_default_credentials.json")
REMOTE_KEY_PATH = "/root/gcp-key.json"
REMOTE_ENV_PATH = "/root/Gemini_Swarm_Core/.env"
PROJECT_ID = "gen-lang-client-0207686288"

def deploy_gcp():
    if not os.path.exists(LOCAL_ADC_PATH):
        safe_print(f"Error: Local Google credentials not found at {LOCAL_ADC_PATH}")
        return
        
    safe_print(f"Connecting to {IP}:{PORT} using key {KEY_PATH}...")
    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)
        safe_print("Connected successfully!")
    except Exception as e:
        safe_print(f"Connection failed: {e}")
        return

    sftp = client.open_sftp()
    
    # 1. Upload the credentials file
    safe_print(f"Uploading local credentials to {REMOTE_KEY_PATH}...")
    sftp.put(LOCAL_ADC_PATH, REMOTE_KEY_PATH)
    sftp.chmod(REMOTE_KEY_PATH, 0o600)
    safe_print("Credentials file uploaded and secured.")

    # 2. Update .env file on VPS
    safe_print(f"Updating {REMOTE_ENV_PATH} on remote host...")
    env_content = ""
    try:
        with sftp.file(REMOTE_ENV_PATH, "r") as f:
            env_content = f.read().decode('utf-8')
    except IOError:
        safe_print("Warning: Remote .env file not found. Creating a new one...")
        
    lines = env_content.split("\n")
    new_lines = []
    has_gcredentials = False
    has_gproject = False
    
    for line in lines:
        stripped = line.strip()
        if stripped.startswith("GOOGLE_APPLICATION_CREDENTIALS="):
            new_lines.append(f'GOOGLE_APPLICATION_CREDENTIALS={REMOTE_KEY_PATH}')
            has_gcredentials = True
        elif stripped.startswith("GCP_PROJECT_ID="):
            new_lines.append(f'GCP_PROJECT_ID={PROJECT_ID}')
            has_gproject = True
        else:
            new_lines.append(line)
            
    if not has_gcredentials:
        new_lines.append(f'GOOGLE_APPLICATION_CREDENTIALS={REMOTE_KEY_PATH}')
    if not has_gproject:
        new_lines.append(f'GCP_PROJECT_ID={PROJECT_ID}')
        
    updated_env = "\n".join(new_lines)
    with sftp.file(REMOTE_ENV_PATH, "w") as f:
        f.write(updated_env.encode('utf-8'))
    safe_print("Remote .env file updated successfully.")
    
    sftp.close()

    # 3. Restart tg-bridge systemd service to reload env
    safe_print("Restarting tg-bridge service to apply new environment variables...")
    stdin, stdout, stderr = client.exec_command("systemctl restart tg-bridge.service && systemctl status tg-bridge.service")
    exit_status = stdout.channel.recv_exit_status()
    out = stdout.read().decode('utf-8')
    err = stderr.read().decode('utf-8')
    if exit_status == 0:
        safe_print("Service restarted successfully!")
        safe_print(f"Status:\n{out}")
    else:
        safe_print("Failed to restart service.")
        safe_print(f"Error:\n{err}")

    client.close()
    safe_print("Google Cloud credentials deployment completed!")

if __name__ == "__main__":
    deploy_gcp()