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

from subprocess import check_output


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 _validate_ip(ip: str) -> bool:
    try:
        socket.inet_aton(ip)
        return False
    except socket.error:
        return True


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

def add_key_to_secrets(key: str):
    

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 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 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()}"
        )

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

    machine_config = get_machine_config(machine_name)
    # We are bootstraping the machine first because we need their ssh keys
    print(f"Bootstrapping machine {args.machine_ip}")
    bootstrap_machine()
    print("Machine bootstrapped")

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

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

    print("Adding ssh key to secrets")
    add_key_to_secrets(ssh_key)
    rekey_secrets()

    # Add the ssh key to keys in secrets/keys.json
    # and rekey the secrets

    # install_machine()

    return 0


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