#!/bin/bash
# TAK Server Health Check Script
set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
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"
}

log_check() {
    echo -e "${BLUE}[CHECK]${NC} $1"
}

# Check if PostgreSQL is running
check_postgres() {
    log_check "Checking PostgreSQL..."
    
    if pg_isready -q -h localhost -p 5432; then
        log_info "PostgreSQL is running"
        return 0
    else
        log_error "PostgreSQL is not running"
        return 1
    fi
}

# Check database connectivity
check_database() {
    log_check "Checking database connectivity..."
    
    if PGPASSWORD="password" psql -h localhost -p 5432 -U martiuser -d cot -c "SELECT 1;" > /dev/null 2>&1; then
        log_info "Database connection successful"
        return 0
    else
        log_error "Database connection failed"
        return 1
    fi
}

# Check PostGIS extension
check_postgis() {
    log_check "Checking PostGIS extension..."
    
    if PGPASSWORD="password" psql -h localhost -p 5432 -U martiuser -d cot -tAc "SELECT 1 FROM pg_extension WHERE extname='postgis'" | grep -q 1; then
        log_info "PostGIS extension is installed"
        return 0
    else
        log_error "PostGIS extension is not installed"
        return 1
    fi
}

# Check TAK server files
check_tak_files() {
    log_check "Checking TAK server files..."
    
    if [ -f "/opt/tak/configureInDocker.sh" ]; then
        log_info "TAK server configuration script found"
        return 0
    else
        log_error "TAK server configuration script not found"
        return 1
    fi
}

# Check TAK server process
check_tak_process() {
    log_check "Checking TAK server process..."
    
    if pgrep -f "takserver" > /dev/null; then
        log_info "TAK server process is running"
        return 0
    else
        log_warn "TAK server process not found"
        return 1
    fi
}

# Check network ports
check_ports() {
    log_check "Checking network ports..."
    
    local ports=(5432 8080 8443 8444 8446)
    local all_good=true
    
    for port in "${ports[@]}"; do
        if netstat -tuln | grep -q ":$port "; then
            log_info "Port $port is listening"
        else
            log_warn "Port $port is not listening"
            all_good=false
        fi
    done
    
    if [ "$all_good" = true ]; then
        return 0
    else
        return 1
    fi
}

# Check disk space
check_disk_space() {
    log_check "Checking disk space..."
    
    local usage=$(df /opt/tak | tail -1 | awk '{print $5}' | sed 's/%//')
    
    if [ "$usage" -lt 80 ]; then
        log_info "Disk space usage: ${usage}% (OK)"
        return 0
    elif [ "$usage" -lt 90 ]; then
        log_warn "Disk space usage: ${usage}% (WARNING)"
        return 1
    else
        log_error "Disk space usage: ${usage}% (CRITICAL)"
        return 1
    fi
}

# Check memory usage
check_memory() {
    log_check "Checking memory usage..."
    
    local mem_info=$(free | grep Mem)
    local total=$(echo $mem_info | awk '{print $2}')
    local used=$(echo $mem_info | awk '{print $3}')
    local usage=$((used * 100 / total))
    
    if [ "$usage" -lt 80 ]; then
        log_info "Memory usage: ${usage}% (OK)"
        return 0
    elif [ "$usage" -lt 90 ]; then
        log_warn "Memory usage: ${usage}% (WARNING)"
        return 1
    else
        log_error "Memory usage: ${usage}% (CRITICAL)"
        return 1
    fi
}

# Run all health checks
run_all_checks() {
    log_info "Running TAK Server Health Checks..."
    echo "=================================="
    
    local failed=0
    
    check_postgres || ((failed++))
    check_database || ((failed++))
    check_postgis || ((failed++))
    check_tak_files || ((failed++))
    check_tak_process || ((failed++))
    check_ports || ((failed++))
    check_disk_space || ((failed++))
    check_memory || ((failed++))
    
    echo "=================================="
    
    if [ $failed -eq 0 ]; then
        log_info "All health checks passed!"
        return 0
    else
        log_error "$failed health check(s) failed"
        return 1
    fi
}

# Main execution
case "${1:-all}" in
    "all")
        run_all_checks
        ;;
    "postgres")
        check_postgres
        ;;
    "database")
        check_database
        ;;
    "postgis")
        check_postgis
        ;;
    "tak-files")
        check_tak_files
        ;;
    "tak-process")
        check_tak_process
        ;;
    "ports")
        check_ports
        ;;
    "disk")
        check_disk_space
        ;;
    "memory")
        check_memory
        ;;
    *)
        echo "Usage: $0 [all|postgres|database|postgis|tak-files|tak-process|ports|disk|memory]"
        echo "  all        - Run all health checks (default)"
        echo "  postgres   - Check PostgreSQL service"
        echo "  database   - Check database connectivity"
        echo "  postgis    - Check PostGIS extension"
        echo "  tak-files  - Check TAK server files"
        echo "  tak-process- Check TAK server process"
        echo "  ports      - Check network ports"
        echo "  disk       - Check disk space"
        echo "  memory     - Check memory usage"
        exit 1
        ;;
esac