#!/bin/bash
###############################################################################
# Safe Package Installation Script
#
# Provides safe python -m pip installation with:
# - Virtual environment detection/creation
# - Dependency resolution
# - Version compatibility checking
# - Rollback on failure
# - Security scanning
#
# Usage:
#   ./scripts/safe_install.sh [package_name] [version]
#   ./scripts/safe_install.sh -r requirements.txt
#   ./scripts/safe_install.sh --all  # Install all project dependencies
###############################################################################

set -e  # Exit on error

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

# Script directory
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"

# Configuration
VENV_DIR="$PROJECT_ROOT/venv"
REQUIREMENTS_FILE="$PROJECT_ROOT/requirements.txt"
BACKUP_DIR="$PROJECT_ROOT/.pip_backups"

###############################################################################
# Helper Functions
###############################################################################

print_header() {
    echo -e "${BLUE}========================================${NC}"
    echo -e "${BLUE}$1${NC}"
    echo -e "${BLUE}========================================${NC}"
}

print_success() {
    echo -e "${GREEN}✓ $1${NC}"
}

print_error() {
    echo -e "${RED}✗ $1${NC}"
}

print_warning() {
    echo -e "${YELLOW}⚠ $1${NC}"
}

print_info() {
    echo -e "${BLUE}ℹ $1${NC}"
}

###############################################################################
# Virtual Environment Management
###############################################################################

check_venv() {
    if [[ "$VIRTUAL_ENV" != "" ]]; then
        print_success "Virtual environment active: $VIRTUAL_ENV"
        return 0
    else
        print_warning "No virtual environment active"
        return 1
    fi
}

create_venv() {
    print_header "Creating Virtual Environment"

    if [ -d "$VENV_DIR" ]; then
        print_info "Virtual environment already exists at: $VENV_DIR"
        read -p "Recreate it? (y/N): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            print_info "Removing old virtual environment..."
            rm -rf "$VENV_DIR"
        else
            return 0
        fi
    fi

    print_info "Creating new virtual environment..."
    python3 -m venv "$VENV_DIR"

    if [ $? -eq 0 ]; then
        print_success "Virtual environment created at: $VENV_DIR"
        print_info "Activate it with: source $VENV_DIR/bin/activate"
        return 0
    else
        print_error "Failed to create virtual environment"
        return 1
    fi
}

activate_venv() {
    if [ ! -d "$VENV_DIR" ]; then
        print_warning "Virtual environment not found at: $VENV_DIR"
        read -p "Create it now? (Y/n): " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Nn]$ ]]; then
            create_venv || return 1
        else
            return 1
        fi
    fi

    # Check if already activated
    if check_venv; then
        return 0
    fi

    print_info "Activating virtual environment..."
    source "$VENV_DIR/bin/activate"

    if check_venv; then
        print_success "Virtual environment activated"
        return 0
    else
        print_error "Failed to activate virtual environment"
        return 1
    fi
}

###############################################################################
# Pip Safety Checks
###############################################################################

check_pip() {
    print_header "Checking pip"

    # Use python -m pip to bypass any aliases
    if ! python -m pip --version &> /dev/null; then
        print_error "pip not found"
        print_info "Installing pip..."

        # Try to install pip
        if command -v python3 &> /dev/null; then
            python3 -m ensurepip --upgrade
        else
            print_error "Python 3 not found. Please install Python 3 first."
            return 1
        fi
    fi

    # Check pip version
    pip_version=$(python -m pip --version | awk '{print $2}')
    print_success "pip version: $pip_version"

    # Upgrade pip if needed
    print_info "Upgrading pip to latest version..."
    python -m pip install --upgrade pip

    return 0
}

create_backup() {
    print_header "Creating Backup"

    mkdir -p "$BACKUP_DIR"

    timestamp=$(date +%Y%m%d_%H%M%S)
    backup_file="$BACKUP_DIR/installed_packages_$timestamp.txt"

    print_info "Saving current package list to: $backup_file"
    python -m pip freeze > "$backup_file"

    print_success "Backup created"
    echo "$backup_file"
}

check_conflicts() {
    local package=$1

    print_header "Checking for Conflicts"

    # Use pip-conflict-checker if available
    if command -v pip-conflict-checker &> /dev/null; then
        pip-conflict-checker "$package"
    else
        # Basic check using pip
        print_info "Running dependency check..."
        python -m pip install --dry-run "$package" 2>&1 | grep -i "conflict\|incompatible" || true
    fi
}

###############################################################################
# Installation Functions
###############################################################################

safe_install_package() {
    local package=$1
    local version=$2

    # Construct package spec
    if [ -n "$version" ]; then
        package_spec="$package==$version"
    else
        package_spec="$package"
    fi

    print_header "Installing: $package_spec"

    # Create backup
    backup_file=$(create_backup)

    # Check for conflicts
    check_conflicts "$package_spec"

    # Ask for confirmation
    read -p "Proceed with installation? (Y/n): " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Nn]$ ]]; then
        print_warning "Installation cancelled"
        return 1
    fi

    # Install
    print_info "Installing $package_spec..."

    if python -m pip install "$package_spec"; then
        print_success "Successfully installed $package_spec"

        # Verify installation
        print_info "Verifying installation..."
        python -m pip show "$package" &> /dev/null

        if [ $? -eq 0 ]; then
            print_success "Package verified"
            return 0
        else
            print_error "Package installation verification failed"
            rollback "$backup_file"
            return 1
        fi
    else
        print_error "Installation failed"
        rollback "$backup_file"
        return 1
    fi
}

install_from_requirements() {
    local req_file=$1

    if [ ! -f "$req_file" ]; then
        print_error "Requirements file not found: $req_file"
        return 1
    fi

    print_header "Installing from: $req_file"

    # Create backup
    backup_file=$(create_backup)

    # Count packages
    package_count=$(grep -c -v '^#' "$req_file" | grep -c -v '^$' || echo 0)
    print_info "Found $package_count packages to install"

    # Show packages
    print_info "Packages to install:"
    grep -v '^#' "$req_file" | grep -v '^$' | sed 's/^/  - /'

    # Ask for confirmation
    read -p "Proceed with installation? (Y/n): " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Nn]$ ]]; then
        print_warning "Installation cancelled"
        return 1
    fi

    # Install
    print_info "Installing packages..."

    if python -m pip install -r "$req_file"; then
        print_success "All packages installed successfully"
        return 0
    else
        print_error "Installation failed"
        rollback "$backup_file"
        return 1
    fi
}

rollback() {
    local backup_file=$1

    print_header "Rolling Back"

    if [ ! -f "$backup_file" ]; then
        print_error "Backup file not found: $backup_file"
        return 1
    fi

    print_warning "Rolling back to previous state..."

    # Uninstall all current packages
    python -m pip freeze | xargs python -m pip uninstall -y

    # Reinstall from backup
    python -m pip install -r "$backup_file"

    if [ $? -eq 0 ]; then
        print_success "Rollback completed"
        return 0
    else
        print_error "Rollback failed"
        return 1
    fi
}

###############################################################################
# Dependency Resolution
###############################################################################

resolve_dependencies() {
    print_header "Resolving Dependencies"

    # Use pip-tools if available
    if command -v pip-compile &> /dev/null; then
        print_info "Using pip-tools to resolve dependencies..."
        pip-compile --upgrade --generate-hashes requirements.in
    else
        print_warning "pip-tools not installed. Install with: python -m pip install pip-tools"
        print_info "Basic dependency check:"
        python -m pip check
    fi
}

###############################################################################
# Security Scanning
###############################################################################

security_scan() {
    print_header "Security Scanning"

    # Use safety if available
    if command -v safety &> /dev/null; then
        print_info "Running safety check..."
        safety check --json | jq '.'
    else
        print_warning "safety not installed. Install with: python -m pip install safety"
        print_info "Skipping security scan"
    fi
}

###############################################################################
# Main Logic
###############################################################################

show_help() {
    cat << EOF
Safe Package Installation Script

Usage:
    $0 [OPTIONS] [PACKAGE] [VERSION]

Options:
    --help              Show this help message
    --create-venv       Create virtual environment only
    --activate-venv     Activate virtual environment (source this script)
    --check             Check pip and environment
    --all               Install all project dependencies
    -r FILE             Install from requirements file
    --resolve           Resolve dependencies (requires pip-tools)
    --scan              Run security scan (requires safety)
    --backup            Create backup of current packages
    --rollback FILE     Rollback to backup file

Examples:
    # Create and activate virtual environment
    $0 --create-venv
    source venv/bin/activate

    # Install a single package
    $0 requests

    # Install specific version
    $0 requests 2.28.0

    # Install from requirements.txt
    $0 -r requirements.txt
    $0 --all  # Uses default requirements.txt

    # Check environment
    $0 --check

    # Security scan
    $0 --scan

EOF
}

main() {
    # Parse arguments
    case "$1" in
        --help|-h)
            show_help
            exit 0
            ;;
        --create-venv)
            create_venv
            exit $?
            ;;
        --check)
            activate_venv || exit 1
            check_pip
            print_info "Python version: $(python --version)"
            print_info "Installed packages:"
            python -m pip list
            exit 0
            ;;
        --all)
            activate_venv || exit 1
            check_pip || exit 1
            install_from_requirements "$REQUIREMENTS_FILE"
            exit $?
            ;;
        -r)
            activate_venv || exit 1
            check_pip || exit 1
            install_from_requirements "$2"
            exit $?
            ;;
        --resolve)
            activate_venv || exit 1
            resolve_dependencies
            exit $?
            ;;
        --scan)
            activate_venv || exit 1
            security_scan
            exit $?
            ;;
        --backup)
            activate_venv || exit 1
            create_backup
            exit $?
            ;;
        --rollback)
            activate_venv || exit 1
            rollback "$2"
            exit $?
            ;;
        "")
            show_help
            exit 1
            ;;
        *)
            # Install package
            activate_venv || exit 1
            check_pip || exit 1
            safe_install_package "$1" "$2"
            exit $?
            ;;
    esac
}

# Run main
main "$@"