#!/usr/bin/env bash

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"

DEFAULT_ENV_FILE=".env.default"
USER_ENV_FILE=".env"

log() {
  printf '%s\n' "$*" >&2
}

die() {
  printf 'Error: %s\n' "$*" >&2
  exit 1
}

detect_compose() {
  if docker compose version >/dev/null 2>&1; then
    COMPOSE_CMD=(docker compose)
    return
  fi

  if command -v docker-compose >/dev/null 2>&1; then
    COMPOSE_CMD=(docker-compose)
    return
  fi

  die "Docker Compose is not available. Install Docker Compose, then run this command again."
}

generate_secret_key() {
  if command -v openssl >/dev/null 2>&1; then
    openssl rand -base64 42
    return
  fi

  if command -v dd >/dev/null 2>&1 && command -v base64 >/dev/null 2>&1; then
    dd if=/dev/urandom bs=42 count=1 2>/dev/null | base64 | tr -d '\n'
    printf '\n'
    return
  fi

  return 1
}

ensure_env_files() {
  [[ -f "$DEFAULT_ENV_FILE" ]] || die "$DEFAULT_ENV_FILE is missing."

  if [[ -f "$USER_ENV_FILE" ]]; then
    return
  fi

  : >"$USER_ENV_FILE"

  if [[ ! -t 0 ]]; then
    log "Created $USER_ENV_FILE for local overrides."
    return
  fi

  printf 'Created %s for local overrides.\n' "$USER_ENV_FILE"
  printf 'Do you need a custom deployment now? (Most users can press Enter to skip.) [y/N] '
  read -r answer

  case "${answer:-}" in
    y | Y | yes | YES | Yes)
      cat <<'EOF'
Edit .env with the settings you want to override, using .env.example as the full reference.
Run ./dify-compose up -d again when you are ready.
EOF
      exit 0
      ;;
  esac
}

user_env_value() {
  local key="$1"
  awk -F= -v target="$key" '
    /^[[:space:]]*#/ || !/=/{ next }
    {
      key = $1
      gsub(/^[[:space:]]+|[[:space:]]+$/, "", key)
      if (key == target) {
        value = substr($0, index($0, "=") + 1)
        gsub(/^[[:space:]]+|[[:space:]]+$/, "", value)
        if ((value ~ /^".*"$/) || (value ~ /^'\''.*'\''$/)) {
          value = substr(value, 2, length(value) - 2)
        }
        result = value
      }
    }
    END { print result }
  ' "$USER_ENV_FILE"
}

set_user_env_value() {
  local key="$1"
  local value="$2"
  local temp_file

  temp_file="$(mktemp "${TMPDIR:-/tmp}/dify-env.XXXXXX")"
  awk -F= -v target="$key" -v replacement="$key=$value" '
    BEGIN { replaced = 0 }
    /^[[:space:]]*#/ || !/=/{ print; next }
    {
      key = $1
      gsub(/^[[:space:]]+|[[:space:]]+$/, "", key)
      if (key == target) {
        if (!replaced) {
          print replacement
          replaced = 1
        }
        next
      }
      print
    }
    END {
      if (!replaced) {
        print replacement
      }
    }
  ' "$USER_ENV_FILE" >"$temp_file"
  mv "$temp_file" "$USER_ENV_FILE"
}

ensure_secret_key() {
  local current_secret_key
  local secret_key

  current_secret_key="$(user_env_value SECRET_KEY)"
  if [[ -n "$current_secret_key" ]]; then
    return
  fi

  secret_key="$(generate_secret_key)" || die "Unable to generate SECRET_KEY. Install openssl or configure SECRET_KEY in .env."
  set_user_env_value SECRET_KEY "$secret_key"
  log "Generated SECRET_KEY in $USER_ENV_FILE."
}

env_value() {
  local key="$1"
  awk -F= -v target="$key" '
    /^[[:space:]]*#/ || !/=/{ next }
    {
      key = $1
      gsub(/^[[:space:]]+|[[:space:]]+$/, "", key)
      if (key == target) {
        value = substr($0, index($0, "=") + 1)
        gsub(/^[[:space:]]+|[[:space:]]+$/, "", value)
        if ((value ~ /^".*"$/) || (value ~ /^'\''.*'\''$/)) {
          value = substr(value, 2, length(value) - 2)
        }
        result = value
      }
    }
    END { print result }
  ' "$DEFAULT_ENV_FILE" "$USER_ENV_FILE"
}

user_overrides() {
  local key="$1"
  grep -Eq "^[[:space:]]*${key}[[:space:]]*=" "$USER_ENV_FILE"
}

write_merged_env() {
  awk '
    function trim(s) {
      sub(/^[[:space:]]+/, "", s)
      sub(/[[:space:]]+$/, "", s)
      return s
    }

    /^[[:space:]]*#/ || !/=/{ next }

    {
      key = $0
      sub(/=.*/, "", key)
      key = trim(key)
      if (key == "") {
        next
      }

      value = substr($0, index($0, "=") + 1)
      value = trim(value)

      if (!(key in seen)) {
        order[++count] = key
        seen[key] = 1
      }

      values[key] = value
    }

    END {
      for (i = 1; i <= count; i++) {
        key = order[i]
        print key "=" values[key]
      }
    }
  ' "$DEFAULT_ENV_FILE" "$USER_ENV_FILE" >"$MERGED_ENV_FILE"
}

set_merged_env_value() {
  local key="$1"
  local value="$2"
  local temp_file

  temp_file="$(mktemp "${TMPDIR:-/tmp}/dify-compose-env.XXXXXX")"
  awk -F= -v target="$key" -v replacement="$key=$value" '
    BEGIN { replaced = 0 }
    /^[[:space:]]*#/ || !/=/{ print; next }
    {
      key = $1
      gsub(/^[[:space:]]+|[[:space:]]+$/, "", key)
      if (key == target) {
        if (!replaced) {
          print replacement
          replaced = 1
        }
        next
      }
      print
    }
    END {
      if (!replaced) {
        print replacement
      }
    }
  ' "$MERGED_ENV_FILE" >"$temp_file"
  mv "$temp_file" "$MERGED_ENV_FILE"
}

set_if_not_overridden() {
  local key="$1"
  local value="$2"

  if user_overrides "$key"; then
    return
  fi

  set_merged_env_value "$key" "$value"
}

metadata_db_host() {
  case "$1" in
    mysql) printf 'db_mysql' ;;
    postgresql | '') printf 'db_postgres' ;;
    *) printf '%s' "$(env_value DB_HOST)" ;;
  esac
}

metadata_db_port() {
  case "$1" in
    mysql) printf '3306' ;;
    postgresql | '') printf '5432' ;;
    *) printf '%s' "$(env_value DB_PORT)" ;;
  esac
}

metadata_db_user() {
  case "$1" in
    mysql) printf 'root' ;;
    postgresql | '') printf 'postgres' ;;
    *) printf '%s' "$(env_value DB_USERNAME)" ;;
  esac
}

build_merged_env() {
  MERGED_ENV_FILE="$(mktemp "${TMPDIR:-/tmp}/dify-compose.XXXXXX")"
  trap 'rm -f "$MERGED_ENV_FILE"' EXIT

  write_merged_env

  local db_type
  local redis_host
  local redis_port
  local redis_username
  local redis_password
  local redis_auth
  local code_execution_api_key
  local weaviate_api_key

  db_type="$(env_value DB_TYPE)"

  set_if_not_overridden DB_HOST "$(metadata_db_host "$db_type")"
  set_if_not_overridden DB_PORT "$(metadata_db_port "$db_type")"
  set_if_not_overridden DB_USERNAME "$(metadata_db_user "$db_type")"

  if ! user_overrides CELERY_BROKER_URL; then
    redis_host="$(env_value REDIS_HOST)"
    redis_port="$(env_value REDIS_PORT)"
    redis_username="$(env_value REDIS_USERNAME)"
    redis_password="$(env_value REDIS_PASSWORD)"
    redis_auth=""

    if [[ -n "$redis_username" && -n "$redis_password" ]]; then
      redis_auth="${redis_username}:${redis_password}@"
    elif [[ -n "$redis_password" ]]; then
      redis_auth=":${redis_password}@"
    elif [[ -n "$redis_username" ]]; then
      redis_auth="${redis_username}@"
    fi

    set_merged_env_value CELERY_BROKER_URL "redis://${redis_auth}${redis_host:-redis}:${redis_port:-6379}/1"
  fi

  if ! user_overrides SANDBOX_API_KEY; then
    code_execution_api_key="$(env_value CODE_EXECUTION_API_KEY)"
    set_if_not_overridden SANDBOX_API_KEY "${code_execution_api_key:-dify-sandbox}"
  fi

  if ! user_overrides WEAVIATE_AUTHENTICATION_APIKEY_ALLOWED_KEYS; then
    weaviate_api_key="$(env_value WEAVIATE_API_KEY)"
    set_if_not_overridden WEAVIATE_AUTHENTICATION_APIKEY_ALLOWED_KEYS \
      "${weaviate_api_key:-WVF5YThaHlkYwhGUSmCRgsX3tD5ngdN8pkih}"
  fi
}

main() {
  detect_compose
  ensure_env_files
  ensure_secret_key
  build_merged_env

  if [[ "$#" -eq 0 ]]; then
    set -- up -d
  fi

  "${COMPOSE_CMD[@]}" --env-file "$MERGED_ENV_FILE" "$@"
}

main "$@"