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

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

export SOC_SHARED_NETWORK="${SOC_SHARED_NETWORK:-soc_shared}"
export IRIS_HTTPS_PORT="${IRIS_HTTPS_PORT:-8443}"
export INTERFACE_HTTPS_PORT="${IRIS_HTTPS_PORT}"
export SHUFFLE_OPENSEARCH_PORT="${SHUFFLE_OPENSEARCH_PORT:-9201}"
export PAGERDUTY_STUB_PORT="${PAGERDUTY_STUB_PORT:-18080}"
export SOC_INTEGRATOR_PORT="${SOC_INTEGRATOR_PORT:-8088}"

if [[ $# -eq 0 ]]; then
  COMMAND="up"
  ARGS=(-d)
else
  COMMAND="$1"
  shift
  ARGS=("$@")
fi

if [[ "${COMMAND}" == "help" || "${COMMAND}" == "--help" || "${COMMAND}" == "-h" ]]; then
  cat <<'EOF'
Usage:
  ./run-combined-stack.sh [command] [target] [options]

Commands:
  up           Start services (default: up -d when no args)
  down         Stop services (requires explicit target)
  logs         View logs
  status       Show container and endpoint status
  cleanup      Remove unused Docker resources (containers/images/cache)
  help         Show this help message

Targets:
  all|--all    All stacks (wazuh, iris, shuffle, pagerduty, integrator, flask-openapi-shuffle)
  wazuh        Wazuh single-node stack
  iris         IRIS-web stack
  shuffle      Shuffle stack
  pagerduty    PagerDuty stub
  integrator   SOC integrator stack
  flask-openapi-shuffle  Flask OpenAPI demo stack

Examples:
  ./run-combined-stack.sh
  ./run-combined-stack.sh up --all -d
  ./run-combined-stack.sh up iris -d
  ./run-combined-stack.sh up flask-openapi-shuffle -d
  ./run-combined-stack.sh down shuffle
  ./run-combined-stack.sh down --all
  ./run-combined-stack.sh logs integrator -f
  ./run-combined-stack.sh logs --all --tail 200
  ./run-combined-stack.sh cleanup
  ./run-combined-stack.sh cleanup --with-volumes
  ./run-combined-stack.sh status
EOF
  exit 0
fi

if [[ "${COMMAND}" == "up" || "${COMMAND}" == "down" || "${COMMAND}" == "logs" ]]; then
  if ! docker network inspect "${SOC_SHARED_NETWORK}" >/dev/null 2>&1; then
    docker network create "${SOC_SHARED_NETWORK}" >/dev/null
  fi
fi

if [[ "${COMMAND}" == "status" ]]; then
  exec "${ROOT_DIR}/soc-status.sh"
fi

run_cleanup() {
  local with_volumes="${1:-false}"

  echo "Pruning stopped containers..."
  docker container prune -f

  echo "Pruning unused images..."
  docker image prune -a -f

  echo "Pruning builder cache..."
  docker builder prune -f

  if [[ "${with_volumes}" == "true" ]]; then
    echo "Pruning unused volumes..."
    docker volume prune -f
  fi

  echo "Current Docker disk usage:"
  docker system df
}

run_wazuh() {
  docker compose \
    --project-name wazuh-single \
    --project-directory "${ROOT_DIR}/wazuh-docker/single-node" \
    -f "${ROOT_DIR}/wazuh-docker/single-node/docker-compose.yml" \
    -f "${ROOT_DIR}/compose-overrides/wazuh.shared-network.yml" \
    "${COMMAND}" ${ARGS[@]+"${ARGS[@]}"}
}

run_iris() {
  docker compose \
    --project-name iris-web \
    --project-directory "${ROOT_DIR}/iris-web" \
    -f "${ROOT_DIR}/iris-web/docker-compose.yml" \
    -f "${ROOT_DIR}/compose-overrides/iris.shared-network.yml" \
    "${COMMAND}" ${ARGS[@]+"${ARGS[@]}"}
}

run_shuffle() {
  docker compose \
    --project-name shuffle \
    --project-directory "${ROOT_DIR}/Shuffle" \
    -f "${ROOT_DIR}/Shuffle/docker-compose.yml" \
    -f "${ROOT_DIR}/compose-overrides/shuffle.shared-network.yml" \
    "${COMMAND}" ${ARGS[@]+"${ARGS[@]}"}
}

run_pagerduty_stub() {
  docker compose \
    --project-name pagerduty-stub \
    --project-directory "${ROOT_DIR}" \
    -f "${ROOT_DIR}/compose-overrides/pagerduty.stub.yml" \
    "${COMMAND}" ${ARGS[@]+"${ARGS[@]}"}
}

run_soc_integrator() {
  docker compose \
    --project-name soc-integrator \
    --project-directory "${ROOT_DIR}/compose-overrides" \
    -f "${ROOT_DIR}/compose-overrides/soc-integrator.yml" \
    "${COMMAND}" ${ARGS[@]+"${ARGS[@]}"}
}

run_flask_openapi_shuffle() {
  docker compose \
    --project-name flask-openapi-shuffle \
    --project-directory "${ROOT_DIR}/flask-openapi-shuffle" \
    -f "${ROOT_DIR}/flask-openapi-shuffle/docker-compose.yml" \
    "${COMMAND}" ${ARGS[@]+"${ARGS[@]}"}
}

run_target() {
  local target="$1"
  case "${target}" in
    wazuh)
      run_wazuh
      ;;
    iris)
      run_iris
      ;;
    shuffle)
      run_shuffle
      ;;
    pagerduty)
      run_pagerduty_stub
      ;;
    integrator)
      run_soc_integrator
      ;;
    flask-openapi-shuffle)
      run_flask_openapi_shuffle
      ;;
    *)
      echo "Unknown target: ${target}"
      echo "Use one of: wazuh, iris, shuffle, pagerduty, integrator, flask-openapi-shuffle"
      exit 1
      ;;
  esac
}

run_all() {
  local mode="$1"

  if [[ "${mode}" == "down" ]]; then
    run_flask_openapi_shuffle
    run_soc_integrator
    run_pagerduty_stub
    run_shuffle
    run_iris
    run_wazuh
  else
    run_wazuh
    run_iris
    run_shuffle
    run_pagerduty_stub
    run_soc_integrator
    run_flask_openapi_shuffle
  fi
}

follow_all_logs() {
  COMMAND="logs"
  ARGS=("-f" "--tail" "${LOG_TAIL:-100}")

  run_wazuh &
  run_iris &
  run_shuffle &
  run_pagerduty_stub &
  run_soc_integrator &
  run_flask_openapi_shuffle &

  trap 'kill 0' INT TERM
  wait
}

run_logs_for_target() {
  local target="${1:-all}"
  case "${target}" in
    wazuh)
      run_wazuh
      ;;
    iris)
      run_iris
      ;;
    shuffle)
      run_shuffle
      ;;
    pagerduty)
      run_pagerduty_stub
      ;;
    integrator)
      run_soc_integrator
      ;;
    flask-openapi-shuffle)
      run_flask_openapi_shuffle
      ;;
    all|--all)
      run_wazuh
      run_iris
      run_shuffle
      run_pagerduty_stub
      run_soc_integrator
      run_flask_openapi_shuffle
      ;;
    *)
      echo "Unknown logs target: ${target}"
      echo "Use one of: wazuh, iris, shuffle, pagerduty, integrator, flask-openapi-shuffle"
      exit 1
      ;;
  esac
}

if [[ "${COMMAND}" == "down" ]]; then
  TARGET="${1:-}"
  if [[ -z "${TARGET}" ]]; then
    echo "Refusing to run 'down' without a target."
    echo "Use one of:"
    echo "  ./run-combined-stack.sh down <wazuh|iris|shuffle|pagerduty|integrator|flask-openapi-shuffle>"
    echo "  ./run-combined-stack.sh down --all"
    exit 1
  fi

  case "${TARGET}" in
    wazuh|iris|shuffle|pagerduty|integrator|flask-openapi-shuffle)
      ARGS=("${ARGS[@]:1}")
      run_target "${TARGET}"
      ;;
    all|--all)
      ARGS=("${ARGS[@]:1}")
      run_all "down"
      ;;
    *)
      run_all "down"
      ;;
  esac
elif [[ "${COMMAND}" == "logs" ]]; then
  LOGS_TARGET="${1:-all}"
  case "${LOGS_TARGET}" in
    wazuh|iris|shuffle|pagerduty|integrator|flask-openapi-shuffle)
      ARGS=("${ARGS[@]:1}")
      run_logs_for_target "${LOGS_TARGET}"
      ;;
    all|--all)
      ARGS=("${ARGS[@]:1}")
      run_logs_for_target "all"
      ;;
    *)
      for arg in ${ARGS[@]+"${ARGS[@]}"}; do
        if [[ "${arg}" == "-f" || "${arg}" == "--follow" ]]; then
          echo "For follow mode, specify one target:"
          echo "./run-combined-stack.sh logs <wazuh|iris|shuffle|pagerduty|integrator|flask-openapi-shuffle> -f"
          exit 1
        fi
      done
      run_logs_for_target "all"
          ;;
  esac
elif [[ "${COMMAND}" == "up" ]]; then
  TARGET="${1:-all}"
  case "${TARGET}" in
    wazuh|iris|shuffle|pagerduty|integrator|flask-openapi-shuffle)
      ARGS=("${ARGS[@]:1}")
      run_target "${TARGET}"
      ;;
    all|--all)
      ARGS=("${ARGS[@]:1}")
      HAS_DETACH="false"
      for arg in ${ARGS[@]+"${ARGS[@]}"}; do
        if [[ "${arg}" == "-d" || "${arg}" == "--detach" ]]; then
          HAS_DETACH="true"
          break
        fi
      done

      if [[ "${HAS_DETACH}" == "true" ]]; then
        run_all "up"
      else
        ARGS+=("-d")
        run_all "up"
        follow_all_logs
      fi
      ;;
    *)
      run_all "up"
      ;;
  esac
elif [[ "${COMMAND}" == "cleanup" ]]; then
  WITH_VOLUMES="false"
  for arg in ${ARGS[@]+"${ARGS[@]}"}; do
    case "${arg}" in
      --with-volumes|-v)
        WITH_VOLUMES="true"
        ;;
      *)
        ;;
    esac
  done
  run_cleanup "${WITH_VOLUMES}"
else
  run_all "up"
fi