#!/usr/bin/env bash

set -euo pipefail
IFS=$'\n\t'

# Constants
readonly NIXOS_RELEASE="24.11"                                      # Home Manager release version (Must match NixOS version)           
readonly GIT_REPO="https://git.mvl.sh/vleeuwenmenno/dotfiles.git"   # Dotfiles repository URL
readonly DOTFILES_PATH="${HOME}/.dotfiles"                          # Dotfiles directory
readonly SETUP_MARKER="${HOME}/.dotfiles-setup"                     # Setup marker file indicates setup has been run

#Color print function, usage: println "message" "color"
println() {
    color=$2
    printfe "%s\n" $color "$1" true
}

# print colored with printf (args: format, color, message ...)
printfe() {
    format=$1
    color=$2
    message=$3
    show_time=true

    # Check if $4 is explicitly set to false, otherwise default to true
    if [ ! -z "$4" ] && [ "$4" == "false" ]; then
        show_time=false
    fi

    red=$(tput setaf 1)
    green=$(tput setaf 2)
    yellow=$(tput setaf 3)
    blue=$(tput setaf 4)
    magenta=$(tput setaf 5)
    cyan=$(tput setaf 6)
    normal=$(tput sgr0)
    grey=$(tput setaf 8)

    case $color in
    "red")
        color=$red
        ;;
    "green")
        color=$green
        ;;
    "yellow")
        color=$yellow
        ;;
    "blue")
        color=$blue
        ;;
    "magenta")
        color=$magenta
        ;;
    "cyan") 
        color=$cyan
        ;;
    "grey")
        color=$grey
        ;;
    *)
        color=$normal
        ;;
    esac

    if [ "$show_time" == "false" ]; then
        printf "$color$format$normal" "$message"
        return
    fi

    printf $grey"%s" "$(date +'%H:%M:%S')"$normal
    
    case $color in
    $green | $cyan | $blue | $magenta | $normal)
        printf "$green INF $normal"
        ;;
    $yellow)
        printf "$yellow WRN $normal"
        ;;
    $red)
        printf "$red ERR $normal"
        ;;
    *)
        printf "$normal"
        ;;
    esac
    printf "$color$format$normal" "$message"
}

# Helper functions
log_info() {
    println "$1" "green"
}

log_success() {
    println "$1" "green"
}

log_error() {
    println "$1" "red" >&2
}

log_warning() {
    println "$1" "yellow" >&2
}

die() {
    log_error "$1"
    exit 1
}

# Request sudo credentials upfront
request_sudo() {
    log_info "This script requires sudo privileges to run, please enter your password."
    sudo -v || die "Failed to obtain sudo privileges"
    
    # Keep sudo credentials refreshed in the background
    (while true; do sudo -v; sleep 50; done) &
    SUDO_KEEPALIVE_PID=$!
    
    # Ensure we kill the keepalive process when the script exits
    trap 'kill $SUDO_KEEPALIVE_PID 2>/dev/null || true' EXIT
}

# Ensure we're running interactively
ensure_interactive() {
    # If stdin is not a terminal, reconnect stdin to /dev/tty
    if [ ! -t 0 ]; then
        exec < /dev/tty || die "Failed to connect to terminal. Please run the script directly instead of piping from curl"
    fi
}

confirm_symlink() {
    local link="$1"
    local msg="$2"
    if [ ! -L "$link" ]; then
        die "$msg"
    fi
}

backup_file() {
    local file="$1"
    local need_sudo="${2:-false}"

    if [ -f "$file" ]; then
        log_info "Backing up $file to $file.bak..."
        if [ "$need_sudo" = "true" ]; then
            sudo mv "$file" "$file.bak" || die "Failed to backup $file (sudo)"
        else
            mv "$file" "$file.bak" || die "Failed to backup $file"
        fi
    fi
}


check_prerequisites() {
    command -v git >/dev/null 2>&1 || die "Git is required but not installed"
    command -v sudo >/dev/null 2>&1 || die "Sudo is required but not installed"
}

validate_hostname() {
    local hostname="$1"
    if [[ -z "$hostname" || ! "$hostname" =~ ^[a-zA-Z0-9_-]+$ || ${#hostname} -gt 64 ]]; then
        return 1
    fi
    return 0
}

update_home_manager_flake() {
    local hostname="$1"
    local isServer="$2"
    local flake_file="$DOTFILES_PATH/config/home-manager/flake.nix"
    
    # Create new configuration entry
    local new_config="        \"$hostname\" = home-manager.lib.homeManagerConfiguration {
          inherit pkgs;
          modules = [ ./home.nix ];
          extraSpecialArgs = {
            inherit pkgs pkgs-unstable;
            isServer = $isServer;
            hostname = \"$hostname\";
          };
        };
        "

    # Create temporary file
    local temp_file=$(mktemp)
    
    # Find the line number where homeConfigurations = { appears
    local config_line=$(grep -n "homeConfigurations = {" "$flake_file" | cut -d: -f1)
    
    if [ -z "$config_line" ]; then
        rm "$temp_file"
        die "Could not find homeConfigurations in flake.nix"
    fi

    # Copy the file up to the line after homeConfigurations = {
    head -n "$config_line" "$flake_file" > "$temp_file"
    
    # Add the new configuration
    echo "$new_config" >> "$temp_file"
    
    # Add the rest of the file starting from the line after homeConfigurations = {
    tail -n +"$((config_line + 1))" "$flake_file" >> "$temp_file"

    # Validate the new file
    if ! nix-shell -p nixfmt --run "nixfmt $temp_file"; then
        rm "$temp_file"
        return 1
    fi

    # Replace original file
    mv "$temp_file" "$flake_file" || return 1
    log_success "Home Manager Flake configuration added successfully."
}

install_nix() {
    if command -v nix-channel >/dev/null 2>&1; then
        log_success "Detected Nix, skipping Nix setup."
        return 0
    fi

    log_info "Nix not detected, installing Nix..."
    curl --proto '=https' --tlsv1.2 -sSf -L https://install.determinate.systems/nix -o install-nix.sh || \
        die "Failed to download Nix installer"
    
    sh install-nix.sh install --no-confirm || die "Failed to install Nix"
    rm install-nix.sh || die "Failed to remove Nix installer"
    . /nix/var/nix/profiles/default/etc/profile.d/nix-daemon.sh || die "Failed to source Nix profile"
}

setup_symlinks() {
    log_info "Setting up symlinks..."

    # Backup and create symlinks for user files
    backup_file "$HOME/.bashrc"
    backup_file "$HOME/.profile"

    if [ -d "$HOME/.config/home-manager" ]; then
        log_info "Backing up ~/.config/home-manager to ~/.config/home-manager.bak..."
        mv "$HOME/.config/home-manager" "$HOME/.config/home-manager.bak" || \
            die "Failed to backup home-manager config"
    fi

    log_info "Linking ~/.config/home-manager to $DOTFILES_PATH/config/home-manager..."
    ln -s "$DOTFILES_PATH/config/home-manager" "$HOME/.config/home-manager" || \
        die "Failed to create home-manager symlink"

    # Verify symlinks
    confirm_symlink "$HOME/.config/home-manager" "Failed to set up home-manager symlink"
    log_success "Symlinks set up successfully."
}

install_home_manager() {
    if command -v home-manager >/dev/null 2>&1; then
        log_success "Home Manager already installed. Skipping..."
        return 0
    fi

    log_info "Installing Home Manager..."
    nix-channel --add "https://github.com/nix-community/home-manager/archive/release-$NIXOS_RELEASE.tar.gz" home-manager || die "Failed to add home-manager channel"
    nix-channel --update || die "Failed to update channels"
    nix-shell '<home-manager>' -A install || die "Failed to install home-manager"
}

prepare_hostname() {
    local hostname_file="$HOME/.hostname"
    local hostname="$1"

    # If a hostname was provided as argument
    if [ -n "$hostname" ]; then
        if ! validate_hostname "$hostname"; then
            die "Invalid hostname provided. Please use a valid hostname."
        fi
        log_info "Using provided hostname: $hostname"
    # Check if hostname is already set
    elif [ -f "$hostname_file" ]; then
        hostname=$(cat "$hostname_file")
        log_success "Hostname already found in $hostname_file. Using $hostname."
        return
    else
        hostname=$(hostname)
        # If hostname_file doesn't exist let's also put the hostname in there
        if [ ! -f "$hostname_file" ]; then
            echo "$hostname" > "$hostname_file" || die "Failed to save hostname"
        fi

        log_warning "No hostname provided. Defaulting to the currently set hostname on the system. ($hostname)"
        return
    fi

    log_info "Setting hostname to $hostname..."
    sudo hostnamectl set-hostname "$hostname" || die "Failed to set hostname"

    echo "$hostname" > "$hostname_file" || die "Failed to save hostname"
    log_success "Hostname set successfully."
}

warning_prompt() {
    echo ""
    log_error "Please ensure you have a backup of your data before proceeding."
    log_error "This script will modify system files and may require sudo permissions."
    echo ""
    log_info "This script has been tested on Ubuntu 22.04, 24.04, 24.10, Pop!_OS 24.04 Alpha 7, Debian 12 and Fedora 41."
    log_info "Setup starts in 10 seconds, to abort use Ctrl+C to exit NOW."
    echo ""
    sleep 10
}

check_selinux() {
    # Check if distro has SELinux at all:
    if [ ! -d /etc/selinux ]; then
        log_success "SELinux not found. Skipping..."
        return 0
    fi

    # Check if getenforce exists, if not it means we don't have SELinux
    if ! command -v getenforce >/dev/null 2>&1; then
        log_success "SELinux not found. Skipping..."
        return 0
    fi

    # Check if getenforce is returning Enforcing
    if [ "$(getenforce)" = "Enforcing" ]; then
        log_warning "SELinux is enabled. Adjusting SELinux to permissive mode..."
        sudo setenforce Permissive || die "Failed to disable SELinux"
        sudo tee /etc/selinux/config << EOF > /dev/null || die "Failed to write to /etc/selinux/config"
SELINUX=permissive
SELINUXTYPE=targeted
EOF
        log_success "SELinux disabled successfully."
    fi
}

setup_ansible() {
    attempt_package_install "ansible"
    attempt_package_install "ansible-lint"
    attempt_package_install "ansible-core"
}

check_command_availibility() {
    local command="$1"
    if ! command -v "$command" >/dev/null 2>&1; then
        die "$command is required but not installed"
    fi
}

attempt_package_install() {
    local package="$1"

    # determine which package manager to use
    local package_manager
    if command -v dnf >/dev/null 2>&1; then
        package_manager="dnf"
    elif command -v apt >/dev/null 2>&1; then
        package_manager="apt"
    elif command -v pacman >/dev/null; then
        package_manager="pacman"
    else
        log_error "No supported package manager was found, aborting setup..."
        exit 1
    fi

    if ! command -v "$package" >/dev/null 2>&1; then
        log_info "Installing $package using $package_manager..."
        if [ "$package_manager" = "dnf" ]; then
            sudo dnf install "$package" -y || die "Failed to install $package"
        elif [ "$package_manager" = "apt" ]; then
            sudo apt install "$package" -y || die "Failed to install $package"
        elif [ "$package_manager" = "pacman" ]; then
            sudo pacman -S "$package" || die "Failed to install $package"
        else
            die "Unsupported package manager: $package_manager"
        fi
    fi
}

# Check compatibility checks for supported distros:
# - Fedora
# - Ubuntu
# - Arch Linux (Untested)
check_compatibility() {
    # Check if we are running under bash
    if [ -z "${BASH_VERSION:-}" ]; then
        die "This script was designed to run using bash, please run using bash"
    fi

    attempt_package_install "awk"
    attempt_package_install "tail"
    attempt_package_install "git"

    # Check if we are a user or root
    if [ "$EUID" -eq 0 ]; then
        die "This script should not be run as root. Please run as a regular user."
    fi

    local distro
    distro=$(awk -F= '/^NAME/{print $2}' /etc/os-release | tr -d '"')

    case "$distro" in
        Fedora*)
            log_success "Detected Fedora. Proceeding with setup..."
            check_command_availibility "dnf"
            ;;
        Ubuntu)
            # Check if we are running either 22.04, 24.04 or 24.10
            if ! grep -q "Ubuntu 22.04" /etc/os-release && ! grep -q "Ubuntu 24.04" /etc/os-release && ! grep -q "Ubuntu 24.10" /etc/os-release; then
                log_warning "Unsupported Ubuntu version detected. Setup may not work as expected."
                log_warning "Supported versions are: Ubuntu 22.04, 24.04, 24.10"
            fi
            log_success "Detected Ubuntu. Proceeding with setup..."
            check_command_availibility "apt"
            ;;
        Debian*)
            log_success "Detected Debian. Proceeding with setup..."
            log_warning "Debian has known issues with ZFS kernel modules, you might need to manually install it to make ZFS work."
            log_warning "Continueing in 5 seconds..."
            sleep 5
            check_command_availibility "apt"
            ;;
        Pop!_OS*)
            log_success "Detected Pop!_OS. Proceeding with setup..."
            log_warning "Only COSMIC alpha is supported, other versions are not tested."
            log_warning "Continueing in 5 seconds..."
            sleep 5
            check_command_availibility "apt"
            ;;
        *)
            die "Unsupported distribution: $distro"
            ;;
    esac
}

# Ensure bash is set as the default shell for the user
ensure_shell() {
    local shell
    shell=$(getent passwd "$USER" | cut -d: -f7)
    if [ "$shell" != "/bin/bash" ]; then
        log_info "Setting default shell to bash..."
        chsh -s /bin/bash || die "Failed to set default shell to bash"
        log_success "Default shell set to bash."
    fi

    # Ensure shell is set for root user
    if [ "$USER" != "root" ]; then
        local root_shell
        root_shell=$(getent passwd root | cut -d: -f7)
        if [ "$root_shell" != "/bin/bash" ]; then
            log_info "Setting default shell for root to bash..."
            sudo chsh -s /bin/bash root || die "Failed to set default shell for root to bash"
            log_success "Default shell for root set to bash."
        fi
    fi
}

main() {
    check_compatibility

    # Check if setup has already been run
    if [ -f "$SETUP_MARKER" ]; then
        log_info "Setup has already been run, exiting..."
        exit 0
    fi

    # Check prerequisites
    check_prerequisites

    # First argument should be the hostname
    local hostname="${1:-}"
    local continue_flag="${2:-}"

    # Request sudo credentials upfront
    request_sudo

    # Clone dotfiles if needed
    if [ ! -d "$DOTFILES_PATH" ]; then
        log_info "Cloning dotfiles repo..."
        echo ""
        git clone "$GIT_REPO" "$DOTFILES_PATH" || die "Failed to clone dotfiles repository"
        echo ""
    fi
    
    if [ "$continue_flag" = "--continue" ]; then
        log_info "Continuing setup..."
    else
        warning_prompt
        prepare_hostname "$hostname"
        check_selinux
        install_nix
    fi

    install_home_manager

    setup_symlinks
    setup_ansible
    ensure_shell

    # Get hostname
    hostname=$(cat "$HOME/.hostname") || die "Failed to read hostname"
    export PATH=$PATH:$DOTFILES_PATH/bin

    # Create new .bashrc with exports for initial `dotf update` command
    cat << EOF >> $HOME/.bashrc
export NIXPKGS_ALLOW_INSECURE=1
export DOTFILES_PATH=${DOTFILES_PATH}
export PATH=\$PATH:\$DOTFILES_PATH/bin
export PATH=\$PATH:\$HOME/.local/bin
EOF
    cp $HOME/.bashrc $HOME/.profile

    # Create setup marker
    touch "$SETUP_MARKER" || die "Failed to create setup marker"

    # Final success message
    echo ""
    log_success "Setup complete. Please logout / restart to continue."
    echo ""
    log_error "!!! Please logout / restart to continue !!!"
    log_error "~~~   Proceed by running 'dotf update'  ~~~"
    echo ""
    log_warning "Note: For servers to be able to load secrets you might want to populate ~/.op_sat with a 1Password Service Account Token."
    echo ""
}

main "$@"