#!/bin/bash
# Database Setup Script for TAK Server
set -e

# Database configuration
DB_NAME="${DB_NAME:-cot}"
DB_USER="${DB_USER:-martiuser}"
DB_PASSWORD="${DB_PASSWORD:-password}"
DB_HOST="${DB_HOST:-localhost}"
DB_PORT="${DB_PORT:-5432}"

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

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

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

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

# Check if PostgreSQL is running
check_postgres() {
    if ! pg_isready -q -h "$DB_HOST" -p "$DB_PORT"; then
        log_error "PostgreSQL is not running or not accessible"
        return 1
    fi
    log_info "PostgreSQL is running"
}

# Create database user if not exists
create_db_user() {
    log_info "Creating database user '$DB_USER'..."
    
    if sudo -u postgres psql -tAc "SELECT 1 FROM pg_roles WHERE rolname='$DB_USER'" | grep -q 1; then
        log_warn "User '$DB_USER' already exists"
    else
        sudo -u postgres createuser --createdb --no-superuser --no-createrole "$DB_USER"
        log_info "User '$DB_USER' created"
    fi
    
    # Set password
    sudo -u postgres psql -c "ALTER USER $DB_USER PASSWORD '$DB_PASSWORD';"
    log_info "Password set for user '$DB_USER'"
}

# Create database if not exists
create_database() {
    log_info "Creating database '$DB_NAME'..."
    
    if sudo -u postgres psql -lqt | cut -d \| -f 1 | grep -qw "$DB_NAME"; then
        log_warn "Database '$DB_NAME' already exists"
    else
        sudo -u postgres createdb -O "$DB_USER" "$DB_NAME"
        log_info "Database '$DB_NAME' created"
    fi
}

# Install PostGIS extension
install_postgis() {
    log_info "Installing PostGIS extension..."
    
    if sudo -u postgres psql -d "$DB_NAME" -tAc "SELECT 1 FROM pg_extension WHERE extname='postgis'" | grep -q 1; then
        log_warn "PostGIS extension already installed"
    else
        sudo -u postgres psql -d "$DB_NAME" -c "CREATE EXTENSION postgis;"
        log_info "PostGIS extension installed"
    fi
}

# Test database connection
test_connection() {
    log_info "Testing database connection..."
    
    if PGPASSWORD="$DB_PASSWORD" psql -h "$DB_HOST" -p "$DB_PORT" -U "$DB_USER" -d "$DB_NAME" -c "SELECT version();" > /dev/null 2>&1; then
        log_info "Database connection successful"
    else
        log_error "Database connection failed"
        return 1
    fi
}

# Main setup function
setup_database() {
    log_info "Starting database setup..."
    
    check_postgres
    create_db_user
    create_database
    install_postgis
    test_connection
    
    log_info "Database setup completed successfully"
}

# Print database configuration
print_config() {
    echo "Database Configuration:"
    echo "  Host: $DB_HOST"
    echo "  Port: $DB_PORT"
    echo "  Database: $DB_NAME"
    echo "  User: $DB_USER"
    echo "  Password: [REDACTED]"
}

# Main execution
case "${1:-setup}" in
    "setup")
        setup_database
        ;;
    "test")
        test_connection
        ;;
    "config")
        print_config
        ;;
    "user")
        create_db_user
        ;;
    "db")
        create_database
        ;;
    "postgis")
        install_postgis
        ;;
    *)
        echo "Usage: $0 [setup|test|config|user|db|postgis]"
        echo "  setup   - Full database setup (default)"
        echo "  test    - Test database connection"
        echo "  config  - Print database configuration"
        echo "  user    - Create database user only"
        echo "  db      - Create database only"
        echo "  postgis - Install PostGIS extension only"
        exit 1
        ;;
esac