#!/usr/bin/env python3
import argparse
import json
import pathlib
import shutil
import socket
import time

from subprocess import check_output

ROOT_DIR = pathlib.Path(__file__).parent.parent
KEYS_FILE = ROOT_DIR / "secrets" / "keys.json"


def _get_available_machines() -> list:
    output = check_output(["nix", "flake", "show", "--json"])
    parsed_output = json.loads(output)
    machines = parsed_output.get("nixosConfigurations", dict()).keys()
    return list(machines)


def _is_valid_ip(ip: str) -> bool:
    try:
        socket.inet_aton(ip)
        return True
    except socket.error:
        return False


def _check_ssh_connection(ip: str) -> bool:
    try:
        check_output(["ssh", f"krop@{ip}", "echo", "Connected"])
        return True
    except Exception:
        return False


def add_key_to_secrets(machine_name: str, key: str):
    keys = json.loads(KEYS_FILE.read_text())
    if keys.get("servers").get(machine_name):
        raise ValueError(f"Key for {machine_name} already exists, remove it first")

    keys["servers"][machine_name] = key

    for secret in keys.get("secrets"):
        keys["secrets"][secret].append(f"servers:{machine_name}")

    KEYS_FILE.write_text(json.dumps(keys, indent=2))


def rekey_secrets():
    agenix_bin = shutil.which("agenix")
    check_output([agenix_bin, "-r"], cwd=ROOT_DIR / "secrets")


def bootstrap_machine(ip: str):
    check_output(
        [
            "nix",
            "run",
            "github:nix-community/nixos-anywhere",
            "--",
            "--flake",
            ".#bootstrap",
            "--target-host",
            f"root@{ip}",
            "--build-on-remote",
        ]
    )


def install_machine(machine_name: str, ip: str):
    check_output(
        [
            "nixos-rebuild",
            "boot",
            "--flake",
            f".#{machine_name}",
            "--fast",
            "--target-host",
            f"krop@{ip}",
            "--build-host",
            f"krop@{ip}",
            "--use-remote-sudo",
        ]
    )


def get_ssh_key(ip: str) -> str:
    """
    This function uses machines ssh-keyscan to get the ssh key and then get the ed25519 key
    """
    ssh_keys = (
        check_output(
            [
                "ssh-keyscan",
                "-q",
                "-t",
                "ed25519",
                ip,
            ]
        )
        .decode("utf-8")
        .strip()
        .splitlines()
    )

    if len(ssh_keys) != 1:
        raise ValueError("Exactly one key should be returned")

    key = ssh_keys.pop().lstrip(f"{ip} ").strip()

    return key


def get_machine_config(machine_name: str) -> dict:
    output = check_output(
        [
            "nix",
            "eval",
            "--json",
            f".#nixosConfigurations.{machine_name}.config.kropcloud",
        ]
    )
    return json.loads(output)


def reboot_machine(ip: str):
    check_output(
        [
            "ssh",
            f"krop@{ip}",
            "sudo",
            "reboot",
        ]
    )


def main() -> int:
    parser = argparse.ArgumentParser(description="Install a machine")
    parser.add_argument(
        "machine_name", type=str, help="The name of the machine to install"
    )
    parser.add_argument("machine_ip", type=str, help="The ip of the machine to install")
    args = parser.parse_args()

    machine_name = args.machine_name
    if machine_name not in _get_available_machines():
        raise ValueError(
            f"Machine {machine_name} not found, available machines are: {_get_available_machines()}"
        )

    machine_ip = args.machine_ip
    if not _is_valid_ip(machine_ip):
        raise ValueError(f"Invalid IP address {machine_ip}")

    print(f"Bootstrapping machine {machine_ip}")
    bootstrap_machine(machine_ip)
    print("Machine bootstrapped")

    print("Waiting for ssh connection")
    while not _check_ssh_connection(machine_ip):
        time.sleep(5)
    print("Machine is up and running")

    print("Getting ssh key")
    ssh_key = get_ssh_key(machine_ip)

    print(f"SSH key: {ssh_key}")
    print("Adding ssh key to secrets")
    add_key_to_secrets(machine_name, ssh_key)
    rekey_secrets()

    print("Installing machine")
    install_machine(machine_name, machine_ip)
    print("Machine installed, rebooting")
    reboot_machine(machine_ip)
    print("")
    return 0


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