#!/bin/bash

# =============================================
# Raspberry Pi Setup Script
# Downloads and configures system scripts from Git
# =============================================

# Configuration
GITHUB_REPO="https://gitea.spiralragetech.com/tiago.aica/scripts/raw/branch/main/debian"
SCRIPT_DIR="/usr/local/bin"
LOG_FILE="/var/log/raspberry_setup.log"

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

# Function to log messages
log() {
    echo -e "$(date '+%Y-%m-%d %H:%M:%S') - $1" | tee -a "$LOG_FILE"
}

# Function to ask for confirmation
ask_confirmation() {
    local message="$1"
    local default="$2"  # "y" or "n"

    while true; do
        if [ "$default" = "y" ]; then
            read -p "$message [Y/n]: " choice < /dev/tty
        else
            read -p "$message [y/N]: " choice < /dev/tty
        fi

        # Default choice if user just presses Enter
        if [ -z "$choice" ]; then
            choice="$default"
        fi

        case "$choice" in
            y|Y|yes|Yes|YES)
                return 0
                ;;
            n|N|no|No|NO)
                return 1
                ;;
            *)
                echo "Please answer yes or no."
                ;;
        esac
    done
}

# Function to validate MAC address format
validate_mac() {
    local mac=$1
    # MAC address validation regex
    local mac_regex='^([0-9A-Fa-f]{2}[:-]){5}([0-9A-Fa-f]{2})$'
    if [[ $mac =~ $mac_regex ]]; then
        return 0
    else
        return 1
    fi
}

# Function to validate IP address format
validate_ip() {
    local ip=$1
    # IP address validation regex
    local ip_regex='^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+(/[0-9]+)?$'
    if [[ $ip =~ $ip_regex ]]; then
        return 0
    else
        return 1
    fi
}

# Function to validate gateway format
validate_gateway() {
    local gateway=$1
    # Gateway validation regex
    local gateway_regex='^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$'
    if [[ $gateway =~ $gateway_regex ]]; then
        return 0
    else
        return 1
    fi
}

# Function to download file from GitLab
download_file() {
    local filename=$1
    local destination=$2

    log "Downloading $filename..."
    if curl -s -o "$destination/$filename" "$GITHUB_REPO/$filename"; then
        log "${GREEN}Successfully downloaded $filename${NC}"
        chmod +x "$destination/$filename"
        return 0
    else
        log "${RED}Failed to download $filename${NC}"
        return 1
    fi
}

# Função para configurar o endereço MAC
configure_mac() {
    local mac_script="$SCRIPT_DIR/set_mac_wifi.sh"
    local new_mac=$1

    if [ ! -f "$mac_script" ]; then
        log "${RED}MAC address script not found${NC}"
        return 1
    fi

    # Validate MAC address format if provided
    if [ -n "$new_mac" ]; then
        if validate_mac "$new_mac"; then
            # Normalize to use colons
            new_mac=$(echo "$new_mac" | tr '-' ':')
            if bash "$mac_script" "$new_mac"; then
                log "${GREEN}MAC address configuration completed successfully with $new_mac${NC}"
                return 0
            else
                log "${RED}Failed to execute MAC address configuration script${NC}"
                return 1
            fi
        else
            log "${RED}Invalid MAC address format: $new_mac. Please use format XX:XX:XX:XX:XX:XX or XX-XX-XX-XX-XX-XX${NC}"
            return 1
        fi
    else
        # If no MAC address was provided, run the script to generate one automatically
        if bash "$mac_script"; then
            log "${GREEN}MAC address configuration completed successfully${NC}"
            return 0
        else
            log "${RED}Failed to execute MAC address configuration script${NC}"
            return 1
        fi
    fi
}

# Função para configurar o endereço MAC
setup_mac_address() {
    if ask_confirmation "Do you want to configure the MAC address for eth0 automatically?" "n"; then
        download_file "set_mac_wifi.sh" "$SCRIPT_DIR" || return 1

        if configure_mac "$1"; then
            log "${GREEN}MAC address setup completed successfully${NC}"
            return 0
        else
            log "${RED}Failed to configure MAC address${NC}"
            return 1
        fi
    else
        log "${YELLOW}Skipping MAC address configuration${NC}"
        return 1
    fi
}

# Function to check if crontab entry exists
crontab_entry_exists() {
    local script_path="$1"
    crontab -l 2>/dev/null | grep -q "@reboot $script_path"
    return $?
}

# Function to setup crontab
setup_crontab() {
    if ask_confirmation "Do you want to setup crontab to run MAC script at boot?" "y"; then
        local mac_script="$SCRIPT_DIR/set_mac_wifi.sh"

        # Download script if not present
        if [ ! -f "$mac_script" ]; then
            download_file "set_mac_wifi.sh" "$SCRIPT_DIR" || return 1
        fi

        if crontab_entry_exists "$mac_script"; then
            log "${YELLOW}Crontab entry for MAC address script already exists${NC}"
            return 0
        fi

        log "Setting up crontab for MAC address script..."
        (crontab -l 2>/dev/null; echo "@reboot $mac_script") | crontab -
        log "${GREEN}Crontab configured successfully${NC}"
        return 0
    else
        log "${YELLOW}Skipping crontab configuration${NC}"
        return 1
    fi
}

# Function to install dependencies and system updates
install_dependencies() {
    if ask_confirmation "Do you want to update the system and install required dependencies?" "y"; then
        log "Updating package lists and upgrading installed packages..."

        # Update package lists
        if ! apt update; then
            log "${RED}Failed to update package lists${NC}"
            return 1
        fi

        # Upgrade installed packages (excluding kernel)
        if ask_confirmation "Do you want to perform a packages upgrade?" "n"; then
            log "Performing package upgrade..."
            if ! apt upgrade -y --without-new-pkgs; then
                log "${RED}Failed to upgrade installed packages${NC}"
                return 1
            fi
        fi

        # Upgrade distribution (including kernel)
        if ask_confirmation "Do you want to perform a distribution upgrade (may include kernel updates)?" "n"; then

            if [ -f /proc/device-tree/model ] && grep -q "Raspberry Pi" /proc/device-tree/model; then
                log "Raspberry Pi detected - using rpi-update for firmware/kernel updates"

                # Install rpi-update if not already installed
                if ! command -v rpi-update >/dev/null 2>&1; then
                    log "Installing rpi-update..."
                    if ! apt install -y rpi-update; then
                        log "${RED}Failed to install rpi-update${NC}"
                        return 1
                    fi
                fi

                # # Update firmware and kernel
                # log "Updating Raspberry Pi firmware and kernel..."
                # if ! rpi-update; then
                #     log "${RED}Failed to update Raspberry Pi firmware/kernel${NC}"
                #     return 1
                # fi
                # log "${GREEN}Raspberry Pi firmware and kernel updated successfully${NC}"
                # log "${YELLOW}Note: A reboot is required for kernel updates to take effect${NC}"

                
                log "Updating Raspberry Pi firmware and kernel..."

                # Install 'yes' command if not available (usually pre-installed)
                if ! command -v yes >/dev/null 2>&1; then
                    if ! apt install -y coreutils; then
                        log "${RED}Failed to install coreutils package (needed for 'yes' command)${NC}"
                        return 1
                    fi
                fi

                # Run rpi-update with all prompts automatically answered with 'y'
                if ! yes | rpi-update; then
                    log "${RED}Failed to update Raspberry Pi firmware/kernel${NC}"
                    
                    log "Trying common distribution upgrade..."
                    if ! apt full-upgrade -y; then
                        log "${RED}Failed to perform distribution upgrade${NC}"
                        return 1
                    fi

                    return 1
                fi

                log "${GREEN}Raspberry Pi firmware and kernel updated successfully${NC}"
                log "${YELLOW}Note: A reboot is required for kernel updates to take effect${NC}"

                
            else
                log "Performing distribution upgrade..."
                if ! apt full-upgrade -y; then
                    log "${RED}Failed to perform distribution upgrade${NC}"
                    return 1
                fi
            fi
        fi

        # Install required dependencies
        if ask_confirmation "Do you want to install the required dependencies (figlet, bc, lsof, net-tools, macchanger)?" "n"; then
            log "Installing required dependencies (figlet, bc, lsof, net-tools, macchanger)..."
            if apt install -y figlet bc lsof net-tools macchanger; then
        # if ask_confirmation "Do you want to install the required dependencies (figlet, bc, lsof, net-tools)?" "n"; then
        #     log "Installing required dependencies (figlet, bc, lsof, net-tools)..."
        #     if apt install -y figlet bc lsof net-tools; then
                log "${GREEN}Dependencies installed successfully${NC}"
                return 0
            else
                log "${RED}Failed to install dependencies${NC}"
                return 1
            fi
        fi

        # Clean up
        if ask_confirmation "Do you want to clean the package cache?" "n"; then
            log "Cleaning up package cache..."
            if apt autoremove -y && apt clean; then
                log "${GREEN}Package cache cleaned successfully${NC}"
            else
                log "${YELLOW}Warning: Failed to clean package cache${NC}"
            fi
        fi
        
    else
        log "${YELLOW}Skipping system update and dependency installation${NC}"
        return 1
    fi
}

# Function to setup GUI manager
setup_gui_manager() {
    if ask_confirmation "Do you want to configure the graphical environment?" "y"; then
        download_file "gui_manager.sh" "$SCRIPT_DIR" || return 1

        local gui_script="$SCRIPT_DIR/gui_manager.sh"

        # If GUI action was provided as parameter, use it
        if [ "$1" = "disable_gui" ]; then
            log "Disabling graphical environment as requested..."
            GUI_ACTION=disable bash "$gui_script"
            # Provide instructions for re-enabling
            log "${YELLOW}Note:${NC} The graphical environment has been disabled."
            log "${YELLOW}To re-enable it later, run:${NC}"
            log "${YELLOW}  sudo GUI_ACTION=enable $gui_script${NC}"
            return 0
        fi

        # Ask if user wants to disable GUI
        if ask_confirmation "Do you want to disable the graphical environment to save RAM?" "n"; then
            log "Disabling graphical environment..."
            GUI_ACTION=disable bash "$gui_script"
            # Provide instructions for re-enabling
            log "${YELLOW}Note:${NC} The graphical environment has been disabled."
            log "${YELLOW}To re-enable it later, run:${NC}"
            log "${YELLOW}  sudo GUI_ACTION=enable $gui_script${NC}"
        else
            log "${YELLOW}Graphical environment will remain enabled${NC}"
        fi
        return 0
    else
        log "${YELLOW}Skipping GUI manager configuration${NC}"
        return 1
    fi
}

# Function to check if welcome script is already in profile
welcome_script_in_profile() {
    local welcome_script="$1"
    grep -q "$welcome_script" /etc/profile
    return $?
}

# Function to configure welcome script
configure_welcome() {
    if ask_confirmation "Do you want to configure the welcome script?" "y"; then
        download_file "welcome.sh" "$SCRIPT_DIR" || return 1

        local welcome_script="$SCRIPT_DIR/welcome.sh"

        if welcome_script_in_profile "$welcome_script"; then
            log "${YELLOW}Welcome script already configured in /etc/profile${NC}"
            return 0
        fi

        log "Configuring welcome script..."
        echo "$welcome_script" | sudo tee -a /etc/profile > /dev/null
        log "${GREEN}Welcome script added to /etc/profile${NC}"
        return 0
    else
        log "${YELLOW}Skipping welcome script configuration${NC}"
        return 1
    fi
}

# Function to setup static IP
setup_static_ip() {
    if ask_confirmation "Do you want to configure a static IP address?" "y"; then
        download_file "set_static_ip.sh" "$SCRIPT_DIR" || return 1

        local ip_script="$SCRIPT_DIR/set_static_ip.sh"
        local ip_address="$1"      # IP address parameter
        local gateway_address="$2" # Gateway parameter

        # If both IP and gateway were provided as parameters, use them
        if [ -n "$ip_address" ] && [ -n "$gateway_address" ]; then
            log "Configuring static IP as requested: $ip_address with gateway $gateway_address"

            # Validate IP and gateway formats
            if validate_ip "$ip_address" && validate_gateway "$gateway_address"; then
                IP_ADDRESS_ENV="$ip_address" GATEWAY_ENV="$gateway_address" bash "$ip_script"
                return 0
            else
                log "${RED}Invalid IP or gateway format. Using default configuration.${NC}"
                return 1
            fi
        fi

        # If only IP was provided, ask for gateway
        if [ -n "$ip_address" ]; then
            # Validate IP format first
            if ! validate_ip "$ip_address"; then
                log "${RED}Invalid IP address format: $ip_address.${NC}"
                return 1
            fi

            # Ask for gateway
            while true; do
                read -p "Enter gateway IP address (format: 192.168.1.1): " gateway_address < /dev/tty
                if validate_gateway "$gateway_address"; then
                    log "Configuring static IP: $ip_address with gateway $gateway_address"
                    IP_ADDRESS_ENV="$ip_address" GATEWAY_ENV="$gateway_address" bash "$ip_script"
                    return 0
                else
                    log "${RED}Invalid gateway format. Please use format like 192.168.1.1${NC}"
                fi
            done
        else
            # Prompt for both IP and gateway
            while true; do
                read -p "Enter static IP address (format: 192.168.1.100 or 192.168.1.100/24): " ip_address < /dev/tty
                if validate_ip "$ip_address"; then
                    while true; do
                        read -p "Enter gateway IP address (format: 192.168.1.1): " gateway_address < /dev/tty
                        if validate_gateway "$gateway_address"; then
                            log "Configuring static IP: $ip_address with gateway $gateway_address"
                            IP_ADDRESS_ENV="$ip_address" GATEWAY_ENV="$gateway_address" bash "$ip_script"
                            return 0
                        else
                            log "${RED}Invalid gateway format. Please use format like 192.168.1.1${NC}"
                        fi
                    done
                else
                    log "${RED}Invalid IP address format. Please use format like 192.168.1.100 or 192.168.1.100/24${NC}"
                fi
            done
        fi
    else
        log "${YELLOW}Skipping static IP configuration${NC}"
        return 1
    fi
}

# Function to upgrade Debian distribution
upgrade_debian_distro() {
    if ! ask_confirmation "Do you want to check for Debian distribution upgrades? (WARNING: This may break Raspberry Pi OS compatibility)" "n"; then
        log "${YELLOW}Skipping Debian distribution upgrade${NC}"
        return 1
    fi

    # Download the upgrade script
    download_file "upgrade_debian.sh" "$SCRIPT_DIR" || {
        log "${RED}Failed to download upgrade_debian.sh${NC}"
        return 1
    }

    # Execute the upgrade script
    log "Starting Debian distribution upgrade..."
    bash "$SCRIPT_DIR/upgrade_debian.sh" || {
        log "${RED}Failed to upgrade Debian distribution${NC}"
        return 1
    }

    return 0
}

# Function to prompt for reboot
prompt_for_reboot() {
    if ask_confirmation "Setup completed. Would you like to reboot now?" "n"; then
        log "${GREEN}Rebooting system...${NC}"
        sleep 2
        reboot
    else
        log "${YELLOW}Skipping reboot. Please remember to reboot later to apply all changes.${NC}"
    fi
}

# Main execution
main() {
    local mac_address=$1
    local gui_action=$2  # GUI
    local ip_address=$3  # Static IP
    local gateway_address=$4 # Gateway IP

    log "Starting Raspberry Pi setup..."

    # Create script directory if it doesn't exist
    mkdir -p "$SCRIPT_DIR"

    # Install dependencies
    install_dependencies
    
    # Check for Debian distribution upgrades
    # upgrade_debian_distro

    # Setup static IP
    setup_static_ip "$ip_address" "$gateway_address"

    # Configure MAC address
    setup_mac_address "$mac_address"

    # Setup GUI manager
    setup_gui_manager "$gui_action"

    # Setup crontab
    setup_crontab

    # Configure welcome script
    configure_welcome

    log "${GREEN}Setup completed successfully!${NC}"

    # Prompt for reboot
    prompt_for_reboot
}

# Call main function with all arguments
main "$@"