"""Install NixOS on a ZFS pool."""

from __future__ import annotations

import curses
import logging
import sys
from argparse import ArgumentParser
from os import getenv
from pathlib import Path
from random import getrandbits
from subprocess import run
from time import sleep
from typing import TYPE_CHECKING

from python.installer.tui import draw_menu
from python.logging_config import configure_logger
from python.process import require_commands, run_output

if TYPE_CHECKING:
    from collections.abc import Sequence

logger = logging.getLogger(__name__)

REQUIRED_COMMANDS = (
    "blkdiscard",
    "cryptsetup",
    "find",
    "lsblk",
    "mkfs.vfat",
    "mount",
    "nixos-generate-config",
    "nixos-install",
    "parted",
    "readlink",
    "zfs",
    "zpool",
)


def partition_disk(disk: str, swap_size: int, reserve: int = 0) -> None:
    """Partition a disk.

    Args:
        disk (str): The disk to partition.
        swap_size (int): The size of the swap partition in GB.
            minimum value is 1.
        reserve (int, optional): The size of the reserve partition in GB. Defaults to 0.
            minimum value is 0.
    """
    logger.info(f"partitioning {disk=}")
    swap_size = max(swap_size, 1)
    reserve = max(reserve, 0)

    run_output(("blkdiscard", "-f", disk))

    if reserve > 0:
        msg = f"Creating swap partition on {disk=} with size {swap_size=}GiB and reserve {reserve=}GiB"
        logger.info(msg)

        swap_start = swap_size + reserve
        swap_partition = f"mkpart swap -{swap_start}GiB -{reserve}GiB "
    else:
        logger.info(f"Creating swap partition on {disk=} with size {swap_size=}GiB")
        swap_start = swap_size
        swap_partition = f"mkpart swap -{swap_start}GiB 100% "

    logger.debug(f"{swap_partition=}")

    create_partitions = (
        "parted",
        "--script",
        "--align=optimal",
        disk,
        "--",
        "mklabel",
        "gpt",
        "mkpart",
        "EFI",
        "1MiB",
        "4GiB",
        "mkpart",
        "root_pool",
        "4GiB",
        f"-{swap_start}GiB",
        *swap_partition.split(),
        "set",
        "1",
        "esp",
        "on",
    )
    run_output(create_partitions)

    logger.info(f"{disk=} successfully partitioned")


def create_zfs_pool(pool_disks: Sequence[str], mnt_dir: str) -> None:
    """Create a ZFS pool.

    Args:
        pool_disks (Sequence[str]): A tuple of disks to use for the pool.
        mnt_dir (str): The mount directory.
    """
    if len(pool_disks) <= 0:
        error = "disks must be a tuple of at least length 1"
        raise ValueError(error)

    zpool_create = [
        "zpool",
        "create",
        "-o",
        "ashift=12",
        "-o",
        "autotrim=on",
        "-R",
        mnt_dir,
        "-O",
        "acltype=posixacl",
        "-O",
        "canmount=off",
        "-O",
        "dnodesize=auto",
        "-O",
        "normalization=formD",
        "-O",
        "relatime=on",
        "-O",
        "xattr=sa",
        "-O",
        "mountpoint=legacy",
        "-O",
        "compression=zstd",
        "-O",
        "atime=off",
        "root_pool",
    ]
    if len(pool_disks) == 1:
        zpool_create.append(pool_disks[0])
    else:
        zpool_create.append("mirror")
        zpool_create.extend(pool_disks)

    run_output(zpool_create)
    zpools = run_output(("zpool", "list", "-o", "name"))
    if "root_pool" not in zpools.splitlines():
        logger.critical("Failed to create root_pool")
        sys.exit(1)


def create_zfs_datasets() -> None:
    """Create ZFS datasets."""
    run_output(("zfs", "create", "-o", "canmount=noauto", "-o", "reservation=10G", "root_pool/root"))
    run_output(("zfs", "create", "root_pool/home"))
    run_output(("zfs", "create", "-o", "reservation=1G", "root_pool/var"))
    run_output(("zfs", "create", "-o", "compression=zstd-9", "-o", "reservation=10G", "root_pool/nix"))
    datasets = run_output(("zfs", "list", "-o", "name"))

    expected_datasets = {
        "root_pool/root",
        "root_pool/home",
        "root_pool/var",
        "root_pool/nix",
    }
    missing_datasets = expected_datasets.difference(datasets.splitlines())
    if missing_datasets:
        logger.critical(f"Failed to create pools {missing_datasets}")
        sys.exit(1)


def get_cpu_manufacturer() -> str:
    """Get the CPU manufacturer."""
    output = Path("/proc/cpuinfo").read_text()

    id_vendor = {"AuthenticAMD": "amd", "GenuineIntel": "intel"}

    for line in output.splitlines():
        if "vendor_id" in line:
            return id_vendor[line.split(": ")[1].strip()]

    error = "Failed to get CPU manufacturer"
    raise RuntimeError(error)


def get_boot_drive_id(disk: str) -> str:
    """Get the boot drive ID."""
    output = run_output(("lsblk", "-o", "UUID", f"{disk}-part1"))
    return output.splitlines()[1]


def create_nix_hardware_file(mnt_dir: str, disks: Sequence[str], encrypt: str | None) -> None:
    """Create a NixOS hardware file."""
    cpu_manufacturer = get_cpu_manufacturer()

    devices = ""
    if encrypt:
        disk = disks[0]

        devices = (
            f'     luks.devices."luks-root-pool-{disk.split("/")[-1]}-part2"'
            "= {\n"
            f'        device = "{disk}-part2";\n'
            "        bypassWorkqueues = true;\n"
            "        allowDiscards = true;\n"
            "      };\n"
        )

    host_id = format(getrandbits(32), "08x")

    nix_hardware = (
        "{ config, lib, modulesPath, ... }:\n"
        "{\n"
        '  imports = [ (modulesPath + "/installer/scan/not-detected.nix") ];\n\n'
        "  boot = {\n"
        "    initrd = {\n"
        '      availableKernelModules = [ \n        "ahci"\n        "ehci_pci"\n        "nvme"\n        "sd_mod"\n'
        '        "usb_storage"\n        "usbhid"\n        "xhci_pci"\n      ];\n'
        "      kernelModules = [ ];\n"
        f" {devices}"
        "    };\n"
        f'    kernelModules = [ "kvm-{cpu_manufacturer}" ];\n'
        "    extraModulePackages = [ ];\n"
        "  };\n\n"
        "  fileSystems = {\n"
        '    "/" = lib.mkDefault {\n      device = "root_pool/root";\n      fsType = "zfs";\n    };\n\n'
        '    "/home" = {\n      device = "root_pool/home";\n      fsType = "zfs";\n    };\n\n'
        '    "/var" = {\n      device = "root_pool/var";\n      fsType = "zfs";\n    };\n\n'
        '    "/nix" = {\n      device = "root_pool/nix";\n      fsType = "zfs";\n    };\n\n'
        '    "/boot" = {\n'
        f'      device = "/dev/disk/by-uuid/{get_boot_drive_id(disks[0])}";\n'
        '      fsType = "vfat";\n      options = [\n        "fmask=0077"\n'
        '        "dmask=0077"\n      ];\n    };\n  };\n\n'
        "  swapDevices = [ ];\n\n"
        "  networking.useDHCP = lib.mkDefault true;\n\n"
        '  nixpkgs.hostPlatform = lib.mkDefault "x86_64-linux";\n'
        f"  hardware.cpu.{cpu_manufacturer}.updateMicrocode = "
        "lib.mkDefault config.hardware.enableRedistributableFirmware;\n"
        f'  networking.hostId = "{host_id}";\n'
        "}\n"
    )

    Path(f"{mnt_dir}/etc/nixos/hardware-configuration.nix").write_text(nix_hardware)


def install_nixos(mnt_dir: str, disks: Sequence[str], encrypt: str | None) -> None:
    """Install NixOS."""
    run_output(("mount", "-o", "X-mount.mkdir", "-t", "zfs", "root_pool/root", mnt_dir))
    run_output(("mount", "-o", "X-mount.mkdir", "-t", "zfs", "root_pool/home", f"{mnt_dir}/home"))
    run_output(("mount", "-o", "X-mount.mkdir", "-t", "zfs", "root_pool/var", f"{mnt_dir}/var"))
    run_output(("mount", "-o", "X-mount.mkdir", "-t", "zfs", "root_pool/nix", f"{mnt_dir}/nix"))

    for disk in disks:
        run_output(("mkfs.vfat", "-n", "EFI", f"{disk}-part1"))

    # set up mirroring afterwards if more than one disk
    run_output(
        (
            "mount",
            "-t",
            "vfat",
            "-o",
            "fmask=0077,dmask=0077,iocharset=iso8859-1,X-mount.mkdir",
            f"{disks[0]}-part1",
            f"{mnt_dir}/boot",
        ),
    )

    run_output(("nixos-generate-config", "--root", mnt_dir))

    create_nix_hardware_file(mnt_dir, disks, encrypt)

    run(("nixos-install", "--root", mnt_dir), check=True)


def installer(
    disks: Sequence[str],
    swap_size: int,
    reserve: int,
    encrypt_key: str | None,
) -> None:
    """Main."""
    logger.info("Starting installation")
    require_commands(REQUIRED_COMMANDS)
    disks = tuple(sorted(disks))

    for disk in disks:
        partition_disk(disk, swap_size, reserve)

        if encrypt_key:
            sleep(1)
            key_input = encrypt_key.encode()
            run(
                ("cryptsetup", "luksFormat", "--type", "luks2", f"{disk}-part2", "-"),
                input=key_input,
                check=True,
            )
            run(
                ("cryptsetup", "luksOpen", f"{disk}-part2", f"luks-root-pool-{disk.split('/')[-1]}-part2", "-"),
                input=key_input,
                check=True,
            )

    # Fixed mount point for the new system; the installer runs as root on a fresh disk
    mnt_dir = "/tmp/nix_install"  # noqa: S108

    Path(mnt_dir).mkdir(parents=True, exist_ok=True)

    if encrypt_key:
        pool_disks = [f"/dev/mapper/luks-root-pool-{disk.split('/')[-1]}-part2" for disk in disks]
    else:
        pool_disks = [f"{disk}-part2" for disk in disks]

    create_zfs_pool(pool_disks, mnt_dir)

    create_zfs_datasets()

    install_nixos(mnt_dir, disks, encrypt_key)

    logger.info("Installation complete")


def main(argv: Sequence[str] | None = None) -> None:
    """Main."""
    parser = ArgumentParser(description="Install this NixOS configuration onto a ZFS root pool.")
    parser.add_argument(
        "--check",
        action="store_true",
        help="validate that the live environment has the external installer commands and exit",
    )
    parser.add_argument("--log-level", default=getenv("LOG_LEVEL", "DEBUG"), help="Python log level")
    args = parser.parse_args(argv)

    configure_logger(args.log_level)

    if args.check:
        require_commands(REQUIRED_COMMANDS)
        logger.info("installer runtime dependencies are available")
        return

    state = curses.wrapper(draw_menu)

    encrypt_key = getenv("ENCRYPT_KEY")

    if not state.selected_device_ids:
        logger.error("No disks selected; exiting without installing")
        sys.exit(1)

    logger.info("installing_nixos")
    logger.info(f"disks: {state.selected_device_ids}")
    logger.info(f"swap_size: {state.swap_size}")
    logger.info(f"reserve: {state.reserve_size}")
    logger.info(f"encrypted: {bool(encrypt_key)}")

    sleep(3)

    installer(
        disks=state.get_selected_devices(),
        swap_size=state.swap_size,
        reserve=state.reserve_size,
        encrypt_key=encrypt_key,
    )


if __name__ == "__main__":
    main()