#!/bin/bash

# PostgreSQL Setup Script for Hero Models Example
# This script checks for PostgreSQL installation and sets up the required database configuration

set -e  # Exit on any error

echo "🚀 Hero Models PostgreSQL Example Setup"
echo "========================================"

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

# Function to print colored output
print_status() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

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

# Detect operating system
detect_os() {
    if [[ "$OSTYPE" == "darwin"* ]]; then
        echo "macos"
    elif [[ "$OSTYPE" == "linux-gnu"* ]]; then
        if command -v apt-get &> /dev/null; then
            echo "ubuntu"
        elif command -v yum &> /dev/null; then
            echo "centos"
        else
            echo "linux"
        fi
    else
        echo "unknown"
    fi
}

# Check if PostgreSQL is installed
check_postgres_installed() {
    # Check if PostgreSQL is in current PATH
    if command -v postgres &> /dev/null || command -v psql &> /dev/null; then
        return 0
    fi
    
    # Check macOS Homebrew installation location
    if [[ $(detect_os) == "macos" ]]; then
        if [[ -f "/opt/homebrew/opt/postgresql@15/bin/postgres" ]] || [[ -f "/opt/homebrew/opt/postgresql@15/bin/psql" ]]; then
            return 0
        fi
        # Also check Intel Mac location
        if [[ -f "/usr/local/opt/postgresql@15/bin/postgres" ]] || [[ -f "/usr/local/opt/postgresql@15/bin/psql" ]]; then
            return 0
        fi
    fi
    
    return 1
}

# Install PostgreSQL based on OS
install_postgres() {
    local os=$(detect_os)
    
    case $os in
        "macos")
            print_status "Installing PostgreSQL on macOS using Homebrew..."
            if ! command -v brew &> /dev/null; then
                print_error "Homebrew is not installed. Please install Homebrew first:"
                echo "  /bin/bash -c \"\$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/HEAD/install.sh)\""
                exit 1
            fi
            brew install postgresql@15
            print_success "PostgreSQL installed successfully"
            ;;
        "ubuntu")
            print_status "Installing PostgreSQL on Ubuntu/Debian..."
            sudo apt-get update
            sudo apt-get install -y postgresql postgresql-contrib
            print_success "PostgreSQL installed successfully"
            ;;
        "centos")
            print_status "Installing PostgreSQL on CentOS/RHEL..."
            sudo yum install -y postgresql-server postgresql-contrib
            sudo postgresql-setup initdb
            print_success "PostgreSQL installed successfully"
            ;;
        *)
            print_error "Unsupported operating system: $os"
            print_error "Please install PostgreSQL manually and run this script again"
            exit 1
            ;;
    esac
}

# Start PostgreSQL service
start_postgres() {
    local os=$(detect_os)
    
    case $os in
        "macos")
            print_status "Starting PostgreSQL service on macOS..."
            brew services start postgresql@15
            # Add PostgreSQL to PATH for this session
            export PATH="/opt/homebrew/opt/postgresql@15/bin:$PATH"
            print_success "PostgreSQL service started"
            ;;
        "ubuntu")
            print_status "Starting PostgreSQL service on Ubuntu/Debian..."
            sudo systemctl start postgresql
            sudo systemctl enable postgresql
            print_success "PostgreSQL service started and enabled"
            ;;
        "centos")
            print_status "Starting PostgreSQL service on CentOS/RHEL..."
            sudo systemctl start postgresql
            sudo systemctl enable postgresql
            print_success "PostgreSQL service started and enabled"
            ;;
    esac
}

# Check if PostgreSQL is running
check_postgres_running() {
    local os=$(detect_os)
    
    # Ensure PostgreSQL binaries are in PATH for macOS
    if [[ $os == "macos" ]]; then
        export PATH="/opt/homebrew/opt/postgresql@15/bin:$PATH"
    fi
    
    if pg_isready -h localhost -p 5432 &> /dev/null; then
        return 0
    else
        return 1
    fi
}

# Setup database user and password
setup_database() {
    print_status "Setting up database user and password..."
    
    local os=$(detect_os)
    
    # Ensure PostgreSQL binaries are in PATH for macOS
    if [[ $os == "macos" ]]; then
        export PATH="/opt/homebrew/opt/postgresql@15/bin:$PATH"
    fi
    
    # Create postgres user if it doesn't exist (mainly for macOS)
    if ! psql -U postgres -c '\q' &> /dev/null; then
        print_status "Creating postgres user..."
        if [[ $os == "macos" ]]; then
            createuser -s postgres 2>/dev/null || true
        else
            sudo -u postgres createuser -s postgres 2>/dev/null || true
        fi
    fi
    
    # Set password for postgres user
    print_status "Setting password for postgres user..."
    if [[ $os == "macos" ]]; then
        psql -U postgres -c "ALTER USER postgres PASSWORD 'test123';" 2>/dev/null || {
            print_warning "Could not set password directly. Trying alternative method..."
            createdb -U postgres postgres 2>/dev/null || true
            psql -U postgres -d postgres -c "ALTER USER postgres PASSWORD 'test123';" || {
                print_error "Failed to set password. You may need to set it manually:"
                echo "  psql -U postgres -c \"ALTER USER postgres PASSWORD 'test123';\""
                return 1
            }
        }
    else
        sudo -u postgres psql -c "ALTER USER postgres PASSWORD 'test123';" || {
            print_error "Failed to set password. You may need to set it manually:"
            echo "  sudo -u postgres psql -c \"ALTER USER postgres PASSWORD 'test123';\""
            return 1
        }
    fi
    
    print_success "Database user configured successfully"
}

# Test database connection
test_connection() {
    print_status "Testing database connection..."
    
    local os=$(detect_os)
    if [[ $os == "macos" ]]; then
        export PATH="/opt/homebrew/opt/postgresql@15/bin:$PATH"
    fi
    
    if PGPASSWORD=test123 psql -h localhost -p 5432 -U postgres -d postgres -c "SELECT version();" &> /dev/null; then
        print_success "Database connection test successful!"
        return 0
    else
        print_error "Database connection test failed"
        return 1
    fi
}

# Add PATH export to shell profile (macOS only)
setup_path_macos() {
    if [[ $(detect_os) == "macos" ]]; then
        local shell_profile=""
        if [[ $SHELL == *"zsh"* ]]; then
            shell_profile="$HOME/.zshrc"
        elif [[ $SHELL == *"bash"* ]]; then
            shell_profile="$HOME/.bash_profile"
        fi
        
        if [[ -n $shell_profile ]]; then
            local path_export='export PATH="/opt/homebrew/opt/postgresql@15/bin:$PATH"'
            if ! grep -q "$path_export" "$shell_profile" 2>/dev/null; then
                print_status "Adding PostgreSQL to PATH in $shell_profile..."
                echo "" >> "$shell_profile"
                echo "# PostgreSQL" >> "$shell_profile"
                echo "$path_export" >> "$shell_profile"
                print_success "PostgreSQL added to PATH. Restart your terminal or run: source $shell_profile"
            fi
        fi
    fi
}

# Prompt user for installation
prompt_install() {
    echo ""
    print_warning "PostgreSQL is not installed on your system."
    echo ""
    print_status "The Hero Models example requires PostgreSQL to be installed and configured."
    echo ""
    echo "Options:"
    echo "  1. Let this script install PostgreSQL automatically"
    echo "  2. Install PostgreSQL manually and run this script again"
    echo ""
    read -p "Would you like this script to install PostgreSQL for you? (y/n): " -n 1 -r
    echo ""
    
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        return 0  # User wants automatic installation
    else
        return 1  # User wants manual installation
    fi
}

# Main setup process
main() {
    echo ""
    
    # Check if PostgreSQL is already installed
    if check_postgres_installed; then
        print_success "PostgreSQL is already installed"
    else
        if prompt_install; then
            print_status "Installing PostgreSQL..."
            install_postgres
        else
            print_status "Please install PostgreSQL manually and run this script again."
            echo ""
            print_status "Manual installation instructions:"
            local os=$(detect_os)
            case $os in
                "macos")
                    echo "  brew install postgresql@15"
                    echo "  brew services start postgresql@15"
                    ;;
                "ubuntu")
                    echo "  sudo apt-get update"
                    echo "  sudo apt-get install postgresql postgresql-contrib"
                    echo "  sudo systemctl start postgresql"
                    ;;
                "centos")
                    echo "  sudo yum install postgresql-server postgresql-contrib"
                    echo "  sudo postgresql-setup initdb"
                    echo "  sudo systemctl start postgresql"
                    ;;
                *)
                    echo "  Please install PostgreSQL using your system's package manager"
                    ;;
            esac
            echo ""
            exit 0
        fi
    fi
    
    # Start PostgreSQL service
    print_status "Checking PostgreSQL service status..."
    if ! check_postgres_running; then
        print_status "PostgreSQL is not running. Starting service..."
        start_postgres
        sleep 2  # Give service time to start
    else
        print_success "PostgreSQL service is already running"
    fi
    
    # Setup database user and password
    setup_database
    
    # Test the connection
    if test_connection; then
        print_success "Setup completed successfully!"
    else
        print_error "Setup completed but connection test failed"
        print_error "You may need to manually configure the database"
        exit 1
    fi
    
    # Setup PATH for macOS
    setup_path_macos
    
    echo ""
    print_success "🎉 PostgreSQL setup complete!"
    echo ""
    print_status "You can now run the example with:"
    echo "  cargo run --example postgres_model_example"
    echo ""
    print_status "Database configuration:"
    echo "  Host: localhost"
    echo "  Port: 5432"
    echo "  Username: postgres"
    echo "  Password: test123"
    echo ""
}

# Run main function
main "$@"