#!/usr/bin/env bash
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "${SCRIPT_DIR}/functions.sh"
resolve_paths
require_env
require_command "yq"
resolve_volume_path

TEMPLATE_COMPOSE="${INIT_DIR}/templates/worker.yaml"
TEMPLATE_CONFIG="${INIT_DIR}/config/worker.yaml"
TEMPLATE_LOG="${INIT_DIR}/config/worker-log.yaml"

require_file "$TEMPLATE_COMPOSE" "Compose template"
require_file "$TEMPLATE_CONFIG" "Config template"
require_file "$TEMPLATE_LOG" "Log template"

WORKERS_DIR="${VOLUME_PATH}/synapse/workers"
mkdir -p "$WORKERS_DIR" "${SCRIPT_DIR}/tmp"

TEMP_COMPOSE="${SCRIPT_DIR}/tmp/compose_workers.yaml"
# Initialize empty compose file for workers
> "$TEMP_COMPOSE"

for ((i=1; i<=WORKER_COUNT; i++)); do
    worker_name="synapse-worker-${i}"
    port=$((9093 + i))
    TEMP_WORKER_COMPOSE="${SCRIPT_DIR}/tmp/${worker_name}.yaml"

    log_info "Generating configuration for ${worker_name}..."

    # 1. Generate Worker Config (synapse config file)
    sed -e "s/{{WORKER_NAME}}/${worker_name}/g" \
        -e "s/{{PORT}}/${port}/g" \
        -e "s/{{SERVER_NAME}}/${HOMESERVER_FQDN}/g" \
        -e "s/{{DB_USER}}/${PG_USER_MATRIX}/g" \
        -e "s/{{DB_PASSWORD}}/${PG_PASSWORD_MATRIX}/g" \
        -e "s/{{DB_NAME}}/${PG_DB_MATRIX}/g" \
        "$TEMPLATE_CONFIG" > "${WORKERS_DIR}/${worker_name}.yaml"

    # 2. Generate Log Config
    sed -e "s/{{WORKER_NAME}}/${worker_name}/g" \
        -e "s/{{MAX_BYTES}}/${LOG_MAX_BYTES:-10485760}/g" \
        -e "s/{{BACKUP_COUNT}}/${LOG_BACKUP_COUNT:-5}/g" \
        -e "s/{{LOG_LEVEL}}/${LOG_LEVEL:-INFO}/g" \
        -e "s/{{SYNAPSE_LOG_LEVEL}}/${SYNAPSE_LOG_LEVEL:-INFO}/g" \
        "$TEMPLATE_LOG" > "${WORKERS_DIR}/${worker_name}.log.config"

    # 3. Generate Temporary Compose Snippet for this worker
    # Start with the base template
    sed -e "s/{{WORKER_NAME}}/${worker_name}/g" \
        -e "s/{{PORT}}/${port}/g" \
        -e "s/{{APP_TYPE}}/synapse.app.generic_worker/g" \
        "$TEMPLATE_COMPOSE" > "$TEMP_WORKER_COMPOSE"

    # 4. Conditionally Inject Traefik Labels using yq
    if [[ "$PROXY" == "TRAEFIK" && -n "$DOMAIN" ]]; then
        log_info "  -> Adding Traefik labels for ${worker_name}"

        # Construct the labels array safely
        # Note: We escape backticks for the shell, but yq handles the YAML structure
        yq -i '.services."'"${worker_name}"'".labels = [
          "traefik.enable=true",
          "traefik.http.services.'"${worker_name}"'.loadbalancer.server.port='"${port}"'",
          "traefik.http.routers.synapse-worker-'"${worker_name}"'.rule=Host(`'"${DOMAIN}"'`) && PathPrefix(`/_matrix`)",
          "traefik.http.routers.synapse-worker-'"${worker_name}"'.entrypoints=websecure",
          "traefik.http.routers.synapse-worker-'"${worker_name}"'.service='"${worker_name}"'",
          "traefik.http.routers.synapse-worker-'"${worker_name}"'.priority=10"
        ]' "$TEMP_WORKER_COMPOSE"
    else
        # Ensure labels key doesn't exist if not using Traefik
        yq -i 'del(.services."'"${worker_name}"'".labels)' "$TEMP_WORKER_COMPOSE" 2>/dev/null || true
    fi

    # 5. Merge this worker into the temporary aggregate compose file
    # If TEMP_COMPOSE is empty, just copy. Otherwise, merge.
    if [[ ! -s "$TEMP_COMPOSE" ]]; then
        cp "$TEMP_WORKER_COMPOSE" "$TEMP_COMPOSE"
    else
        # Merge: existing content * new worker content
        yq eval-all "select(fileIndex == 0) * select(fileIndex == 1)" \
            "$TEMP_COMPOSE" "$TEMP_WORKER_COMPOSE" > "${TEMP_COMPOSE}.merged"
        mv "${TEMP_COMPOSE}.merged" "$TEMP_COMPOSE"
    fi

    # Cleanup temp worker file
    rm -f "$TEMP_WORKER_COMPOSE"
done

# 6. Final Merge into main compose.yaml
if [[ -f "$COMPOSE_FILE" ]]; then
    # Check if workers already exist to avoid duplicates
    if ! grep -qE "^  synapse-worker-[0-9]+:" "$COMPOSE_FILE"; then
        # Merge main compose with generated workers
        yq eval-all "select(fileIndex == 0) * select(fileIndex == 1)" \
            "$COMPOSE_FILE" "$TEMP_COMPOSE" > "${COMPOSE_FILE}.new"

        if mv "${COMPOSE_FILE}.new" "$COMPOSE_FILE"; then
            log_info "Appended ${WORKER_COUNT} worker(s) to compose.yaml"
        else
            log_error "Failed to update compose.yaml"
            exit 1
        fi
    else
        log_warn "Some workers already exist in compose.yaml. Manual merge required."
        cp "$TEMP_COMPOSE" "${SCRIPT_DIR}/tmp/compose_workers.yaml"
        log_warn "Saved new worker definitions to ${SCRIPT_DIR}/tmp/compose_workers.yaml"
    fi
else
    # If no main compose exists, the generated file becomes the main one
    mv "$TEMP_COMPOSE" "$COMPOSE_FILE"
    log_info "Created new compose.yaml with ${WORKER_COUNT} worker(s)"
fi

rm -f "$TEMP_COMPOSE"
log_info "Worker generation complete"