#!/bin/bash

set -e
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
DOTFILES_DIR="$(dirname "$SCRIPT_DIR")"
DOTFILES_SUBDIR="$DOTFILES_DIR/dotfiles"

# Colors for output
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
RED='\033[0;31m'
NC='\033[0m' # No Color

echo -e "${GREEN}Starting system setup...${NC}"

# Detect OS
if [ -f /etc/os-release ]; then
    . /etc/os-release
    OS_NAME=$ID
    OS_VERSION=$VERSION_ID
else
    echo -e "${YELLOW}Unable to detect OS, assuming Debian/Ubuntu...${NC}"
    OS_NAME="ubuntu"
    OS_VERSION="22.04"
fi

echo -e "${GREEN}Detected OS: ${OS_NAME} ${OS_VERSION}${NC}"

# Function to determine the package manager to use
determine_pkg_manager() {
  if command -v nala &> /dev/null; then
    echo "nala"
  elif [ "$OS_NAME" = "fedora" ]; then
    echo "dnf"
  else
    echo "apt"
  fi
}

# Set up package management based on OS
if [ "$OS_NAME" = "fedora" ]; then
    echo -e "${YELLOW}Setting up Fedora repositories and package management...${NC}"

    # Install prerequisites for Fedora
    sudo dnf install -y wget gpg

    # Setup VS Code repository for Fedora
    echo -e "${YELLOW}Setting up VS Code repository for Fedora...${NC}"
    sudo rpm --import https://packages.microsoft.com/keys/microsoft.asc
    sudo sh -c 'echo -e "[code]\nname=Visual Studio Code\nbaseurl=https://packages.microsoft.com/yumrepos/vscode\nenabled=1\ngpgcheck=1\ngpgkey=https://packages.microsoft.com/keys/microsoft.asc" > /etc/yum.repos.d/vscode.repo'

    # Setup GitHub CLI repository for Fedora
    echo -e "${YELLOW}Setting up GitHub CLI repository for Fedora...${NC}"
    sudo dnf install -y 'dnf-command(config-manager)'
    # Use a different approach to add the GitHub CLI repo to avoid the "--add-repo" error
    sudo curl -fsSL https://cli.github.com/packages/rpm/gh-cli.repo -o /etc/yum.repos.d/gh-cli.repo

    # Update package lists
    echo -e "${YELLOW}Updating package lists for Fedora...${NC}"
    sudo dnf check-update -y || true

else
    # Add apt repositories for Debian/Ubuntu
    echo -e "${YELLOW}Setting up apt repositories...${NC}"

    # Install prerequisites
    sudo apt-get install -y wget gpg apt-transport-https

    # Setup Nala repository - Try multiple methods
    echo -e "${YELLOW}Setting up Nala repository...${NC}"

    # First check if Nala is available in standard repositories (newer Ubuntu versions)
    if apt-cache show nala &>/dev/null; then
        echo -e "${GREEN}Nala is available in standard repositories${NC}"
    else
        echo -e "${YELLOW}Nala not found in standard repositories. Trying alternative installation methods...${NC}"

        # Check Ubuntu version
        if grep -q "noble\|lunar\|mantic\|jammy" /etc/os-release; then
            echo -e "${GREEN}Ubuntu $(grep VERSION_CODENAME /etc/os-release | cut -d= -f2) detected. Ensuring universe repository is enabled...${NC}"
            # Make sure universe repository is enabled
            sudo apt-get update
            sudo apt-get install -y software-properties-common
            sudo add-apt-repository -y universe
        else
            # For older Ubuntu versions try to install directly
            echo -e "${YELLOW}Older Ubuntu version detected. Trying to install Nala directly...${NC}"
            sudo apt-get update
            sudo apt-get install -y software-properties-common
        fi
    fi

    # Setup VS Code repository
    echo -e "${YELLOW}Setting up VS Code repository...${NC}"
    wget -qO- https://packages.microsoft.com/keys/microsoft.asc | gpg --dearmor > /tmp/packages.microsoft.gpg
    sudo install -D -o root -g root -m 644 /tmp/packages.microsoft.gpg /etc/apt/keyrings/packages.microsoft.gpg
    sudo sh -c 'echo "deb [arch=amd64,arm64,armhf signed-by=/etc/apt/keyrings/packages.microsoft.gpg] https://packages.microsoft.com/repos/code stable main" > /etc/apt/sources.list.d/vscode.list'
    rm /tmp/packages.microsoft.gpg

    # Setup GitHub CLI repository
    echo -e "${YELLOW}Setting up GitHub CLI repository...${NC}"
    sudo mkdir -p -m 755 /etc/apt/keyrings
    out=$(mktemp) && wget -nv -O$out https://cli.github.com/packages/githubcli-archive-keyring.gpg \
        && sudo install -D -o root -g root -m 644 "$out" /etc/apt/keyrings/githubcli-archive-keyring.gpg \
        && rm "$out" \
        && echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/githubcli-archive-keyring.gpg] https://cli.github.com/packages stable main" | sudo tee /etc/apt/sources.list.d/github-cli.list > /dev/null

    # Update package lists (ignoring previous errors)
    echo -e "${YELLOW}Updating package lists...${NC}"
    sudo apt update || echo -e "${YELLOW}Warning: apt update had some errors, but continuing...${NC}"

    # Install Nala if not already installed
    echo -e "${YELLOW}Checking for Nala package manager...${NC}"
    if ! command -v nala &> /dev/null; then
        echo -e "${YELLOW}Installing Nala package manager...${NC}"
        if sudo apt-get install -y nala; then
            echo -e "${GREEN}Nala installed successfully!${NC}"
            # Update PKG_MANAGER since we now have nala
            PKG_MANAGER="nala"
        else
            echo -e "${YELLOW}Failed to install Nala. Continuing with standard apt...${NC}"
        fi
    else
        echo -e "${GREEN}Nala is already installed!${NC}"
    fi
fi

# Determine which package manager to use
PKG_MANAGER=$(determine_pkg_manager)
echo -e "${GREEN}Using package manager: $PKG_MANAGER${NC}"

# Load packages from package list
mapfile -t pkgs < <(grep -v '^//' "$SCRIPT_DIR/packages.list" | grep -v -e '^$')

# Map Debian/Ubuntu package names to Fedora equivalents if needed
declare -A fedora_pkg_map
fedora_pkg_map["bat"]="bat"
fedora_pkg_map["fd-find"]="fd-find"
# Add more package mappings as needed

# Process the package list based on OS
install_pkg_list=()
for pkg in "${pkgs[@]}"; do
    # Skip nala package on non-Debian/Ubuntu systems
    if [ "$pkg" = "nala" ] && [ "$OS_NAME" != "ubuntu" ] && [ "$OS_NAME" != "debian" ]; then
        continue
    fi

    # Check if we need to map the package name for Fedora
    if [ "$OS_NAME" = "fedora" ] && [[ -n "${fedora_pkg_map[$pkg]}" ]]; then
        install_pkg_list+=("${fedora_pkg_map[$pkg]}")
    else
        install_pkg_list+=("$pkg")
    fi
done

# Install packages using the determined package manager
echo -e "${YELLOW}Installing packages from packages.list with $PKG_MANAGER...${NC}"

case $PKG_MANAGER in
    nala)
        # Configure Nala mirrors
        echo -e "${YELLOW}Configuring Nala mirrors...${NC}"
        # First, remove any existing nala sources list to avoid issues with invalid mirrors
        sudo rm -f /etc/apt/sources.list.d/nala-sources.list 2>/dev/null

        # Try to fetch mirrors with less aggressive settings
        if ! sudo nala fetch --auto --fetches 1 --country auto; then
            echo -e "${YELLOW}Mirror selection failed, continuing with system default mirrors...${NC}"
            # Remove any potentially corrupted Nala sources
            sudo rm -f /etc/apt/sources.list.d/nala-sources.list 2>/dev/null
        fi

        # Install packages using Nala
        sudo nala install -y "${install_pkg_list[@]}" || {
            echo -e "${YELLOW}Nala install failed, falling back to apt...${NC}"
            sudo apt-get install -y "${install_pkg_list[@]}"
        }
        ;;
    dnf)
        # Install packages using DNF
        sudo dnf install -y "${install_pkg_list[@]}"
        ;;
    apt)
        # Install packages using APT
        sudo apt-get install -y "${install_pkg_list[@]}"
        ;;
esac

echo -e "${GREEN}Package installation completed for $OS_NAME $OS_VERSION.${NC}"

# Install Zsh if not already installed
echo -e "${YELLOW}Installing Zsh...${NC}"
if ! command -v zsh &> /dev/null; then
    case $PKG_MANAGER in
        nala)
            sudo nala install -y zsh
            ;;
        dnf)
            sudo dnf install -y zsh
            ;;
        apt)
            sudo apt-get install -y zsh
            ;;
    esac
fi

# Install Oh My Zsh
echo -e "${YELLOW}Installing Oh My Zsh...${NC}"
if [ ! -d "$HOME/.oh-my-zsh" ]; then
    sh -c "$(curl -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)" "" --unattended
else
    echo "Oh My Zsh already installed"
fi

# Install zoxide
echo -e "${YELLOW}Installing zoxide...${NC}"
if ! command -v zoxide &> /dev/null; then
    curl -sS https://raw.githubusercontent.com/ajeetdsouza/zoxide/main/install.sh | bash

    # Ensure .local/bin is in PATH
    if [[ ":$PATH:" != *":$HOME/.local/bin:"* ]]; then
        echo 'export PATH="$PATH:$HOME/.local/bin"' >> "$HOME/.bashrc"
        export PATH="$PATH:$HOME/.local/bin"
    fi
fi

# Install nvm (Node Version Manager)
echo -e "${YELLOW}Installing nvm...${NC}"
if [ ! -d "$HOME/.nvm" ]; then
    curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.39.5/install.sh | bash
fi

# Load nvm regardless of whether it was just installed or already existed
export NVM_DIR="$HOME/.nvm"
[ -s "$NVM_DIR/nvm.sh" ] && \. "$NVM_DIR/nvm.sh"
[ -s "$NVM_DIR/bash_completion" ] && \. "$NVM_DIR/bash_completion"

# Install and set up Node.js only if nvm is available
if command -v nvm &>/dev/null; then
    # Install latest LTS version of Node.js
    nvm install --lts
    nvm use --lts
    # Set the LTS version as default
    nvm alias default 'lts/*'
else
    echo -e "${YELLOW}Warning: nvm installation may require a new shell session${NC}"
fi

# Install Lazydocker (not available in apt repositories)
echo -e "${YELLOW}Installing Lazydocker...${NC}"
if ! command -v lazydocker &> /dev/null; then
    LAZYDOCKER_VERSION=$(curl -s "https://api.github.com/repos/jesseduffield/lazydocker/releases/latest" | grep -Po '"tag_name": "v\K[^"]*')
    curl -Lo lazydocker.tar.gz "https://github.com/jesseduffield/lazydocker/releases/latest/download/lazydocker_${LAZYDOCKER_VERSION}_Linux_x86_64.tar.gz"
    mkdir -p lazydocker-temp
    tar xf lazydocker.tar.gz -C lazydocker-temp
    sudo mv lazydocker-temp/lazydocker /usr/local/bin
    rm -rf lazydocker-temp lazydocker.tar.gz
    echo -e "${GREEN}Lazydocker installed successfully!${NC}"
else
    echo -e "Lazydocker is already installed"
fi

# Define a reusable function for cloning Zsh plugins
clone_zsh_plugin() {
    local plugin_url=$1
    local plugin_dir=$2

    if [ ! -d "$plugin_dir" ]; then
        git clone "$plugin_url" "$plugin_dir"
    fi
}

# Install Zsh plugins using the reusable function
echo -e "${YELLOW}Installing Zsh plugins...${NC}"
ZSH_CUSTOM="$HOME/.oh-my-zsh/custom"
PLUGINS_DIR="$ZSH_CUSTOM/plugins"
clone_zsh_plugin "https://github.com/zsh-users/zsh-autosuggestions" "$PLUGINS_DIR/zsh-autosuggestions"
clone_zsh_plugin "https://github.com/zsh-users/zsh-syntax-highlighting" "$PLUGINS_DIR/zsh-syntax-highlighting"
clone_zsh_plugin "https://github.com/MichaelAquilina/zsh-you-should-use" "$PLUGINS_DIR/zsh-you-should-use"

# Set up dotfiles
echo -e "${YELLOW}Setting up dotfiles...${NC}"
# Consolidate symbolic link creation for dotfiles
ln -sf "$DOTFILES_SUBDIR/.zshrc" "$HOME/.zshrc"
ln -sf "$DOTFILES_SUBDIR/.nanorc" "$HOME/.nanorc" 2>/dev/null || true
ln -sf "$DOTFILES_SUBDIR/.profile" "$HOME/.profile" 2>/dev/null || true
ln -sf "$DOTFILES_SUBDIR/.gitconfig" "$HOME/.gitconfig" 2>/dev/null || true

if [ "$OS_NAME" = "fedora" ]; then
    echo -e "${YELLOW}Setting up Fedora-specific configurations...${NC}"
    echo -e "${GREEN}Fedora-specific setup completed.${NC}"
fi

# Set Zsh as default shell if not already set
if [[ "$SHELL" != *"zsh"* ]]; then
    echo -e "${YELLOW}Setting Zsh as default shell...${NC}"
    ZSH_PATH=$(which zsh)
    if [ -f "$ZSH_PATH" ]; then
        # Check if zsh is already in /etc/shells
        if ! grep -q "$ZSH_PATH" /etc/shells; then
            echo -e "${YELLOW}Adding $ZSH_PATH to /etc/shells...${NC}"
            echo "$ZSH_PATH" | sudo tee -a /etc/shells
        fi
        # Set as default shell
        chsh -s "$ZSH_PATH"
    else
        echo -e "${RED}Zsh not found at $ZSH_PATH. Please set it as your default shell manually.${NC}"
    fi
fi

echo -e "${GREEN}Setup completed successfully for $OS_NAME $OS_VERSION!${NC}"
echo -e "${YELLOW}Note: You may need to log out and log back in for all changes to take effect.${NC}"