#!/usr/bin/env python3
"""Upgrade RB3011 ENJOY from ROS 6.49 to ROS 7 (arm). Does not touch RB4011."""
from __future__ import annotations

import time
import urllib.request
from pathlib import Path

import paramiko

MK = ("192.168.88.249", "admin", "1234")
VER = "7.21.5"
NPK_NAME = f"routeros-arm-{VER}.npk"
NPK_URL = f"https://download.mikrotik.com/routeros/{VER}/{NPK_NAME}"
LOCAL = Path(r"c:\laragon\www\light-otantik-nux\scripts") / NPK_NAME


def wait_prompt(shell, timeout=25) -> str:
    buf = b""
    end = time.time() + timeout
    while time.time() < end:
        time.sleep(0.12)
        if shell.recv_ready():
            buf += shell.recv(65535)
            end = time.time() + 1.8
            if b"> " in buf[-250:] or b"y/N" in buf.lower()[-80:] or b"y/n" in buf.lower()[-80:]:
                break
    return buf.decode("utf-8", "replace")


def cmd(shell, line: str, timeout=25) -> str:
    shell.send(line + "\r")
    out = wait_prompt(shell, timeout)
    print("\n===== " + line + " =====")
    print(out[-2500:])
    return out


def ssh():
    mk = paramiko.SSHClient()
    mk.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    mk.connect(
        MK[0],
        username=MK[1],
        password=MK[2],
        timeout=20,
        allow_agent=False,
        look_for_keys=False,
    )
    return mk


def ping_up(timeout=180) -> bool:
    import socket

    end = time.time() + timeout
    while time.time() < end:
        s = socket.socket()
        s.settimeout(3)
        try:
            s.connect((MK[0], 22))
            s.close()
            return True
        except OSError:
            time.sleep(4)
        finally:
            try:
                s.close()
            except OSError:
                pass
    return False


def main() -> None:
    if not LOCAL.exists() or LOCAL.stat().st_size < 1_000_000:
        print("downloading", NPK_URL)
        urllib.request.urlretrieve(NPK_URL, LOCAL)
    print("local npk", LOCAL.stat().st_size)

    mk = ssh()
    shell = mk.invoke_shell(width=200, height=50)
    wait_prompt(shell, 12)
    cmd(shell, "/system identity print")
    cmd(shell, "/system resource print")
    cmd(shell, "/system backup save name=before-ros7")
    cmd(shell, "/export file=before-ros7")
    cmd(shell, "/file print")
    mk.close()

    print("uploading npk via sftp...")
    mk = ssh()
    sftp = mk.open_sftp()
    sftp.put(str(LOCAL), NPK_NAME)
    sftp.close()
    mk.close()
    print("uploaded")

    mk = ssh()
    shell = mk.invoke_shell(width=200, height=50)
    wait_prompt(shell, 12)
    cmd(shell, "/file print where name~\"routeros\"")
    cmd(shell, "/file print where name~\"before-ros7\"")
    print("rebooting to install ROS", VER)
    shell.send("/system reboot\r")
    time.sleep(1)
    out = wait_prompt(shell, 8)
    print(out[-800:])
    if "y/N" in out or "y/n" in out.lower() or "Reboot" in out:
        shell.send("y\r")
        time.sleep(2)
    try:
        mk.close()
    except Exception:
        pass

    print("waiting for SSH after reboot...")
    time.sleep(25)
    if not ping_up(240):
        raise SystemExit("RB3011 did not come back on SSH")
    print("SSH is up, waiting extra 20s for boot...")
    time.sleep(20)

    mk = ssh()
    shell = mk.invoke_shell(width=200, height=50)
    wait_prompt(shell, 20)
    cmd(shell, "/system identity print")
    cmd(shell, "/system resource print")
    cmd(shell, "/interface wireguard print")
    mk.close()
    print("UPGRADE_DONE")


if __name__ == "__main__":
    main()