# 2026-03-29 — handler.py: установка apt-пакетов на ВМ по SSH.
# sless_job runtime: python3.11, entrypoint: handler.install
#
# event_json:
#   packages: ["git", "curl", ...]  — список пакетов (обязательно)
#   update:   true/false            — apt-get update перед install (default: true)
#
# env_vars:
#   VM_IP:    внешний IP ВМ
#   SSH_USER: логин (ubuntu)
#   SSH_KEY:  содержимое приватного SSH-ключа (PEM)

import os, io, time
import paramiko


def _load_key(content):
    """Загрузить SSH-ключ (Ed25519 / RSA / ECDSA)."""
    for cls in (paramiko.Ed25519Key, paramiko.RSAKey, paramiko.ECDSAKey):
        try:
            return cls.from_private_key(io.StringIO(content))
        except Exception:
            pass
    raise ValueError("Неподдерживаемый тип SSH-ключа")


def _ssh_connect(retries=5, delay=10):
    """Подключение к ВМ с retry — ВМ может ещё загружаться."""
    key = _load_key(os.environ["SSH_KEY"])
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    last_err = None
    for attempt in range(retries):
        try:
            client.connect(
                hostname=os.environ["VM_IP"],
                username=os.environ["SSH_USER"],
                pkey=key,
                timeout=15,
            )
            return client
        except Exception as e:
            last_err = e
            if attempt < retries - 1:
                time.sleep(delay)
    raise RuntimeError(f"SSH не удалось после {retries} попыток: {last_err}")


def _run(client, cmd, timeout=120, check=True):
    """Выполнить команду, вернуть (exit_code, stdout, stderr)."""
    _, stdout, stderr = client.exec_command(cmd, timeout=timeout)
    code = stdout.channel.recv_exit_status()
    out = stdout.read().decode(errors="replace").strip()
    err = stderr.read().decode(errors="replace").strip()
    if check and code != 0:
        raise RuntimeError(f"Ошибка (exit {code}):\n{cmd}\nstderr: {err}")
    return code, out, err


def _wait_apt_lock(client, attempts=20, delay=10):
    """Ждать завершения cloud-init и убить авто-обновления. Ubuntu 22.04+."""
    # Шаг 1: Ждём завершения cloud-init — он держит apt при первом старте VM
    _run(client, "timeout 300 sudo cloud-init status --wait 2>/dev/null; true", check=False, timeout=310)
    # Шаг 2: Mask (не просто disable) — systemd не сможет перезапустить
    _run(client, "sudo systemctl mask unattended-upgrades apt-daily.service apt-daily-upgrade.service apt-daily.timer apt-daily-upgrade.timer 2>/dev/null; true", check=False)
    _run(client, "sudo systemctl stop unattended-upgrades apt-daily.service apt-daily-upgrade.service 2>/dev/null; true", check=False)
    # Шаг 3: Добить оставшиеся apt/dpkg процессы
    _run(client, "sudo pkill -9 -x unattended-upgrades apt-get apt dpkg 2>/dev/null; true", check=False)
    _run(client, "sudo kill -9 $(sudo lsof -t /var/lib/dpkg/lock-frontend 2>/dev/null) 2>/dev/null; true", check=False)
    # Шаг 4: Убрать стейл-локи и починить dpkg
    _run(client, "sudo rm -f /var/lib/dpkg/lock-frontend /var/lib/dpkg/lock /var/cache/apt/archives/lock /var/lib/apt/lists/lock 2>/dev/null; true", check=False)
    _run(client, "sudo dpkg --configure -a 2>/dev/null; true", check=False)
    time.sleep(3)

    locks = [
        "/var/lib/dpkg/lock-frontend",
        "/var/lib/dpkg/lock",
        "/var/lib/apt/lists/lock",
    ]
    for i in range(attempts):
        all_free = all(
            _run(client, f"sudo flock -n {lock} true 2>/dev/null", check=False)[0] == 0
            for lock in locks
        )
        if all_free:
            return
        # Повторить убийство процессов удерживающих lock
        _run(client, "sudo pkill -9 -x apt-get apt dpkg 2>/dev/null; true", check=False)
        _run(client, "sudo kill -9 $(sudo lsof -t /var/lib/dpkg/lock-frontend 2>/dev/null) 2>/dev/null; true", check=False)
        if i < attempts - 1:
            time.sleep(delay)
    raise RuntimeError("apt lock занят слишком долго — проверьте процессы на ВМ")


def install(event):
    """Установить apt-пакеты. Идемпотентно — повторный запуск безопасен."""
    packages = event.get("packages", [])
    if not packages:
        return {"status": "skipped", "reason": "packages list is empty"}

    do_update = event.get("update", True)

    client = _ssh_connect()
    try:
        _wait_apt_lock(client)

        if do_update:
            _run(client, "sudo DEBIAN_FRONTEND=noninteractive apt-get -o DPkg::Lock::Timeout=600 update -qq", timeout=420)

        pkg_str = " ".join(packages)
        _run(
            client,
            f"sudo DEBIAN_FRONTEND=noninteractive apt-get -o DPkg::Lock::Timeout=600 install -y -qq {pkg_str}",
            timeout=300,
        )

        # Проверить что установилось
        installed, missing = [], []
        for pkg in packages:
            code, _, _ = _run(
                client,
                f"dpkg -l {pkg} 2>/dev/null | grep -q '^ii'",
                check=False,
            )
            (installed if code == 0 else missing).append(pkg)

        return {
            "status": "ok" if not missing else "partial",
            "installed": installed,
            "missing": missing,
        }
    finally:
        client.close()