#!/bin/bash
# Package Installation Script with Vendor Repository Support
# Handles package installation across different distributions with inline vendor repository setup

set -euo pipefail

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

# Configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ROOT_DIR="$(dirname "$SCRIPT_DIR")"

# Check if we're running in a Packer environment (temporary location)
if [[ "$SCRIPT_DIR" == "/tmp" ]] && [[ -d "/tmp/linux_automation" ]]; then
    ROOT_DIR="/tmp/linux_automation"
elif [[ "$SCRIPT_DIR" == *"/tmp/linux_automation/scripts" ]]; then
    ROOT_DIR="/tmp/linux_automation"
fi

CONFIG_DIR="$ROOT_DIR/config/packages"

# Logging functions
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

log_step() {
    echo -e "${BLUE}[STEP]${NC} $1"
}

# Detect package manager
detect_package_manager() {
    if command -v apt-get &> /dev/null; then
        echo "apt"
    elif command -v dnf &> /dev/null; then
        echo "dnf"
    elif command -v yum &> /dev/null; then
        echo "yum"
    elif command -v zypper &> /dev/null; then
        echo "zypper"
    else
        log_error "No supported package manager found"
        return 1
    fi
}

# Detect distribution
detect_distribution() {
    if [[ -f /etc/os-release ]]; then
        source /etc/os-release
        echo "$ID"
    else
        log_error "Cannot detect distribution"
        return 1
    fi
}

# Parse YAML (basic parser for package lists)
parse_yaml_list() {
    local yaml_file=$1
    local package_manager=$2
    
    if [[ ! -f "$yaml_file" ]]; then
        log_error "Configuration file not found: $yaml_file"
        return 1
    fi
    
    # Extract packages for the specified package manager
    awk "
    /^${package_manager}_packages:/ { in_section = 1; next }
    /^[a-zA-Z]/ && in_section { in_section = 0 }
    in_section && /^  - / { gsub(/^  - /, \"\"); print }
    " "$yaml_file"
}

#=============================================================================
# Vendor Repository Setup Functions
#=============================================================================

# Setup vendor repositories for APT (Ubuntu/Debian)
setup_vendor_repos_apt() {
    log_step "Setting up vendor repositories for APT"
    
    # Detect distribution
    local distribution
    distribution=$(detect_distribution)
    
    # Install prerequisites
    log_info "Installing repository prerequisites..."
    apt-get update -qq
    apt-get install -y ca-certificates curl gnupg lsb-release apt-transport-https software-properties-common
    
    # Create keyrings directory (modern approach)
    mkdir -p /etc/apt/keyrings

    # Note: Docker repository setup removed - Docker is now installed via Ansible playbook
    # See ansible/playbooks/docker-compose.yml for Docker installation

    # HashiCorp Repository  
    log_info "Setting up HashiCorp repository..."
    curl -fsSL https://apt.releases.hashicorp.com/gpg | gpg --dearmor -o /usr/share/keyrings/hashicorp-archive-keyring.gpg
    echo "deb [arch=amd64 signed-by=/usr/share/keyrings/hashicorp-archive-keyring.gpg] https://apt.releases.hashicorp.com $(lsb_release -cs) main" | tee /etc/apt/sources.list.d/hashicorp.list > /dev/null
    
    # Kubernetes Repository
    log_info "Setting up Kubernetes repository..."
    curl -fsSL https://pkgs.k8s.io/core:/stable:/v1.28/deb/Release.key | gpg --dearmor -o /etc/apt/keyrings/kubernetes-apt-keyring.gpg
    echo "deb [signed-by=/etc/apt/keyrings/kubernetes-apt-keyring.gpg] https://pkgs.k8s.io/core:/stable:/v1.28/deb/ /" | tee /etc/apt/sources.list.d/kubernetes.list > /dev/null
    
    # Microsoft Repository (Azure CLI)
    log_info "Setting up Microsoft repository..."
    curl -fsSL https://packages.microsoft.com/keys/microsoft.asc | gpg --dearmor -o /etc/apt/keyrings/microsoft.gpg
    echo "deb [arch=amd64,arm64,armhf signed-by=/etc/apt/keyrings/microsoft.gpg] https://packages.microsoft.com/repos/azure-cli/ $(lsb_release -cs) main" | tee /etc/apt/sources.list.d/azure-cli.list > /dev/null
    
    # Update package lists with new repositories
    log_info "Updating package lists with new repositories..."
    apt-get update -qq
    
    log_info "Vendor repositories configured successfully"
}

# Setup vendor repositories for DNF/YUM (RHEL-family)
setup_vendor_repos_dnf() {
    local distribution=$1
    log_step "Setting up vendor repositories for DNF/YUM"
    
    # Determine if we're using dnf or yum
    local pkg_cmd="dnf"
    if ! command -v dnf &> /dev/null && command -v yum &> /dev/null; then
        pkg_cmd="yum"
    fi
    
    # EPEL Repository (Essential for RHEL-family, but NOT Fedora)
    if [[ "$distribution" != "fedora" ]]; then
        log_info "Installing EPEL repository..."
        if ! rpm -q epel-release &> /dev/null; then
            $pkg_cmd install -y epel-release
        else
            log_info "EPEL already installed"
        fi
    else
        log_info "Skipping EPEL (not needed on Fedora)"
    fi

    # Note: Docker repository setup removed - Docker is now installed via Ansible playbook
    # See ansible/playbooks/docker-compose.yml for Docker installation

    # HashiCorp Repository
    log_info "Setting up HashiCorp repository..."
    if [[ ! -f /etc/yum.repos.d/hashicorp.repo ]]; then
        if [[ "$distribution" == "fedora" ]]; then
            # Fedora uses Fedora-specific HashiCorp repo
            curl -fsSL https://rpm.releases.hashicorp.com/fedora/hashicorp.repo -o /etc/yum.repos.d/hashicorp.repo
        else
            # RHEL-family
            $pkg_cmd config-manager --add-repo https://rpm.releases.hashicorp.com/RHEL/hashicorp.repo || {
                curl -fsSL https://rpm.releases.hashicorp.com/RHEL/hashicorp.repo -o /etc/yum.repos.d/hashicorp.repo
            }
        fi
    fi
    
    # Kubernetes Repository
    log_info "Setting up Kubernetes repository..."
    if [[ ! -f /etc/yum.repos.d/kubernetes.repo ]]; then
        cat <<EOF | tee /etc/yum.repos.d/kubernetes.repo > /dev/null
[kubernetes]
name=Kubernetes
baseurl=https://pkgs.k8s.io/core:/stable:/v1.28/rpm/
enabled=1
gpgcheck=1
gpgkey=https://pkgs.k8s.io/core:/stable:/v1.28/rpm/repodata/repomd.xml.key
EOF
    fi
    
    # Microsoft Repository (Azure CLI)
    log_info "Setting up Microsoft repository..."
    if [[ ! -f /etc/yum.repos.d/azure-cli.repo ]]; then
        rpm --import https://packages.microsoft.com/keys/microsoft.asc || true
        cat <<EOF | tee /etc/yum.repos.d/azure-cli.repo > /dev/null
[azure-cli]
name=Azure CLI
baseurl=https://packages.microsoft.com/yumrepos/azure-cli
enabled=1
gpgcheck=1
gpgkey=https://packages.microsoft.com/keys/microsoft.asc
EOF
    fi
    
    log_info "Vendor repositories configured successfully"
}

# Setup vendor repositories for zypper (openSUSE)
setup_vendor_repos_zypper() {
    log_step "Setting up vendor repositories for zypper (openSUSE)"

    # Note: Many vendor packages are available through Open Build Service (OBS)
    # We'll use OBS repositories where appropriate for better compatibility

    # Note: Docker repository setup removed - Docker is now installed via Ansible playbook
    # See ansible/playbooks/docker-compose.yml for Docker installation

    # HashiCorp - Install via direct download (no official openSUSE repo)
    log_info "Note: HashiCorp tools will be installed from official binaries"
    # We'll handle these differently in package installation
    
    # Kubernetes - Use official Kubernetes repo
    log_info "Setting up Kubernetes repository..."
    if ! zypper lr | grep -q "kubernetes"; then
        cat <<EOF | tee /etc/zypp/repos.d/kubernetes.repo > /dev/null
[kubernetes]
name=Kubernetes
baseurl=https://pkgs.k8s.io/core:/stable:/v1.28/rpm/
enabled=1
gpgcheck=1
gpgkey=https://pkgs.k8s.io/core:/stable:/v1.28/rpm/repodata/repomd.xml.key
EOF
    fi
    
    # Microsoft Azure CLI - Use Microsoft repo (works with openSUSE)
    log_info "Setting up Microsoft repository..."
    if ! zypper lr | grep -q "azure-cli"; then
        rpm --import https://packages.microsoft.com/keys/microsoft.asc || true
        
        cat <<EOF | tee /etc/zypp/repos.d/azure-cli.repo > /dev/null
[azure-cli]
name=Azure CLI
baseurl=https://packages.microsoft.com/yumrepos/azure-cli
enabled=1
gpgcheck=1
gpgkey=https://packages.microsoft.com/keys/microsoft.asc
EOF
    fi
    
    # Refresh repositories
    log_info "Refreshing zypper repositories..."
    zypper --gpg-auto-import-keys refresh || log_warn "Some repositories may have failed to refresh"
    
    log_info "Vendor repositories configured successfully"
}

# Main vendor repository setup dispatcher
setup_vendor_repos() {
    local package_manager=$1
    local distribution=$2
    
    log_step "Setting up vendor repositories for extended packages"
    
    case "$package_manager" in
        "apt")
            setup_vendor_repos_apt
            ;;
        "dnf"|"yum")
            setup_vendor_repos_dnf "$distribution"
            ;;
        "zypper")
            setup_vendor_repos_zypper
            ;;
        *)
            log_warn "Vendor repository setup not supported for $package_manager"
            ;;
    esac
}

#=============================================================================
# Package Installation Functions
#=============================================================================

# Map package names for distribution compatibility
map_package_name() {
    local package=$1
    local distribution=$2
    
    # Distribution-specific package mappings
    case "$distribution" in
        "rocky"|"almalinux"|"centos"|"rhel")
            case "$package" in
                "redhat-lsb-core")
                    # Skip this package on Rocky/Alma as it's not available/needed
                    echo ""
                    return
                    ;;
                "python3-virtualenv")
                    # Rocky 9 uses different package name
                    echo "python3-pip"  # pip includes venv functionality
                    return
                    ;;
                "azure-cli")
                    # Azure CLI not available/compatible with Rocky 9 repos
                    # Skip it - can be installed via pip if needed
                    echo ""
                    return
                    ;;
            esac
            ;;
    esac
    
    # Return original package name if no mapping needed
    echo "$package"
}

# Install packages
install_packages() {
    local package_manager=$1
    local config_file=$2
    local distribution=$3
    
    log_step "Installing packages using $package_manager"
    
    # Get package list for this package manager
    local packages
    packages=$(parse_yaml_list "$config_file" "${package_manager}")
    
    if [[ -z "$packages" ]]; then
        log_warn "No packages found for $package_manager in $config_file"
        return 0
    fi
    
    # Convert to array and apply distribution-specific mappings
    local package_array=()
    while IFS= read -r package; do
        if [[ -n "$package" ]]; then
            mapped_package=$(map_package_name "$package" "$distribution")
            if [[ -n "$mapped_package" ]]; then
                package_array+=("$mapped_package")
            else
                log_info "Skipping package '$package' (not available for $distribution)"
            fi
        fi
    done <<< "$packages"
    
    if [[ ${#package_array[@]} -eq 0 ]]; then
        log_warn "No valid packages to install"
        return 0
    fi
    
    log_info "Installing ${#package_array[@]} packages..."
    
    # Install packages individually to handle missing packages gracefully
    local failed_packages=()
    local success_count=0
    
    for package in "${package_array[@]}"; do
        # Check for custom binary installations
        case "$package" in
            "helm")
                if install_helm; then
                    ((success_count++)) || true
                else
                    failed_packages+=("$package")
                    log_warn "Failed to install custom binary: $package"
                fi
                continue
                ;;
            "k9s")
                if install_k9s; then
                    ((success_count++)) || true
                else
                    failed_packages+=("$package")
                    log_warn "Failed to install custom binary: $package"
                fi
                continue
                ;;
            "awscli")
                if install_awscli; then
                    ((success_count++)) || true
                else
                    failed_packages+=("$package")
                    log_warn "Failed to install custom binary: $package"
                fi
                continue
                ;;
        esac
        
        # Standard package manager installation
        case "$package_manager" in
            "apt")
                if apt-get install -y "$package" 2>&1 | tee -a /tmp/package-install.log; then
                    ((success_count++)) || true
                else
                    failed_packages+=("$package")
                    log_warn "Failed to install package: $package"
                fi
                ;;
            "dnf")
                if dnf install -y "$package" 2>&1 | tee -a /tmp/package-install.log; then
                    ((success_count++)) || true
                else
                    failed_packages+=("$package")
                    log_warn "Failed to install package: $package"
                fi
                ;;
            "yum")
                if yum install -y "$package" 2>&1 | tee -a /tmp/package-install.log; then
                    ((success_count++)) || true
                else
                    failed_packages+=("$package")
                    log_warn "Failed to install package: $package"
                fi
                ;;
            "zypper")
                if zypper install -y "$package" 2>&1 | tee -a /tmp/package-install.log; then
                    ((success_count++)) || true
                else
                    failed_packages+=("$package")
                    log_warn "Failed to install package: $package"
                fi
                ;;
            *)
                log_error "Package installation not implemented for $package_manager"
                return 1
                ;;
        esac
    done
    
    # Report results
    log_info "Successfully installed $success_count packages"
    
    if [[ ${#failed_packages[@]} -gt 0 ]]; then
        log_warn "Failed to install ${#failed_packages[@]} packages: ${failed_packages[*]}"
        log_info "Continuing with remaining installation steps..."
    fi
}

# Map service names for distribution compatibility
map_service_name() {
    local service=$1
    local distribution=$2
    local package_manager=$3
    
    # Distribution-specific service name mappings
    case "$distribution" in
        "rocky"|"almalinux"|"centos"|"rhel"|"fedora")
            case "$service" in
                "ssh")
                    echo "sshd"
                    return
                    ;;
                "cron")
                    echo "crond"
                    return
                    ;;
            esac
            ;;
        "ubuntu"|"debian")
            # Ubuntu/Debian use standard names, but chrony service name differs
            case "$service" in
                "chronyd")
                    # On Ubuntu 24.04+, chrony service is named 'chrony'
                    if systemctl list-unit-files | grep -q "^chrony.service"; then
                        echo "chrony"
                    else
                        echo "chronyd"
                    fi
                    return
                    ;;
            esac
            ;;
    esac
    
    # Return original service name if no mapping needed
    echo "$service"
}

# Enable services
enable_services() {
    local config_file=$1
    local distribution=$2
    local package_manager=$3
    
    log_step "Enabling services"
    
    # Get service list
    local services
    services=$(awk '/^services:/ { in_section = 1; next } /^[a-zA-Z]/ && in_section { in_section = 0 } in_section && /^  - / { gsub(/^  - /, ""); print }' "$config_file")
    
    if [[ -z "$services" ]]; then
        log_info "No services to enable"
        return 0
    fi
    
    while IFS= read -r service; do
        if [[ -n "$service" ]]; then
            # Map service name for distribution
            local mapped_service
            mapped_service=$(map_service_name "$service" "$distribution" "$package_manager")
            
            log_info "Enabling service: $mapped_service"
            if systemctl enable "$mapped_service" 2>/dev/null; then
                log_info "✓ Enabled: $mapped_service"
            else
                log_warn "Failed to enable service: $mapped_service (may not exist or already enabled)"
            fi
        fi
    done <<< "$services"
}

#=============================================================================
# Custom Binary Installation Functions
#=============================================================================

# Install Helm - Kubernetes package manager
install_helm() {
    log_step "Installing Helm from GitHub releases"
    
    # Detect architecture
    local arch
    arch=$(uname -m)
    case "$arch" in
        x86_64)
            arch="amd64"
            ;;
        aarch64|arm64)
            arch="arm64"
            ;;
        *)
            log_error "Unsupported architecture for Helm: $arch"
            return 1
            ;;
    esac
    
    # Get latest version from GitHub API
    local latest_version
    latest_version=$(curl -fsSL https://api.github.com/repos/helm/helm/releases/latest | grep '"tag_name"' | sed -E 's/.*"v([^"]+)".*/\1/')
    
    if [[ -z "$latest_version" ]]; then
        log_error "Failed to determine latest Helm version"
        return 1
    fi
    
    log_info "Installing Helm version $latest_version for $arch"
    
    # Download URL
    local download_url="https://get.helm.sh/helm-v${latest_version}-linux-${arch}.tar.gz"
    local temp_dir
    temp_dir=$(mktemp -d)
    
    # Download and extract
    if curl -fsSL "$download_url" -o "$temp_dir/helm.tar.gz"; then
        tar -xzf "$temp_dir/helm.tar.gz" -C "$temp_dir"
        
        # Install binary
        if [[ -f "$temp_dir/linux-${arch}/helm" ]]; then
            mv "$temp_dir/linux-${arch}/helm" /usr/local/bin/helm
            chmod +x /usr/local/bin/helm
            
            # Verify installation
            if /usr/local/bin/helm version --short; then
                log_info "✓ Helm installed successfully"
            else
                log_error "Helm binary installed but failed verification"
                rm -rf "$temp_dir"
                return 1
            fi
        else
            log_error "Helm binary not found in downloaded archive"
            rm -rf "$temp_dir"
            return 1
        fi
    else
        log_error "Failed to download Helm from $download_url"
        rm -rf "$temp_dir"
        return 1
    fi
    
    # Cleanup
    rm -rf "$temp_dir"
    return 0
}

# Install k9s - Kubernetes CLI UI
install_k9s() {
    log_step "Installing k9s from GitHub releases"
    
    # Detect architecture
    local arch
    arch=$(uname -m)
    case "$arch" in
        x86_64)
            arch="amd64"
            ;;
        aarch64|arm64)
            arch="arm64"
            ;;
        *)
            log_error "Unsupported architecture for k9s: $arch"
            return 1
            ;;
    esac
    
    # Get latest version from GitHub API
    local latest_version
    latest_version=$(curl -fsSL https://api.github.com/repos/derailed/k9s/releases/latest | grep '"tag_name"' | sed -E 's/.*"v([^"]+)".*/\1/')
    
    if [[ -z "$latest_version" ]]; then
        log_error "Failed to determine latest k9s version"
        return 1
    fi
    
    log_info "Installing k9s version $latest_version for $arch"
    
    # Download URL
    local download_url="https://github.com/derailed/k9s/releases/download/v${latest_version}/k9s_Linux_${arch}.tar.gz"
    local temp_dir
    temp_dir=$(mktemp -d)
    
    # Download and extract
    if curl -fsSL "$download_url" -o "$temp_dir/k9s.tar.gz"; then
        tar -xzf "$temp_dir/k9s.tar.gz" -C "$temp_dir"
        
        # Install binary
        if [[ -f "$temp_dir/k9s" ]]; then
            mv "$temp_dir/k9s" /usr/local/bin/k9s
            chmod +x /usr/local/bin/k9s
            
            # Verify installation
            if /usr/local/bin/k9s version --short 2>&1; then
                log_info "✓ k9s installed successfully"
            else
                log_error "k9s binary installed but failed verification"
                rm -rf "$temp_dir"
                return 1
            fi
        else
            log_error "k9s binary not found in downloaded archive"
            rm -rf "$temp_dir"
            return 1
        fi
    else
        log_error "Failed to download k9s from $download_url"
        rm -rf "$temp_dir"
        return 1
    fi
    
    # Cleanup
    rm -rf "$temp_dir"
    return 0
}

# Install AWS CLI v2
install_awscli() {
    log_step "Installing AWS CLI v2 from AWS releases"
    
    # Detect architecture
    local arch
    arch=$(uname -m)
    case "$arch" in
        x86_64)
            arch="x86_64"
            ;;
        aarch64|arm64)
            arch="aarch64"
            ;;
        *)
            log_error "Unsupported architecture for AWS CLI: $arch"
            return 1
            ;;
    esac
    
    log_info "Installing AWS CLI v2 for $arch"
    
    # Download URL for AWS CLI v2
    local download_url="https://awscli.amazonaws.com/awscli-exe-linux-${arch}.zip"
    local temp_dir
    temp_dir=$(mktemp -d)
    
    # Download and extract
    if curl -fsSL "$download_url" -o "$temp_dir/awscliv2.zip"; then
        unzip -q "$temp_dir/awscliv2.zip" -d "$temp_dir"
        
        # Install AWS CLI
        if [[ -d "$temp_dir/aws" && -f "$temp_dir/aws/install" ]]; then
            # Run the AWS installer
            "$temp_dir/aws/install" --update
            
            # Verify installation
            if aws --version 2>&1 | grep -q "aws-cli"; then
                log_info "✓ AWS CLI installed successfully"
            else
                log_error "AWS CLI installation failed verification"
                rm -rf "$temp_dir"
                return 1
            fi
        else
            log_error "AWS CLI installer not found in downloaded archive"
            rm -rf "$temp_dir"
            return 1
        fi
    else
        log_error "Failed to download AWS CLI from $download_url"
        rm -rf "$temp_dir"
        return 1
    fi
    
    # Cleanup
    rm -rf "$temp_dir"
    return 0
}

#=============================================================================

# Configure system PATH for custom binaries
configure_system_path() {
    log_step "Configuring system PATH for custom binaries"

    # Create profile.d script to ensure /usr/local/bin is in PATH
    cat > /etc/profile.d/custom-path.sh <<'EOF'
# Add /usr/local/bin to PATH for custom installed binaries
# (helm, k9s, aws-cli, etc.)
if ! echo "$PATH" | grep -q "/usr/local/bin"; then
    export PATH="/usr/local/bin:$PATH"
fi
EOF

    chmod +x /etc/profile.d/custom-path.sh

    log_info "✓ System PATH configured to include /usr/local/bin"
    log_info "Custom binaries (helm, k9s, aws) will be available after login"
}

# Main Installation Function
#=============================================================================

install_package_profile() {
    local profile=$1
    local config_file="$CONFIG_DIR/${profile}-packages.yml"

    if [[ ! -f "$config_file" ]]; then
        log_error "Package configuration file not found: $config_file"
        return 1
    fi
    
    log_step "Installing $profile package profile"
    
    # Detect system
    local package_manager
    package_manager=$(detect_package_manager)
    local distribution
    distribution=$(detect_distribution)
    
    log_info "Detected: $distribution with $package_manager package manager"
    
    # Setup vendor repositories (only for extended and security profiles)
    if [[ "$profile" == "extended" || "$profile" == "security" ]]; then
        setup_vendor_repos "$package_manager" "$distribution"
    fi
    
    # Install packages
    install_packages "$package_manager" "$config_file" "$distribution"

    # Enable services
    enable_services "$config_file" "$distribution" "$package_manager"

    # Configure PATH for custom binaries (extended/security profiles with custom binaries)
    if [[ "$profile" == "extended" || "$profile" == "security" ]]; then
        configure_system_path
    fi

    log_info "$profile package profile installation completed"
}

#=============================================================================
# Script Entry Point
#=============================================================================

# Usage
usage() {
    cat <<EOF
Usage: $0 PROFILE

Install packages based on predefined profiles with vendor repository support.

Profiles:
  base        Essential packages for all templates (default repos only)
  extended    Base packages + DevOps tools (includes vendor repos)
  security    Extended packages + security hardening (includes vendor repos)

Vendor Repositories (extended/security only):
  - Docker (docker-ce, docker-compose, buildx)
  - HashiCorp (terraform, packer, vault, consul)
  - Kubernetes (kubectl)
  - Microsoft (azure-cli)
  - EPEL (RHEL-family only - ansible and utilities)

Examples:
  $0 base
  $0 extended
  $0 security

EOF
}

# Main
main() {
    if [[ $# -lt 1 ]]; then
        usage
        exit 1
    fi
    
    local profile=$1
    
    case "$profile" in
        "base"|"extended"|"security")
            install_package_profile "$profile"
            ;;
        *)
            log_error "Unknown profile: $profile"
            usage
            exit 1
            ;;
    esac
}

# Run main function if script is executed directly
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi
