#!/usr/bin/env python3
"""Simple SSH helper for deployment (paramiko)."""
from __future__ import annotations

import argparse
import base64
import sys
from pathlib import Path

import paramiko


def connect(host: str, user: str, password: str) -> paramiko.SSHClient:
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    client.connect(
        host,
        username=user,
        password=password,
        timeout=30,
        allow_agent=False,
        look_for_keys=False,
    )
    return client


def _exec(client: paramiko.SSHClient, command: str, timeout: int = 120) -> tuple[int, str, str]:
    stdin, stdout, stderr = client.exec_command(command, timeout=timeout, get_pty=True)
    out = stdout.read().decode(errors="replace")
    err = stderr.read().decode(errors="replace")
    code = stdout.channel.recv_exit_status()
    return code, out, err


def run(
    host: str,
    user: str,
    password: str,
    command: str,
    *,
    sudo: bool = False,
    timeout: int = 120,
) -> int:
    client = connect(host, user, password)
    try:
        if sudo:
            b64 = base64.b64encode(command.encode()).decode()
            full = (
                f"echo {password!r} | sudo -S -p '' bash -c "
                f"'echo {b64} | base64 -d | bash'"
            )
        else:
            full = f"bash -lc {command!r}"
        code, out, err = _exec(client, full, timeout=timeout)
        sys.stdout.write(out)
        if err.strip():
            sys.stderr.write(err)
        return code
    finally:
        client.close()


def upload(
    host: str,
    user: str,
    password: str,
    local: str,
    remote: str,
    *,
    sudo: bool = False,
) -> int:
    data = Path(local).read_bytes()
    b64 = base64.b64encode(data).decode()
    remote_q = remote.replace("'", "'\\''")
    write_cmd = f"mkdir -p \"$(dirname '{remote_q}')\" && echo {b64} | base64 -d > '{remote_q}'"
    return run(host, user, password, write_cmd, sudo=sudo, timeout=300)


def main() -> int:
    parser = argparse.ArgumentParser()
    sub = parser.add_subparsers(dest="action", required=True)

    p_run = sub.add_parser("run")
    p_run.add_argument("--host", required=True)
    p_run.add_argument("--user", required=True)
    p_run.add_argument("--password", required=True)
    p_run.add_argument("--sudo", action="store_true")
    p_run.add_argument("--timeout", type=int, default=120)
    p_run.add_argument("command")

    p_up = sub.add_parser("upload")
    p_up.add_argument("--host", required=True)
    p_up.add_argument("--user", required=True)
    p_up.add_argument("--password", required=True)
    p_up.add_argument("--local", required=True)
    p_up.add_argument("--remote", required=True)
    p_up.add_argument("--sudo", action="store_true")

    args = parser.parse_args()
    if args.action == "run":
        return run(
            args.host,
            args.user,
            args.password,
            args.command,
            sudo=args.sudo,
            timeout=args.timeout,
        )
    return upload(
        args.host,
        args.user,
        args.password,
        args.local,
        args.remote,
        sudo=args.sudo,
    )


if __name__ == "__main__":
    raise SystemExit(main())