#!/bin/bash
###############################################################################
# GigLez Database Setup Script
#
# Creates PostgreSQL database, user, and schema for GigLez platform
#
# Usage:
#   sudo ./scripts/setup_database.sh
#
# This script must be run with sudo to access PostgreSQL
###############################################################################

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
DB_NAME="giglez"
DB_USER="giglez_user"
PASSWORD_FILE="$PROJECT_ROOT/.db_password"
SCHEMA_FILE="$PROJECT_ROOT/scripts/create_schema.sql"

###############################################################################
# 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}"
}

###############################################################################
# Checks
###############################################################################

check_sudo() {
    if [ "$EUID" -ne 0 ]; then
        print_error "This script must be run with sudo"
        echo "Usage: sudo $0"
        exit 1
    fi
    print_success "Running with sudo privileges"
}

check_postgresql() {
    print_header "Checking PostgreSQL Installation"

    if ! command -v psql &> /dev/null; then
        print_error "PostgreSQL is not installed"
        print_info "Install with: sudo apt-get install postgresql postgresql-contrib postgis"
        exit 1
    fi

    print_success "PostgreSQL is installed"

    # Check if PostgreSQL service is running
    if ! sudo systemctl is-active --quiet postgresql; then
        print_warning "PostgreSQL service is not running"
        print_info "Starting PostgreSQL service..."
        sudo systemctl start postgresql
        sleep 2
    fi

    print_success "PostgreSQL service is running"
}

check_password_file() {
    print_header "Checking Password File"

    if [ ! -f "$PASSWORD_FILE" ]; then
        print_error "Password file not found: $PASSWORD_FILE"
        exit 1
    fi

    DB_PASSWORD=$(cat "$PASSWORD_FILE" | tr -d '\n\r')

    if [ -z "$DB_PASSWORD" ]; then
        print_error "Password file is empty"
        exit 1
    fi

    print_success "Password loaded from file"
}

###############################################################################
# Database Setup
###############################################################################

create_database() {
    print_header "Creating Database"

    # Check if database already exists
    if sudo -u postgres psql -lqt | cut -d \| -f 1 | grep -qw "$DB_NAME"; then
        print_warning "Database '$DB_NAME' already exists"
        read -p "Drop and recreate? (y/N): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            print_info "Dropping existing database..."
            sudo -u postgres psql -c "DROP DATABASE IF EXISTS $DB_NAME;"
            print_success "Existing database dropped"
        else
            print_info "Keeping existing database"
            return 0
        fi
    fi

    # Create database
    print_info "Creating database '$DB_NAME'..."
    sudo -u postgres psql -c "CREATE DATABASE $DB_NAME;"
    print_success "Database created"
}

create_user() {
    print_header "Creating Database User"

    # Check if user already exists
    if sudo -u postgres psql -tAc "SELECT 1 FROM pg_roles WHERE rolname='$DB_USER'" | grep -q 1; then
        print_warning "User '$DB_USER' already exists"
        read -p "Drop and recreate? (y/N): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            print_info "Dropping existing user..."
            sudo -u postgres psql -c "DROP USER IF EXISTS $DB_USER;"
            print_success "Existing user dropped"
        else
            print_info "Updating password for existing user..."
            sudo -u postgres psql -c "ALTER USER $DB_USER WITH PASSWORD '$DB_PASSWORD';"
            print_success "Password updated"
            return 0
        fi
    fi

    # Create user with password
    print_info "Creating user '$DB_USER'..."
    sudo -u postgres psql -c "CREATE USER $DB_USER WITH PASSWORD '$DB_PASSWORD';"
    print_success "User created"
}

grant_privileges() {
    print_header "Granting Privileges"

    print_info "Granting all privileges on database '$DB_NAME' to '$DB_USER'..."
    sudo -u postgres psql -c "GRANT ALL PRIVILEGES ON DATABASE $DB_NAME TO $DB_USER;"

    # Grant schema privileges
    sudo -u postgres psql -d $DB_NAME -c "GRANT ALL ON SCHEMA public TO $DB_USER;"
    sudo -u postgres psql -d $DB_NAME -c "GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA public TO $DB_USER;"
    sudo -u postgres psql -d $DB_NAME -c "GRANT ALL PRIVILEGES ON ALL SEQUENCES IN SCHEMA public TO $DB_USER;"

    # Set default privileges for future objects
    sudo -u postgres psql -d $DB_NAME -c "ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON TABLES TO $DB_USER;"
    sudo -u postgres psql -d $DB_NAME -c "ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT ALL ON SEQUENCES TO $DB_USER;"

    print_success "Privileges granted"
}

enable_postgis() {
    print_header "Enabling PostGIS Extension"

    # Check if PostGIS is available
    if ! dpkg -l | grep -q postgis; then
        print_warning "PostGIS is not installed"
        print_info "Install with: sudo apt-get install postgis postgresql-15-postgis-3"
        print_info "Skipping PostGIS extension..."
        return 0
    fi

    print_info "Enabling PostGIS extension..."
    sudo -u postgres psql -d $DB_NAME -c "CREATE EXTENSION IF NOT EXISTS postgis;"
    print_success "PostGIS extension enabled"
}

create_schema() {
    print_header "Creating Database Schema"

    if [ -f "$SCHEMA_FILE" ]; then
        print_info "Running schema file: $SCHEMA_FILE"
        PGPASSWORD="$DB_PASSWORD" psql -U $DB_USER -d $DB_NAME -h localhost -f "$SCHEMA_FILE"
        print_success "Schema created from file"
    else
        print_warning "Schema file not found: $SCHEMA_FILE"
        print_info "You'll need to run migrations manually"
    fi
}

test_connection() {
    print_header "Testing Connection"

    print_info "Testing connection as '$DB_USER'..."

    if PGPASSWORD="$DB_PASSWORD" psql -U $DB_USER -d $DB_NAME -h localhost -c "SELECT version();" &> /dev/null; then
        print_success "Connection successful!"

        # Show connection string
        print_info ""
        print_info "Connection details:"
        echo "  Database: $DB_NAME"
        echo "  User: $DB_USER"
        echo "  Host: localhost"
        echo "  Port: 5432"
        echo ""
        echo "  Connection string:"
        echo "  postgresql://$DB_USER:$DB_PASSWORD@localhost:5432/$DB_NAME"
    else
        print_error "Connection failed"
        exit 1
    fi
}

###############################################################################
# Main
###############################################################################

main() {
    print_header "GigLez Database Setup"
    echo ""

    # Run checks
    check_sudo
    check_postgresql
    check_password_file

    echo ""

    # Setup database
    create_database
    create_user
    grant_privileges
    enable_postgis

    # Create schema if file exists
    if [ -f "$SCHEMA_FILE" ]; then
        create_schema
    fi

    # Test connection
    test_connection

    echo ""
    print_header "Setup Complete!"
    print_success "Database is ready for use"
    echo ""
    print_info "Next steps:"
    echo "  1. Run import script to load test data:"
    echo "     cd $PROJECT_ROOT"
    echo "     source venv/bin/activate"
    echo "     python scripts/import_with_african_gps.py --limit 5000 --no-matching"
    echo ""
    echo "  2. Start the API server:"
    echo "     uvicorn src.api.main:app --reload"
    echo ""
}

# Run main
main