#!/usr/bin/env python3
"""SSH/SFTP helper for the Ubuntu production server."""
from __future__ import annotations

import argparse
import os
import sys
import time

import paramiko

HOST = "10.15.15.89"
USER = "ubuntu"
PASSWORD = "ubuntu"
PORT = 22


def connect() -> paramiko.SSHClient:
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    client.connect(
        HOST,
        port=PORT,
        username=USER,
        password=PASSWORD,
        timeout=20,
        banner_timeout=30,
        allow_agent=False,
        look_for_keys=False,
    )
    return client


def run(cmd: str, timeout: int = 300, sudo: bool = False) -> tuple[int, str]:
    if sudo:
        # -S reads password from stdin; -p '' hides prompt
        wrapped = f"sudo -S -p '' bash -lc {repr(cmd)}"
    else:
        wrapped = f"bash -lc {repr(cmd)}"
    client = connect()
    try:
        stdin, stdout, stderr = client.exec_command(wrapped, timeout=timeout, get_pty=True)
        if sudo:
            stdin.write(PASSWORD + "\n")
            stdin.flush()
        out = stdout.read().decode("utf-8", errors="replace")
        err = stderr.read().decode("utf-8", errors="replace")
        code = stdout.channel.recv_exit_status()
        text = out
        if err.strip() and err.strip() not in text:
            text += ("\n" if text else "") + err
        # strip sudo password leak if any
        lines = []
        for line in text.splitlines():
            if line.strip() == PASSWORD:
                continue
            lines.append(line)
        return code, "\n".join(lines) + ("\n" if lines else "")
    finally:
        client.close()


def put(local: str, remote: str) -> None:
    client = connect()
    try:
        sftp = client.open_sftp()
        remote_dir = os.path.posixpath.dirname(remote)
        # mkdir -p via ssh is simpler
        client.exec_command(f"mkdir -p {remote_dir}")
        time.sleep(0.2)
        sftp.put(local, remote)
        sftp.close()
    finally:
        client.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("cmd", nargs="?", default="uname -a")
    parser.add_argument("--timeout", type=int, default=300)
    parser.add_argument("--sudo", action="store_true")
    parser.add_argument("--put", nargs=2, metavar=("LOCAL", "REMOTE"))
    args = parser.parse_args()
    if args.put:
        put(args.put[0], args.put[1])
        print(f"uploaded {args.put[0]} -> {args.put[1]}")
        sys.exit(0)
    code, out = run(args.cmd, timeout=args.timeout, sudo=args.sudo)
    sys.stdout.buffer.write(out.encode("utf-8", errors="replace"))
    sys.exit(code)