#!/bin/bash

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

# Configuration
ENVIRONMENT=${1:-local}
VERBOSE=${2:-false}

# Determine compose file
if [ "$ENVIRONMENT" == "production" ]; then
    COMPOSE_FILE="docker-compose.production.yml"
    FRONTEND_PORT=80
else
    COMPOSE_FILE="docker-compose.yml"
    FRONTEND_PORT=3000
fi

# Exit code
EXIT_CODE=0

# Function to print colored messages
print_info() {
    echo -e "${BLUE}ℹ️  $1${NC}"
}

print_success() {
    echo -e "${GREEN}✅ $1${NC}"
}

print_warning() {
    echo -e "${YELLOW}⚠️  $1${NC}"
}

print_error() {
    echo -e "${RED}❌ $1${NC}"
}

# Function to check container status
check_container_status() {
    local service=$1
    local status=$(docker-compose -f "$COMPOSE_FILE" ps "$service" 2>/dev/null | tail -n +3 | awk '{print $1, $4, $5, $6, $7}' | grep "$service" || echo "")
    
    if [ -z "$status" ]; then
        return 1
    fi
    
    if echo "$status" | grep -q "Up\|healthy"; then
        return 0
    else
        return 1
    fi
}

# Function to check HTTP endpoint
check_http_endpoint() {
    local service=$1
    local url=$2
    local timeout=${3:-5}
    
    if curl -sf --max-time "$timeout" "$url" >/dev/null 2>&1; then
        return 0
    else
        return 1
    fi
}

# Function to check service health
check_service_health() {
    local service=$1
    local endpoint=$2
    local port=${3:-}
    
    print_info "Checking ${service}..."
    
    # Check container status
    if ! check_container_status "$service"; then
        print_error "${service}: Container is not running or unhealthy"
        return 1
    fi
    
    # Check HTTP endpoint if provided
    if [ -n "$endpoint" ]; then
        if check_http_endpoint "$service" "$endpoint"; then
            print_success "${service}: Healthy (HTTP endpoint responding)"
            return 0
        else
            print_warning "${service}: Container running but HTTP endpoint not responding"
            return 1
        fi
    fi
    
    # If no endpoint, just check container status
    print_success "${service}: Healthy (container running)"
    return 0
}

# Function to check database connectivity
check_database() {
    local service=$1
    local user=$2
    local db=$3
    
    print_info "Checking ${service} database connectivity..."
    
    if docker-compose -f "$COMPOSE_FILE" exec -T "$service" psql -U "$user" -d "$db" -c "SELECT 1;" >/dev/null 2>&1; then
        print_success "${service}: Database connection successful"
        return 0
    else
        print_error "${service}: Database connection failed"
        return 1
    fi
}

# Function to check Redis connectivity
check_redis() {
    local service=$1
    
    print_info "Checking ${service} connectivity..."
    
    if docker-compose -f "$COMPOSE_FILE" exec -T "$service" redis-cli ping >/dev/null 2>&1; then
        print_success "${service}: Redis connection successful"
        return 0
    else
        print_error "${service}: Redis connection failed"
        return 1
    fi
}

# Main health check
main() {
    print_info "🏥 Checking health of all services for ${ENVIRONMENT} environment..."
    echo ""
    
    # Check prerequisites
    if [ ! -f "$COMPOSE_FILE" ]; then
        print_error "$COMPOSE_FILE not found. Aborting.";
        exit 1;
    fi
    
    # Check PostgreSQL
    if ! check_service_health "postgres" ""; then
        EXIT_CODE=1
    else
        if [ "$ENVIRONMENT" == "local" ]; then
            check_database "postgres" "veza_user" "veza_local" || EXIT_CODE=1
        fi
    fi
    echo ""
    
    # Check Redis
    if ! check_service_health "redis" ""; then
        EXIT_CODE=1
    else
        check_redis "redis" || EXIT_CODE=1
    fi
    echo ""
    
    # Check Backend API
    if ! check_service_health "backend-api" "http://localhost:8080/health"; then
        EXIT_CODE=1
    fi
    echo ""
    
    # Check Chat Server
    if ! check_service_health "chat-server" "http://localhost:8081/health"; then
        EXIT_CODE=1
    fi
    echo ""
    
    # Check Stream Server
    if ! check_service_health "stream-server" "http://localhost:8082/health"; then
        EXIT_CODE=1
    fi
    echo ""
    
    # Check Frontend
    if ! check_service_health "frontend" "http://localhost:${FRONTEND_PORT}/health"; then
        EXIT_CODE=1
    fi
    echo ""
    
    # Display summary
    if [ $EXIT_CODE -eq 0 ]; then
        print_success "🎉 All services are healthy!"
        echo ""
        
        if [ "$VERBOSE" == "true" ] || [ "$VERBOSE" == "-v" ]; then
            print_info "Detailed service status:"
            docker-compose -f "$COMPOSE_FILE" ps
        fi
    else
        print_error "❌ Some services are unhealthy!"
        echo ""
        print_info "Detailed service status:"
        docker-compose -f "$COMPOSE_FILE" ps
        echo ""
        print_warning "Run with verbose flag (-v) for more details"
    fi
    
    exit $EXIT_CODE
}

# Run main function
main