#!/bin/bash
set -e

# ANSI color codes
BLUE='\033[0;34m'
GREEN='\033[0;32m'
RED='\033[0;31m'
NC='\033[0m'

# Check for root
if [ "$(id -u)" -ne 0 ]; then
    echo -e "${RED}This script must be run as root${NC}"
    exit 1
fi

REPO_DIR="/root/.nixos-utils"

# List and select NixOS containers
declare -a valid_ctids
echo -e "${BLUE}Available NixOS Containers:${NC}"
printf "%-8s %s\n" "VMID" "Name"
echo "----------------"

for conf in /etc/pve/lxc/*.conf; do
    if grep -q "^ostype: nixos" "$conf"; then
        vmid=$(basename "$conf" .conf)
        name=$(grep "^hostname:" "$conf" | cut -d' ' -f2)
        printf "%-8s %s\n" "$vmid" "$name"
        valid_ctids+=("$vmid")
    fi
done

if [ ${#valid_ctids[@]} -eq 0 ]; then
    echo -e "${RED}No NixOS containers found!${NC}" >&2
    exit 1
fi

# Container selection
while true; do
    echo -e "\nEnter the CTID of the container you want to configure:"
    read -r CTID

    valid_selection=0
    for ctid in "${valid_ctids[@]}"; do
        if [ "$ctid" = "$CTID" ]; then
            valid_selection=1
            break
        fi
    done

    if [ $valid_selection -eq 1 ]; then
        break
    else
        echo -e "${RED}Error: Invalid CTID selected! Please try again.${NC}" >&2
    fi
done

echo -e "${GREEN}Selected container: $CTID${NC}"

# Password setup prompt
echo -e "\nWould you like to set a password for the 'nixos' user? (y/n)"
read -r set_password

if [[ "$set_password" =~ ^[Yy]$ ]]; then
    while true; do
        echo -e "\nEnter new password for 'nixos' user:"
        read -s password
        echo -e "\nConfirm password:"
        read -s password_confirm

        if [ "$password" = "$password_confirm" ]; then
            echo -e "\n${GREEN}Password confirmed${NC}"
            break
        else
            echo -e "\n${RED}Passwords do not match. Please try again.${NC}"
            echo -e "Press Enter to continue or Ctrl+C to exit"
            read -r
        fi
    done
fi

echo -e "\nConfiguring NixOS container $CTID..."

# Setup LXC terminal configuration
PVE_CONFIG="/etc/pve/lxc/${CTID}.conf"
LXC_CONFIG="/var/lib/lxc/${CTID}/config"

# Update PVE config
if [ -f "$PVE_CONFIG" ]; then
    if ! grep -q "^lxc.init.cmd:" "$PVE_CONFIG"; then
        echo "lxc.init.cmd: /run/current-system/sw/bin/bash" >> "$PVE_CONFIG"
    else
        sed -i "s|^lxc.init.cmd:.*|lxc.init.cmd: /run/current-system/sw/bin/bash|" "$PVE_CONFIG"
    fi

    if ! grep -q "^cmode:" "$PVE_CONFIG"; then
        echo "cmode: shell" >> "$PVE_CONFIG"
    else
        sed -i "s|^cmode:.*|cmode: shell|" "$PVE_CONFIG"
    fi
fi

# Update LXC config
if [ -f "$LXC_CONFIG" ]; then
    if ! grep -q "^lxc.init.cmd" "$LXC_CONFIG"; then
        echo "lxc.init.cmd = /run/current-system/sw/bin/bash" >> "$LXC_CONFIG"
    else
        sed -i "s|^lxc.init.cmd.*|lxc.init.cmd = /run/current-system/sw/bin/bash|" "$LXC_CONFIG"
    fi

    if ! grep -q "^lxc.environment = TERM" "$LXC_CONFIG"; then
        echo "lxc.environment = TERM=linux" >> "$LXC_CONFIG"
    else
        sed -i "s|^lxc.environment = TERM.*|lxc.environment = TERM=linux|" "$LXC_CONFIG"
    fi

    if ! grep -q "^lxc.environment = PATH" "$LXC_CONFIG"; then
        echo "lxc.environment = PATH=/run/current-system/sw/bin" >> "$LXC_CONFIG"
    else
        sed -i "s|^lxc.environment = PATH.*|lxc.environment = PATH=/run/current-system/sw/bin|" "$LXC_CONFIG"
    fi
fi

# Mount and copy configs
MOUNT_POINT="/root/nixos-config-${CTID}"
CONFIG_FILE="/etc/pve/lxc/${CTID}.conf"

ROOTFS_LINE=$(grep "^rootfs:" "$CONFIG_FILE")
VOLUME_NAME=$(echo "$ROOTFS_LINE" | sed 's/rootfs: local-lvm:\([^,]*\).*/\1/')
DEVICE_PATH="/dev/pve/${VOLUME_NAME}"

mkdir -p "$MOUNT_POINT"
mount "$DEVICE_PATH" "$MOUNT_POINT"
mkdir -p "${MOUNT_POINT}/etc/nixos"

# Copy files and fix permissions for unprivileged container access
cp -r "${REPO_DIR}/nix-config/"* "${MOUNT_POINT}/etc/nixos/"
chown -R 100000:100000 "${MOUNT_POINT}/etc/nixos"
chmod -R 755 "${MOUNT_POINT}/etc/nixos"

# Set password if requested
if [[ "$set_password" =~ ^[Yy]$ ]]; then
    pct exec ${CTID} -- echo "nixos:${password}" | chpasswd
fi

umount "$MOUNT_POINT"
rmdir "$MOUNT_POINT"

# Rebuild NixOS configuration
echo "Rebuilding NixOS configuration..."
pct exec ${CTID} -- nix-channel --update
pct exec ${CTID} -- nixos-rebuild switch -I nixpkgs=/nix/var/nix/profiles/per-user/root/channels/nixos -I nixos-config=/etc/nixos/configuration.nix

echo -e "${GREEN}NixOS container $CTID has been configured successfully!${NC}"