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

GIT_ROOT=$(cd "${BASH_SOURCE%/*}" && git rev-parse --show-toplevel)
source "${GIT_ROOT}/scripts/colors.sh"
source "${GIT_ROOT}/scripts/codecov.sh"

if [[ $UNIT_TEST_RERUN_FAILS == 'true' ]]; then
  GOTESTSUM_EXTRAFLAGS="${GOTESTSUM_EXTRAFLAGS} --rerun-fails"
elif [[ $UNIT_TEST_FAILFAST == 'true' ]]; then
  GOTEST_EXTRAFLAGS="${GOTEST_EXTRAFLAGS} -failfast"
fi

if [[ $UNIT_TEST_USE_DEVELOPMENT_LOGGER == 'false' ]]; then
  if [[ -z $BUILD_TAGS ]]; then
    BUILD_TAGS="test_silent"
  else
    BUILD_TAGS="${BUILD_TAGS} test_silent"
  fi
fi

if [[ -z "${UNIT_TEST_COUNT}" ]]; then
  UNIT_TEST_COUNT=1
fi

redirect_stdout() {
  output_file=$1

  if [[ "${CI}" == 'true' ]]; then
    cat > "${output_file}";
  else
    tee "${output_file}";
  fi
}

run_test_for_packages() {
  local packages="$1"
  local iteration="$2"
  local count="$3"
  local single_timeout="$4"
  local log_message="$5"
  local run_filter="${6:-}"

  local output_file="test_${iteration}.log"
  local coverage_file="test_${iteration}.coverage.out"
  local report_file="report_${iteration}.xml"
  local rerun_report_file="report_rerun_fails_${iteration}.txt"
  local exit_code_file="exit_code_${iteration}.txt"
  local timeout="$(( single_timeout * count))m"

  if [[ "${UNIT_TEST_DRY_RUN}" == 'true' ]]; then
    echo -e "${GRN}Dry run ${iteration}. message:${RST} ${log_message}\n"\
    "${YLW}Dry run ${iteration}. packages:${RST} ${packages}\n"\
    "${YLW}Dry run ${iteration}. count:${RST} ${count}\n"\
    "${YLW}Dry run ${iteration}. timeout:${RST} ${timeout}\n"\
    "${YLW}Dry run ${iteration}. run_filter:${RST} ${run_filter:-<none>}"
    return 0
  fi

  echo -e "${GRN}Testing:${RST} ${log_message}. Iteration ${iteration}. -test.count=${count}. Timeout: ${timeout}"

  gotestsum_flags="${GOTESTSUM_EXTRAFLAGS}"
  if [[ "${CI}" == 'true' ]]; then
    gotestsum_flags="${gotestsum_flags} --junitfile=${report_file} --rerun-fails-report=${rerun_report_file}"
  fi

  # Prepare env variables for `test_with_coverage.sh`
  export TEST_WITH_COVERAGE_PACKAGES="${packages}"
  export TEST_WITH_COVERAGE_COUNT="${count}"
  export TEST_WITH_COVERAGE_REPORTS_DIR="$(mktemp -d)"

  # Cleanup previous coverage reports
  rm -f "${TEST_WITH_COVERAGE_REPORTS_DIR}/coverage.out.rerun.*"

  # Run tests
  local test_extra_flags="${GOTEST_EXTRAFLAGS}"
  if [[ -n "${run_filter}" ]]; then
    test_extra_flags="${test_extra_flags} -run ${run_filter}"
  fi

  gotestsum --packages="${packages}" ${gotestsum_flags} --raw-command -- \
    ./scripts/test_with_coverage.sh \
    ${test_extra_flags} \
    -timeout "${timeout}" \
    -tags "${BUILD_TAGS}" | \
    redirect_stdout "${output_file}"

  local go_test_exit=$?

  # Merge package coverage results
  go tool gocovmerge ${TEST_WITH_COVERAGE_REPORTS_DIR}/coverage.out.rerun.* > ${coverage_file}
  rm -f "${TEST_WITH_COVERAGE_REPORTS_DIR}/coverage.out.rerun.*"

  echo "${go_test_exit}" > "${exit_code_file}"
  if [[ "${go_test_exit}" -ne 0 ]]; then
    if [[ "${CI}" == 'true' ]]; then
      echo -e "${YLW}Failed, see the log:${RST} ${BLD}${output_file}${RST}"
    fi
  fi

  return ${go_test_exit}
}

rm -rf ./**/*.coverage.out

echo -e "${GRN}Testing HEAD:${RST} $(git rev-parse HEAD)"

DEFAULT_TIMEOUT_MINUTES=5
PROTOCOL_TIMEOUT_MINUTES=45
DEFAULT_PROTOCOL_SHARDS=5
MIN_PROTOCOL_SHARDS=1
MAX_PROTOCOL_SHARDS=32

resolve_protocol_shards() {
  local shards

  if [[ -n "${UNIT_TEST_PROTOCOL_SHARDS}" ]]; then
    shards="${UNIT_TEST_PROTOCOL_SHARDS}"
  else
    if command -v nproc >/dev/null 2>&1; then
      shards="$(nproc)"
      echo -e "${GRN}Protocol shards:${RST} auto-detected ${shards} from nproc (override with UNIT_TEST_PROTOCOL_SHARDS)" >&2
    else
      shards="${DEFAULT_PROTOCOL_SHARDS}"
      echo -e "${GRN}Protocol shards:${RST} nproc not found; defaulting to ${shards} (override with UNIT_TEST_PROTOCOL_SHARDS)" >&2
    fi
  fi

  if ! [[ "${shards}" =~ ^[0-9]+$ ]]; then
    echo -e "${YLW}Invalid UNIT_TEST_PROTOCOL_SHARDS='${shards}', using ${DEFAULT_PROTOCOL_SHARDS}${RST}" >&2
    shards="${DEFAULT_PROTOCOL_SHARDS}"
  fi
  if [[ "${shards}" -lt "${MIN_PROTOCOL_SHARDS}" ]]; then
    echo -e "${YLW}Protocol shards=${shards} is too low, clamping to ${MIN_PROTOCOL_SHARDS}${RST}" >&2
    shards="${MIN_PROTOCOL_SHARDS}"
  fi
  if [[ "${shards}" -gt "${MAX_PROTOCOL_SHARDS}" ]]; then
    echo -e "${YLW}Protocol shards=${shards} is too high, clamping to ${MAX_PROTOCOL_SHARDS}${RST}" >&2
    shards="${MAX_PROTOCOL_SHARDS}"
  fi

  echo "${shards}"
}

HAS_PROTOCOL_PACKAGE=true
if [[ $(echo "${UNIT_TEST_PACKAGES}" | grep -E '\s?\S+protocol\s+') == "" ]]; then
  HAS_PROTOCOL_PACKAGE=false
fi

bg_pids=()

if [[ $HAS_PROTOCOL_PACKAGE == 'false' ]]; then
  # This is the default single-line flow for testing all packages
  # The `else` branch is temporary and will be removed once the `protocol` package runtime is optimized.
  run_test_for_packages "${UNIT_TEST_PACKAGES}" "0" "${UNIT_TEST_COUNT}" "${DEFAULT_TIMEOUT_MINUTES}" "All packages"
else
  # Spawn a process to test all packages except `protocol`
  UNIT_TEST_PACKAGES_FILTERED=$(echo "${UNIT_TEST_PACKAGES}" | tr ' ' '\n' | grep -v '/protocol$' | tr '\n' ' ')
  run_test_for_packages "${UNIT_TEST_PACKAGES_FILTERED}" "0" "${UNIT_TEST_COUNT}" "${DEFAULT_TIMEOUT_MINUTES}" "All packages except 'protocol'" &
  bg_pids+=("$!")

  if [[ $UNIT_TEST_COUNT -eq 1 ]]; then
    protocol_shards="$(resolve_protocol_shards)"

    protocol_list_cmd=(go test -list '^Test')
    if [[ -n "${BUILD_TAGS}" ]]; then
      protocol_list_cmd+=(-tags "${BUILD_TAGS}")
    fi
    protocol_list_cmd+=(github.com/status-im/status-go/internal/protocol)

    mapfile -t protocol_tests < <("${protocol_list_cmd[@]}" | awk '/^Test/ { print $0 }' | sort -u)
    if [[ ${#protocol_tests[@]} -eq 0 ]]; then
      echo -e "${YLW}No protocol tests discovered, running package without sharding${RST}"
      run_test_for_packages github.com/status-im/status-go/internal/protocol "1" 1 "${PROTOCOL_TIMEOUT_MINUTES}" "Only 'protocol' package (fallback, no discovered tests)" &
      bg_pids+=("$!")
    else
      echo -e "${GRN}Protocol sharding:${RST} ${#protocol_tests[@]} top-level tests across ${protocol_shards} shards"
      declare -a protocol_shard_filters
      for ((i=0; i<protocol_shards; i++)); do
        protocol_shard_filters[i]=""
      done

      for ((i=0; i<${#protocol_tests[@]}; i++)); do
        shard_idx=$((i % protocol_shards))
        if [[ -z "${protocol_shard_filters[$shard_idx]}" ]]; then
          protocol_shard_filters[$shard_idx]="${protocol_tests[$i]}"
        else
          protocol_shard_filters[$shard_idx]="${protocol_shard_filters[$shard_idx]}|${protocol_tests[$i]}"
        fi
      done

      protocol_iteration=1
      for ((i=0; i<protocol_shards; i++)); do
        if [[ -z "${protocol_shard_filters[$i]}" ]]; then
          continue
        fi
        protocol_filter="^(${protocol_shard_filters[$i]})$"
        run_test_for_packages github.com/status-im/status-go/internal/protocol "${protocol_iteration}" 1 "${PROTOCOL_TIMEOUT_MINUTES}" "Only 'protocol' package (shard $((i + 1))/${protocol_shards})" "${protocol_filter}" &
        bg_pids+=("$!")
        protocol_iteration=$((protocol_iteration + 1))
      done
    fi
  else
    # Spawn separate processes to run `protocol` package for each nightly iteration
    for ((i=1; i<=UNIT_TEST_COUNT; i++)); do
      run_test_for_packages github.com/status-im/status-go/internal/protocol "${i}" 1 "${PROTOCOL_TIMEOUT_MINUTES}" "Only 'protocol' package" &
      bg_pids+=("$!")
    done
  fi
fi

# Wait for jobs and handle failfast
if [[ $UNIT_TEST_FAILFAST == 'true' ]]; then
  while [[ ${#bg_pids[@]} -gt 0 ]]; do
    wait -n
    exit_code=$?
    if [[ $exit_code -ne 0 ]]; then
      # Kill all remaining background jobs
      for kill_pid in "${bg_pids[@]}"; do
        kill "$kill_pid" 2>/dev/null
      done
      echo -e "${RED}Failfast: Stopping all tests due to failure${RST}"
      exit $exit_code
    fi
    # Remove finished jobs from bg_pids
    remaining_bg_pids=()
    for pid in "${bg_pids[@]}"; do
      if kill -0 "$pid" 2>/dev/null; then
        remaining_bg_pids+=("$pid")
      fi
    done
    bg_pids=("${remaining_bg_pids[@]}")
  done
else
  wait
fi

if [[ "${UNIT_TEST_DRY_RUN}" == 'true' ]]; then
  echo -e "${GRN}Dry run finished${RST}"
  exit 0
fi

# When running in PRs (count=1), early exit if any test failed.
# When running nightly (count>1), generate test stats ant coverage reports anyway.
if [[ $UNIT_TEST_COUNT -eq 1 ]]; then
  for exit_code_file in "${GIT_ROOT}"/exit_code_*.txt; do
    read exit_code < "${exit_code_file}"
    if [[ "${exit_code}" -ne 0 ]]; then
      echo -e "${RED}Testing failed${RST}, exit code: ${exit_code}"
      exit ${exit_code}
    fi
  done
fi

# Gather test coverage results
merged_coverage_report="coverage_merged.out"
coverage_reports=$(find . -iname "*.coverage.out")
rm -f ${merged_coverage_report}

echo -e "${GRN}Gathering test coverage results: ${RST} output: ${merged_coverage_report}, input: ${coverage_reports}"
echo $coverage_reports | xargs go tool gocovmerge > ${merged_coverage_report}

# Generate HTML coverage report
convert_coverage_to_html ${merged_coverage_report} "test-coverage.html"

if [[ $UNIT_TEST_REPORT_CODECOV == 'true' ]]; then
  report_to_codecov "report_*.xml" ${merged_coverage_report} "unit"
fi

# Generate report with test stats
shopt -s globstar nullglob # Enable recursive globbing
if [[ "${UNIT_TEST_COUNT}" -gt 1 ]]; then
  for exit_code_file in "${GIT_ROOT}"/**/exit_code_*.txt; do
    read exit_code < "${exit_code_file}"
    if [[ "${exit_code}" -ne 0 ]]; then
      echo -e "${GRN}Generating test stats${RST}, exit code: ${exit_code}"
      mkdir -p "${GIT_ROOT}/reports"
      "${GIT_ROOT}/scripts/test_stats.py" | tee "${GIT_ROOT}/reports/test_stats.txt"
      exit ${exit_code}
    fi
  done
fi

echo -e "${GRN}Testing finished${RST}"