#!/usr/bin/env bash

# Exit on error. Append "|| true" if you expect an error.
set -o errexit
# Exit on error inside any functions or subshells.
set -o errtrace
# Do not allow use of undefined vars. Use ${VAR:-} to use an undefined VAR
set -o nounset
# Catch error in pipe chains
set -o pipefail

# Script constants
readonly SCRIPT_NAME=$(basename "${0}")
readonly SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
readonly COMPOSE_FILE="docker-compose.yml"

# Color constants - using tput for better terminal compatibility
if [[ -t 1 ]]; then
    readonly COLOR_RESET="$(tput sgr0)"
    readonly COLOR_INFO="$(tput setaf 6)"    # Cyan
    readonly COLOR_SUCCESS="$(tput setaf 2)"  # Green
    readonly COLOR_WARNING="$(tput setaf 3)"  # Yellow
    readonly COLOR_ERROR="$(tput setaf 1)"    # Red
    readonly COLOR_DEBUG="$(tput setaf 5)"    # Magenta
else
    readonly COLOR_RESET=""
    readonly COLOR_INFO=""
    readonly COLOR_SUCCESS=""
    readonly COLOR_WARNING=""
    readonly COLOR_ERROR=""
    readonly COLOR_DEBUG=""
fi

show_spinner() {
    local -r pid="$1"
    local -r delay=0.1
    local spinstr='/-\|'
    printf "  "
    tput sc
    
    while kill -0 "${pid}" 2>/dev/null; do
        local temp=${spinstr#?}
        tput rc
        printf "%s[%c]%s" "${COLOR_WARNING}" "${spinstr}" "${COLOR_RESET}"
        local spinstr=${temp}${spinstr%"${temp}"}
        sleep "${delay}"
    done
    tput rc
    printf "     \r"
}

log() {
    local -r level="$1"
    local -r message="$2"
    local -r timestamp=$(date +"%Y-%m-%d %H:%M:%S")
    local -r use_printf="${3:-false}"
    local color

    case "${level}" in
        INFO)    color="${COLOR_INFO}" ;;
        SUCCESS) color="${COLOR_SUCCESS}" ;;
        WARNING) color="${COLOR_WARNING}" ;;
        ERROR)   color="${COLOR_ERROR}" ;;
        DEBUG)   color="${COLOR_DEBUG}" ;;
        *)       color="${COLOR_RESET}" ;;
    esac

    if [[ "${use_printf}" == true ]]; then
        printf "%s[%s] [%s]:%s %s" "${color}" "${timestamp}" "${level}" "${COLOR_RESET}" "${message}"
    else
        printf "%s[%s] [%s]:%s %s\n" "${color}" "${timestamp}" "${level}" "${COLOR_RESET}" "${message}"
    fi
}

log_info()    { log "INFO" "$1" "${2:-false}"; }
log_success() { log "SUCCESS" "$1" "${2:-false}"; }
log_warning() { log "WARNING" "$1" "${2:-false}"; }
log_error()   { log "ERROR" "$1" "${2:-false}"; }
log_debug()   { log "DEBUG" "$1" "${2:-false}"; }

# Error handler
trap 'error_handler $? $LINENO $BASH_LINENO "$BASH_COMMAND" $(printf "::%s" ${FUNCNAME[@]:-})' ERR

error_handler() {
    local -r exit_code="$1"
    local -r line_number="$2"
    local -r bash_lineno="$3"
    local -r command="$4"
    local -r function_trace="$5"
    local -r error_message="Error in ${SCRIPT_NAME}: line ${line_number}, command '${command}' exited with status ${exit_code}"
    
    log_error "${error_message}"
    exit "${exit_code}"
}

# Docker compose wrapper
docker_compose() {
    local -r compose_file="$1"
    local -r command="$2"
    
    if ! docker compose -f "${compose_file}" "${command}" &>/dev/null; then
        log_error "Failed to execute: docker compose -f ${compose_file} ${command}"
        return 1
    fi
}

get_container_stats() {
    local running_count=0
    local not_running_count=0
    local partially_running_count=0

    while IFS= read -r -d '' dir; do
        local compose_path="${dir}/${COMPOSE_FILE}"
        if [[ -f "${compose_path}" ]]; then
            local dir_name=$(basename "${dir}")
            
            # Get expected container count from compose file
            local expected_count
            expected_count=$(grep -c "image:" "${compose_path}" 2>/dev/null || echo "0")
            
            # Get actual running container count
            local running_containers
            running_containers=$(docker compose -f "${compose_path}" ps | grep -c "Up" 2>/dev/null || echo "0")
            
            # Ensure counts are integers
            expected_count=$(echo "$expected_count" | tr -cd '0-9')
            running_containers=$(echo "$running_containers" | tr -cd '0-9')
            
            # Default to 0 if empty
            expected_count=${expected_count:-0}
            running_containers=${running_containers:-0}
            
            if [ "$expected_count" -eq "$running_containers" ]; then
                log_success "${dir_name}: ${running_containers}/${expected_count}"
                running_count=$((running_count + 1))
            elif [ "$running_containers" -eq 0 ]; then
                log_error "${dir_name}: ${running_containers}/${expected_count}"
                not_running_count=$((not_running_count + 1))
            else
                log_warning "${dir_name}: ${running_containers}/${expected_count}"
                partially_running_count=$((partially_running_count + 1))
            fi
        fi
    done < <(find . -maxdepth 1 -type d -print0)

    log_info "Summary:"
    log_info "  Fully running: ${running_count}"
    log_info "  Not running: ${not_running_count}"
    log_info "  Partially running: ${partially_running_count}"
}

list_containers() {
    local -r show_status="${1:-false}"
    log_info "Running containers: $(docker ps -q | wc -l)"
    
    if [[ "${show_status}" == true ]]; then
        get_container_stats
    fi
}

stop_containers() {
    log_warning "Stopping all containers"
    
    while IFS= read -r -d '' dir; do
        local compose_path="${dir}/${COMPOSE_FILE}"
        if [[ -f "${compose_path}" ]]; then
            local dir_name=$(basename "${dir}")
            log_info "Stopping ${dir_name}" true
            docker_compose "${compose_path}" down --remove-orphans & show_spinner $!
            echo
        fi
    done < <(find . -maxdepth 1 -type d -print0)
    
    list_containers
}

pull_containers() {
    log_warning "Pulling images for all containers"
    
    while IFS= read -r -d '' dir; do
        local compose_path="${dir}/${COMPOSE_FILE}"
        if [[ -f "${compose_path}" ]]; then
            local dir_name=$(basename "${dir}")
            log_info "Pulling ${dir_name}" true
            docker_compose "${compose_path}" pull & show_spinner $!
            echo
        fi
    done < <(find . -maxdepth 1 -type d -print0)
}

start_containers() {
    log_warning "Starting up all containers"
    
    while IFS= read -r -d '' dir; do
        local compose_path="${dir}/${COMPOSE_FILE}"
        if [[ -f "${compose_path}" ]]; then
            local dir_name=$(basename "${dir}")
            log_info "Starting ${dir_name}"
            docker compose -f "${compose_path}" up -d
        fi
    done < <(find . -maxdepth 1 -type d -print0)
    
    list_containers
}

show_usage() {
    cat << EOF
Usage: ${SCRIPT_NAME} [options]

Docker container management script

Options:
    -h, --help      Show this help message
    -p, --pull      Pull the images for all containers
    -s, --start     Start all containers
    -r, --restart   Restart all containers
    -u, --update    Update all containers (stop, pull, start)
    -d, --down      Stop all containers
    -l, --list      List containers and their status

Examples:
    ${SCRIPT_NAME} --start    # Start all containers
    ${SCRIPT_NAME} --update   # Update all containers
    ${SCRIPT_NAME} --list     # Show container status
EOF
}

main() {
    local args=()
    while [[ $# -gt 0 ]]; do
        case "$1" in
            -h|--help)    show_usage; exit 0 ;;
            -p|--pull)    pull_containers; exit 0 ;;
            -s|--start)   start_containers; exit 0 ;;
            -r|--restart) stop_containers; start_containers; exit 0 ;;
            -u|--update)  stop_containers; pull_containers; start_containers; exit 0 ;;
            -d|--down)    stop_containers; exit 0 ;;
            -l|--list)    list_containers true; exit 0 ;;
            *)           log_error "Invalid option: $1"; show_usage; exit 1 ;;
        esac
    done
}

# Execute main function
main "$@"