#!/bin/bash
set -e

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

echo -e "${GREEN}[Postgre-TLS] Starting PostgreSQL with encryption...${NC}"

# Function to setup SSL certificates from Docker Swarm secrets
setup_ssl_certificates() {
    echo -e "${YELLOW}[Postgre-TLS] Setting up SSL certificates from Docker Swarm secrets...${NC}"
    
    # Debug: list contents of /run/secrets
    echo -e "${YELLOW}[Postgre-TLS] Listing contents of /run/secrets:${NC}"
    ls -la /run/secrets/ || echo "No /run/secrets directory found"
    
    # Check if running in Docker Swarm mode
    if [ -f "/run/secrets/server_crt" ] || [ -f "/run/secrets/server.crt" ]; then
        echo -e "${YELLOW}[Postgre-TLS] Using SSL certificates from Docker Swarm secrets...${NC}"
        
        # Copy certificates from Docker secrets to expected locations
        if [ -f "/run/secrets/ca_crt" ]; then
            cp /run/secrets/ca_crt /var/lib/postgresql/ssl/ca.crt
        elif [ -f "/run/secrets/ca.crt" ]; then
            cp /run/secrets/ca.crt /var/lib/postgresql/ssl/ca.crt
        fi
        
        if [ -f "/run/secrets/server_crt" ]; then
            cp /run/secrets/server_crt /var/lib/postgresql/ssl/server.crt
        elif [ -f "/run/secrets/server.crt" ]; then
            cp /run/secrets/server.crt /var/lib/postgresql/ssl/server.crt
        fi
        
        if [ -f "/run/secrets/server_key" ]; then
            cp /run/secrets/server_key /var/lib/postgresql/ssl/server.key
        elif [ -f "/run/secrets/server.key" ]; then
            cp /run/secrets/server.key /var/lib/postgresql/ssl/server.key
        fi
        
        # Set proper permissions
        chmod 600 /var/lib/postgresql/ssl/server.key
        chmod 644 /var/lib/postgresql/ssl/server.crt
        chmod 644 /var/lib/postgresql/ssl/ca.crt
        
        # Set ownership to postgres user
        chown postgres:postgres /var/lib/postgresql/ssl/*
        
        echo -e "${GREEN}[Postgre-TLS] SSL certificates from Docker Swarm configured successfully${NC}"
    elif [ "$ENABLE_FALLBACK_SSL" = "true" ]; then
        echo -e "${YELLOW}[Postgre-TLS] Generating fallback SSL certificates for local development...${NC}"
        
        # Generate CA private key
        openssl genrsa -out /var/lib/postgresql/ssl/ca.key 2048
        
        # Generate CA certificate
        openssl req -new -x509 -key /var/lib/postgresql/ssl/ca.key \
            -out /var/lib/postgresql/ssl/ca.crt -days 365 \
            -subj "/C=US/ST=State/L=City/O=Postgre-TLS/CN=Postgre-TLS-CA"
        
        # Generate server private key
        openssl genrsa -out /var/lib/postgresql/ssl/server.key 2048
        
        # Generate server certificate signing request
        openssl req -new -key /var/lib/postgresql/ssl/server.key \
            -out /var/lib/postgresql/ssl/server.csr \
            -subj "/C=US/ST=State/L=City/O=Postgre-TLS/CN=localhost"
        
        # Generate server certificate signed by CA
        openssl x509 -req -in /var/lib/postgresql/ssl/server.csr \
            -CA /var/lib/postgresql/ssl/ca.crt \
            -CAkey /var/lib/postgresql/ssl/ca.key \
            -CAcreateserial -out /var/lib/postgresql/ssl/server.crt \
            -days 365
        
        # Clean up CSR
        rm /var/lib/postgresql/ssl/server.csr
        
        # Set proper permissions
        chmod 600 /var/lib/postgresql/ssl/server.key
        chmod 644 /var/lib/postgresql/ssl/server.crt
        chmod 644 /var/lib/postgresql/ssl/ca.crt
        
        # Set ownership to postgres user
        chown postgres:postgres /var/lib/postgresql/ssl/*
        
        echo -e "${GREEN}[Postgre-TLS] Fallback SSL certificates generated successfully${NC}"
    else
        echo -e "${RED}[Postgre-TLS] Error: SSL certificates must be provided via Docker Swarm secrets. Fallback generation is not supported.${NC}"
        exit 1
    fi
}

# Function to check if database needs initialization
needs_initialization() {
    if [ ! -s "$PGDATA/PG_VERSION" ]; then
        return 0  # true - needs initialization
    else
        return 1  # false - already initialized
    fi
}

# Function to load secrets from Docker Swarm or environment
load_secrets() {
    echo -e "${YELLOW}[Postgre-TLS] Loading secrets...${NC}"
    
    # Check if secrets are available
    if [ -f "/run/secrets/postgres_password" ]; then
        echo -e "${YELLOW}[Postgre-TLS] Loading secrets from /run/secrets...${NC}"
        export POSTGRES_PASSWORD=$(cat /run/secrets/postgres_password)
        export BACKUP_ENCRYPTION_KEY=$(cat /run/secrets/backup_encryption_key 2>/dev/null || echo "")
        export JWT_SECRET=$(cat /run/secrets/jwt_secret 2>/dev/null || echo "")
        export APP_SECRET_KEY=$(cat /run/secrets/app_secret_key 2>/dev/null || echo "")
        echo -e "${GREEN}[Postgre-TLS] Secrets loaded from /run/secrets${NC}"
        echo "Loaded POSTGRES_PASSWORD: $POSTGRES_PASSWORD"
    else
        echo -e "${YELLOW}[Postgre-TLS] Using environment variables for secrets${NC}"
        export POSTGRES_PASSWORD="${POSTGRES_PASSWORD:-change_me_in_production}"
    fi
    
    # Set default values for other variables
    export POSTGRES_DB="${POSTGRES_DB:-ploughgres}"
    export POSTGRES_USER="${POSTGRES_USER:-ploughgres_user}"
}

# Function to initialize database with encryption
initialize_database() {
    echo -e "${YELLOW}[Postgre-TLS] Initializing database...${NC}"
    
    # Create a temporary password file with proper ownership
    su-exec postgres sh -c "echo '$POSTGRES_PASSWORD' > /tmp/pgpass"
    su-exec postgres chmod 600 /tmp/pgpass
    
    # Initialize the database with proper encoding and authentication
    # Create the main database during initdb
    export POSTGRES_DB="${POSTGRES_DB:-ploughgres}"
    su-exec postgres initdb \
        --pgdata="$PGDATA" \
        --username="$POSTGRES_USER" \
        --pwfile=/tmp/pgpass \
        --auth-local=scram-sha-256 \
        --auth-host=scram-sha-256 \
        --encoding=UTF8 \
        --locale=C.UTF-8 \
        --data-checksums
    
    # Clean up password file
    rm -f /tmp/pgpass
    
    echo -e "${GREEN}[Postgre-TLS] Database initialized successfully${NC}"
}

# Function to start PostgreSQL
start_postgresql() {
    echo -e "${YELLOW}[Postgre-TLS] Starting PostgreSQL server...${NC}"
    
    # Start PostgreSQL with custom configuration
    exec su-exec postgres postgres \
        -c config_file=/etc/postgresql/postgresql.conf \
        -c ssl=on \
        -c ssl_cert_file=/var/lib/postgresql/ssl/server.crt \
        -c ssl_key_file=/var/lib/postgresql/ssl/server.key \
        -c ssl_ca_file=/var/lib/postgresql/ssl/ca.crt \
        -c shared_preload_libraries=pg_stat_statements \
        -c log_statement=ddl \
        -c log_destination=stderr \
        -c logging_collector=on \
        -c log_directory=/var/log/postgresql \
        -c log_filename=postgresql.log
}

# Main execution
main() {
    # Set PGDATA if not set
    export PGDATA="${PGDATA:-/var/lib/postgresql/data}"
    
    # Ensure data directory exists and has proper permissions
    # mkdir -p "$PGDATA"
    # chown postgres:postgres "$PGDATA"
    # chmod 700 "$PGDATA"
    
    # Load secrets from Docker Swarm or environment
    load_secrets
    
    # Setup SSL certificates
    setup_ssl_certificates
    
    # Initialize database if needed
    if needs_initialization; then
        echo -e "${YELLOW}[Postgre-TLS] Database not found, initializing...${NC}"
        initialize_database
        
        # Create temporary pg_hba.conf for trust authentication during init
        echo 'local all all trust' > /tmp/pg_hba.conf
        echo 'host all all 127.0.0.1/32 trust' >> /tmp/pg_hba.conf
        echo 'host all all ::1/128 trust' >> /tmp/pg_hba.conf
        
        # Start PostgreSQL temporarily with temporary hba for initialization
        echo -e "${YELLOW}[Postgre-TLS] Starting PostgreSQL temporarily for initialization...${NC}"
        su-exec postgres postgres \
            -c config_file=/etc/postgresql/postgresql.conf \
            -c ssl=off \
            -c listen_addresses=localhost \
            -c port=5432 \
            -c hba_file=/tmp/pg_hba.conf &
        
        # Wait for PostgreSQL to start
        until su-exec postgres pg_isready -h localhost -p 5432; do
            echo -e "${YELLOW}[Postgre-TLS] Waiting for PostgreSQL to start...${NC}"
            sleep 1
        done
        
        # Create the application database if it doesn't exist
        echo -e "${YELLOW}[Postgre-TLS] Creating database: $POSTGRES_DB${NC}"
        su-exec postgres createdb -h localhost -p 5432 -U "$POSTGRES_USER" "$POSTGRES_DB" 2>/dev/null || echo -e "${GREEN}[Postgre-TLS] Database $POSTGRES_DB already exists${NC}"
        
        # Run initialization scripts
        for f in /docker-entrypoint-initdb.d/*; do
            case "$f" in
                *.sh)
                    echo -e "${YELLOW}[Postgre-TLS] Running $f${NC}"
                    bash "$f"
                    ;;
                *.sql)
                    echo -e "${YELLOW}[Postgre-TLS] Running $f${NC}"
                    su-exec postgres psql -v ON_ERROR_STOP=1 -h localhost -p 5432 --username "$POSTGRES_USER" --dbname "$POSTGRES_DB" < "$f"
                    ;;
                *)
                    echo -e "${YELLOW}[Postgre-TLS] Ignoring $f${NC}"
                    ;;
            esac
        done
        
        # Stop the temporary PostgreSQL instance
        su-exec postgres pg_ctl -D "$PGDATA" -m fast -w stop
        
        # Configure final pg_hba.conf for SSL connections
        echo -e "${YELLOW}[Postgre-TLS] Configuring pg_hba.conf for SSL...${NC}"
        cat > "$PGDATA/pg_hba.conf" <<EOF
# TYPE  DATABASE        USER            ADDRESS                 METHOD
local   all             all                                     scram-sha-256
hostssl all             all             0.0.0.0/0               scram-sha-256
hostssl all             all             ::/0                    scram-sha-256
EOF
        chown postgres:postgres "$PGDATA/pg_hba.conf"
        chmod 600 "$PGDATA/pg_hba.conf"
        
        echo -e "${GREEN}[Postgre-TLS] Database initialization completed${NC}"
    else
        echo -e "${GREEN}[Postgre-TLS] Database already initialized${NC}"
    fi
    
    # Start PostgreSQL with full configuration
    start_postgresql
}

# Run main function
main "$@"