#!/usr/bin/env bash
# Nomarchy Fresh Installation Script
# For use from the live ISO to install on a new disk

set -euo pipefail

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
NC='\033[0m'
BOLD='\033[1m'

# Installation target
TARGET_DISK=""
USE_LUKS=true
LUKS_PARTITION=""
ROOT_PARTITION=""
EFI_PARTITION=""

print_banner() {
    echo -e "${CYAN}"
    cat << 'EOF'
    _   __                          __
   / | / /___  ____ ___  ____ _____/ /___  __
  /  |/ / __ \/ __ `__ \/ __ `/ __/ / __/ / /
 / /|  / /_/ / / / / / / /_/ / / / / /__ / /
/_/ |_/\____/_/ /_/ /_/\__,_/_/ /_/\___//_/

        Fresh Installation
EOF
    echo -e "${NC}"
    echo ""
}

check_root() {
    if [ "$(id -u)" != "0" ]; then
        echo -e "${RED}Error: This installer must be run as root.${NC}"
        echo "Run: sudo nomarchy-fresh-install"
        exit 1
    fi
}

check_uefi() {
    if [ ! -d /sys/firmware/efi ]; then
        echo -e "${YELLOW}Warning: System is not booted in UEFI mode.${NC}"
        echo "Nomarchy requires UEFI. Please reboot in UEFI mode."
        read -rp "Continue anyway (legacy BIOS)? [y/N]: " cont
        if [[ ! "$cont" =~ ^[Yy] ]]; then
            exit 1
        fi
    else
        echo -e "${GREEN}UEFI mode detected.${NC}"
    fi
}

list_disks() {
    echo -e "${BOLD}Available disks:${NC}"
    echo ""
    lsblk -d -o NAME,SIZE,MODEL,TYPE | grep -E "disk|NAME"
    echo ""
}

select_disk() {
    list_disks

    echo -e "${YELLOW}WARNING: The selected disk will be completely erased!${NC}"
    echo ""
    read -rp "Enter disk to install to (e.g., sda, vda, nvme0n1): " disk_input

    if [ -z "$disk_input" ]; then
        echo -e "${RED}No disk selected. Aborting.${NC}"
        exit 1
    fi

    # Handle nvme vs sata naming
    if [[ "$disk_input" == nvme* ]]; then
        TARGET_DISK="/dev/${disk_input}"
        EFI_PARTITION="${TARGET_DISK}p1"
        ROOT_PARTITION="${TARGET_DISK}p2"
    else
        TARGET_DISK="/dev/${disk_input}"
        EFI_PARTITION="${TARGET_DISK}1"
        ROOT_PARTITION="${TARGET_DISK}2"
    fi

    if [ ! -b "$TARGET_DISK" ]; then
        echo -e "${RED}Error: ${TARGET_DISK} is not a valid block device.${NC}"
        exit 1
    fi

    echo ""
    echo -e "${RED}${BOLD}WARNING: ALL DATA ON ${TARGET_DISK} WILL BE DESTROYED!${NC}"
    echo ""
    read -rp "Type 'yes' to confirm: " confirm

    if [ "$confirm" != "yes" ]; then
        echo "Aborting."
        exit 1
    fi
}

prompt_encryption() {
    echo ""
    echo -e "${BOLD}Disk Encryption${NC}"
    echo "LUKS encryption protects your data if the device is lost or stolen."
    read -rp "Enable LUKS encryption? [Y/n]: " luks_choice
    luks_choice="${luks_choice:-y}"

    if [[ "$luks_choice" =~ ^[Yy] ]]; then
        USE_LUKS=true
        echo -e "${GREEN}LUKS encryption enabled.${NC}"
    else
        USE_LUKS=false
        echo -e "${YELLOW}Installing without encryption.${NC}"
    fi
    echo ""
}

partition_disk() {
    echo -e "${BLUE}Partitioning ${TARGET_DISK}...${NC}"

    # Wipe existing partition table
    wipefs -a "$TARGET_DISK"

    # Create GPT partition table with EFI and root partitions
    parted -s "$TARGET_DISK" mklabel gpt
    parted -s "$TARGET_DISK" mkpart ESP fat32 1MiB 512MiB
    parted -s "$TARGET_DISK" set 1 esp on
    parted -s "$TARGET_DISK" mkpart primary 512MiB 100%

    # Wait for partitions to appear
    sleep 2
    partprobe "$TARGET_DISK"
    sleep 1

    echo -e "${GREEN}Partitioning complete.${NC}"
}

setup_luks() {
    if [ "$USE_LUKS" = true ]; then
        echo ""
        echo -e "${BLUE}Setting up LUKS encryption...${NC}"
        echo -e "${YELLOW}You will be prompted to enter a passphrase.${NC}"
        echo ""

        cryptsetup luksFormat --type luks2 "$ROOT_PARTITION"
        cryptsetup open "$ROOT_PARTITION" cryptroot

        LUKS_PARTITION="$ROOT_PARTITION"
        ROOT_PARTITION="/dev/mapper/cryptroot"

        echo -e "${GREEN}LUKS encryption configured.${NC}"
    fi
}

format_filesystems() {
    echo ""
    echo -e "${BLUE}Formatting filesystems...${NC}"

    # Format EFI partition
    mkfs.fat -F32 "$EFI_PARTITION"

    # Format root partition with ext4
    mkfs.ext4 -L nixos "$ROOT_PARTITION"

    echo -e "${GREEN}Filesystems formatted.${NC}"
}

mount_filesystems() {
    echo ""
    echo -e "${BLUE}Mounting filesystems...${NC}"

    mount "$ROOT_PARTITION" /mnt
    mkdir -p /mnt/boot
    mount "$EFI_PARTITION" /mnt/boot

    echo -e "${GREEN}Filesystems mounted.${NC}"
}

generate_hardware_config() {
    echo ""
    echo -e "${BLUE}Generating hardware configuration...${NC}"

    nixos-generate-config --root /mnt

    echo -e "${GREEN}Hardware configuration generated.${NC}"
}

# User configuration prompts (similar to install.sh)
prompt_username() {
    echo ""
    echo -e "${BOLD}Username${NC}"
    read -rp "Enter username: " username
    username="${username:-user}"
}

prompt_hostname() {
    echo ""
    echo -e "${BOLD}Hostname${NC}"
    read -rp "Enter hostname: " hostname_input
    hostname_input="${hostname_input:-nomarchy}"
}

prompt_timezone() {
    echo ""
    echo -e "${BOLD}Timezone${NC}"
    echo "Examples: America/New_York, Europe/London, Asia/Tokyo"
    read -rp "Enter timezone [UTC]: " timezone
    timezone="${timezone:-UTC}"
}

prompt_locale() {
    echo ""
    echo -e "${BOLD}Locale${NC}"
    read -rp "Enter locale [en_US.UTF-8]: " locale
    locale="${locale:-en_US.UTF-8}"
}

prompt_keyboard() {
    echo ""
    echo -e "${BOLD}Keyboard Layout${NC}"
    read -rp "Enter keyboard layout [us]: " keyboard
    keyboard="${keyboard:-us}"
}

prompt_bootloader() {
    echo ""
    echo -e "${BOLD}Bootloader${NC}"
    echo "1) systemd-boot - Simple, reliable"
    echo "2) Limine - Modern, themed to match Nomarchy"
    read -rp "Choose bootloader [1/2]: " bootloader_choice
    bootloader_choice="${bootloader_choice:-1}"

    if [[ "$bootloader_choice" == "2" ]]; then
        bootloader="limine"
    else
        bootloader="systemd-boot"
    fi
}

prompt_features() {
    echo ""
    echo -e "${BOLD}Optional Features${NC}"

    read -rp "Enable Plymouth boot splash? [y/N]: " plymouth
    [[ "$plymouth" =~ ^[Yy] ]] && enable_plymouth="true" || enable_plymouth="false"

    read -rp "Enable Syncthing? [Y/n]: " syncthing
    syncthing="${syncthing:-y}"
    [[ "$syncthing" =~ ^[Yy] ]] && enable_syncthing="true" || enable_syncthing="false"

    read -rp "Enable printing support? [Y/n]: " printing
    printing="${printing:-y}"
    [[ "$printing" =~ ^[Yy] ]] && enable_printing="true" || enable_printing="false"
}

generate_config() {
    echo ""
    echo -e "${BLUE}Generating Nomarchy configuration...${NC}"

    local config_dir="/mnt/etc/nomarchy"
    mkdir -p "$config_dir"

    # Keyboard layouts as Nix list
    local kb_layouts_nix
    kb_layouts_nix=$(echo "$keyboard" | sed 's/,/" "/g' | sed 's/^/"/' | sed 's/$/"/')

    cat > "$config_dir/config.nix" << EOF
# Nomarchy Configuration
# Generated during installation on $(date)
{
  username = "${username}";
  hostname = "${hostname_input}";
  timezone = "${timezone}";
  locale = "${locale}";
  keyboardLayouts = [${kb_layouts_nix}];
  keyboardVariants = [""];

  location = {
    lat = 0;
    lon = 0;
  };

  enableSyncthing = ${enable_syncthing};
  enableMullvad = false;
  enablePrinting = ${enable_printing};
  enableBluetooth = true;

  bootloader = "${bootloader}";
  enablePlymouth = ${enable_plymouth};

  enableMitigationsOff = false;

  theme = "classical";
}
EOF

    # Generate flake.nix
    cat > "$config_dir/flake.nix" << 'FLAKEEOF'
{
  description = "Nomarchy NixOS Configuration";

  inputs = {
    nomarchy.url = "github:blu/nomarchy";
    nixpkgs.follows = "nomarchy/nixpkgs";
  };

  outputs = { self, nomarchy, nixpkgs, ... }@inputs: let
    userConfig = import ./config.nix;
  in {
    nixosConfigurations.${userConfig.hostname} = nomarchy.lib.mkHost {
      system = "x86_64-linux";
      config = userConfig;
      extraModules = [
        /etc/nixos/hardware-configuration.nix
      ];
    };
  };
}
FLAKEEOF

    # Build LUKS config if needed
    local luks_config=""
    if [ "$USE_LUKS" = true ]; then
        local luks_uuid
        luks_uuid=$(blkid -s UUID -o value "$LUKS_PARTITION")
        luks_config="
  # LUKS encryption
  boot.initrd.luks.devices.\"cryptroot\" = {
    device = \"/dev/disk/by-uuid/${luks_uuid}\";
    preLVM = true;
  };"
    fi

    # Create configuration.nix with LUKS config included
    cat > "/mnt/etc/nixos/configuration.nix" << EOF
# This file is managed by Nomarchy
# Edit /etc/nomarchy/config.nix instead
{ ... }: {
  imports = [ ./hardware-configuration.nix ];

  # Minimal config for initial boot - will be replaced by flake
  boot.loader.systemd-boot.enable = true;
  boot.loader.efi.canTouchEfiVariables = true;

  networking.hostName = "${hostname_input}";
  time.timeZone = "${timezone}";

  users.users.${username} = {
    isNormalUser = true;
    extraGroups = [ "wheel" "networkmanager" ];
    initialPassword = "changeme";
  };

  services.openssh.enable = true;
  networking.networkmanager.enable = true;
${luks_config}
  system.stateVersion = "24.11";
}
EOF

    echo -e "${GREEN}Configuration generated.${NC}"
}

run_install() {
    echo ""
    echo -e "${BLUE}Installing NixOS...${NC}"
    echo "This may take a while."
    echo ""

    NIX_CONFIG="experimental-features = nix-command flakes" nixos-install --no-root-passwd

    echo ""
    echo -e "${GREEN}Installation complete!${NC}"
}

set_user_password() {
    echo ""
    echo -e "${BOLD}Set password for ${username}${NC}"
    nixos-enter --root /mnt -c "passwd ${username}"
}

show_summary() {
    echo ""
    echo -e "${BOLD}Installation Summary${NC}"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo -e "Disk:        ${CYAN}${TARGET_DISK}${NC}"
    echo -e "Encryption:  ${CYAN}${USE_LUKS}${NC}"
    echo -e "Username:    ${CYAN}${username}${NC}"
    echo -e "Hostname:    ${CYAN}${hostname_input}${NC}"
    echo -e "Timezone:    ${CYAN}${timezone}${NC}"
    echo -e "Bootloader:  ${CYAN}${bootloader}${NC}"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""
}

finish() {
    echo ""
    echo -e "${GREEN}${BOLD}Nomarchy installation complete!${NC}"
    echo ""
    echo "Next steps:"
    echo "  1. Reboot: shutdown -r now"
    echo "  2. Log in as '${username}' with the password you set"
    echo "  3. Select 'Hyprland' from the session menu"
    echo "  4. Press Super+/ for Quick Actions"
    echo ""
    echo "To rebuild with flake (after first boot):"
    echo "  cd /etc/nomarchy && sudo nixos-rebuild switch --flake .#${hostname_input}"
    echo ""
}

main() {
    print_banner
    check_root
    check_uefi

    select_disk
    prompt_encryption
    partition_disk
    setup_luks
    format_filesystems
    mount_filesystems
    generate_hardware_config

    # User configuration
    prompt_username
    prompt_hostname
    prompt_timezone
    prompt_locale
    prompt_keyboard
    prompt_bootloader
    prompt_features

    show_summary

    read -rp "Proceed with installation? [Y/n]: " proceed
    proceed="${proceed:-y}"
    if [[ ! "$proceed" =~ ^[Yy] ]]; then
        echo "Installation cancelled."
        exit 0
    fi

    generate_config
    run_install
    set_user_password
    finish
}

main "$@"