#!/usr/bin/env python3
import json
import os
import os.path
import subprocess
import sys

progname = sys.argv[0]
os.chdir(os.path.dirname(os.path.realpath(__file__)))


def log(msg):
    sys.stderr.write(f"{progname}: {msg.strip() + '\n'}")


def tofu_output():
    completed = subprocess.run(
        [
            "tofu",
            "output",
            "-json",
        ],
        stdout=subprocess.PIPE,
        stderr=sys.stderr,
        text=True,
        check=True,
        cwd="../infra",
    )
    output = json.loads(completed.stdout)
    n = len(output.keys())
    log(f"{n} output{'s' if n > 1 else ''} from tofu")
    return {key: output[key]["value"] for key in output.keys()}


def nixos_cfg_from_fun(nixosfun, special_args):
    completed = subprocess.run(
        [
            "nix",
            "eval",
            "--raw",
            nixosfun,
            "--apply",
            f"f: (f (builtins.fromJSON ''{json.dumps(special_args)}'')).config.system.build.toplevel.drvPath",
        ],
        check=True,
        stdout=subprocess.PIPE,
        stderr=sys.stderr,
        text=True,
    )
    log(f"retrieved derivation {completed.stdout}")
    return completed.stdout + "^out"


def nix_build(installable):
    completed = subprocess.run(
        [
            "nix",
            "build",
            "--log-format",
            "raw",
            "--no-link",
            "--print-out-paths",
            installable,
        ],
        check=True,
        stdout=subprocess.PIPE,
        stderr=sys.stderr,
        text=True,
    )
    build_path = completed.stdout.strip()
    log(f"built {build_path}")
    return build_path


def nix_copy(installable, target):
    completed = subprocess.run(
        [
            "nix",
            "copy",
            "--verbose",
            "--substitute-on-destination",
            "--to",
            f"ssh://{target}",
            installable,
        ],
        stdin=sys.stdin,
        stdout=sys.stdout,
        stderr=sys.stderr,
        check=True,
    )
    log(f"copied build to {target}")


def nixos_switch(target, buildPath):
    completed = subprocess.run(
        [
            "env",
            "LANG=",
            "LC_ALL=C.UTF-8",
            "ssh",
            target,
            f"""
        sudo nix-env -p /nix/var/nix/profiles/system --set {buildPath}
        sudo /nix/var/nix/profiles/system/bin/switch-to-configuration switch
        """,
        ],
        stdin=sys.stdin,
        stdout=sys.stdout,
        stderr=sys.stderr,
        check=True,
    )


def build(nixosfunattr, special_args=None):
    drv_path = nixos_cfg_from_fun(
        nixosfunattr, special_args if special_args else tofu_output()
    )
    return nix_build(drv_path)


def apply_(build_path, target):
    nix_copy(build_path, target)
    nixos_switch(target, build_path)


def usage():
    sys.stderr.write(f"usage: {progname} [build|apply]\n")


NIXOSFUNATTR = "..#nixosFunctions.arm"
TARGET = "root@ampere"


def main():
    subcommand: str
    try:
        subcommand = sys.argv[1]
    except IndexError:
        subcommand = "apply"
    try:
        match subcommand:
            case "build":
                build(NIXOSFUNATTR)
            case "apply":
                path = build(NIXOSFUNATTR)
                apply_(path, TARGET)
            case _:
                usage()
                sys.exit(127)
    except subprocess.CalledProcessError:
        sys.exit(1)


if __name__ == "__main__":
    main()