#!/usr/bin/env bash

set -e

help() {
  echo "Usage: $0 [-e <config|compile|clean>] [-c <string>] [-t <string>] [-j <number>] [-p <string>] [-f]" 1>&2
  echo 1>&2
  echo "  - e - Parameter for esphome command. Default compile. Common alternative is config." 1>&2
  echo "  - c - Component folder name to test. Default *. E.g. '-c logger'." 1>&2
  echo "  - t - Target name to test. Put '-t list' to display all possibilities. E.g. '-t esp32-s2-idf-51'." 1>&2
  echo "  - j - Number of parallel jobs. Default is number of CPU cores." 1>&2
  echo "  - p - Platform filter. E.g. '-p esp32' to test only ESP32 platforms." 1>&2
  echo "  - f - Fail fast. Exit on first failure." 1>&2
  echo "  - b - Build cache directory. E.g. '-b /tmp/esphome_cache'." 1>&2
  exit 1
}

# Parse parameter:
esphome_command="compile"
target_component="*"
num_jobs=$(nproc 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 4)
platform_filter=""
fail_fast=false
build_cache_dir=""

while getopts e:c:t:j:p:b:fh flag
do
    case $flag in
        e) esphome_command=${OPTARG};;
        c) target_component=${OPTARG};;
        t) requested_target_platform=${OPTARG};;
        j) num_jobs=${OPTARG};;
        p) platform_filter=${OPTARG};;
        f) fail_fast=true;;
        b) build_cache_dir=${OPTARG};;
        h) help;;
        \?) help;;
    esac
done

cd "$(dirname "$0")/.."

if ! [ -d "./tests/test_build_components/build" ]; then
  mkdir -p ./tests/test_build_components/build
fi

# Export build cache directory if specified
if [ -n "$build_cache_dir" ]; then
  export PLATFORMIO_BUILD_CACHE_DIR="$build_cache_dir"
  mkdir -p "$build_cache_dir"
  echo "Using build cache directory: $build_cache_dir"
fi

# Track PIDs for parallel execution
pids=()
failed_builds=()
build_count=0
total_builds=0

# Function to wait for jobs and handle failures
wait_for_jobs() {
  local max_jobs=$1
  while [ ${#pids[@]} -ge $max_jobs ]; do
    for i in "${!pids[@]}"; do
      if ! kill -0 "${pids[$i]}" 2>/dev/null; then
        wait "${pids[$i]}"
        exit_code=$?
        if [ $exit_code -ne 0 ]; then
          failed_builds+=("${build_info[$i]}")
          if [ "$fail_fast" = true ]; then
            echo "Build failed, exiting due to fail-fast mode"
            # Kill remaining jobs
            for pid in "${pids[@]}"; do
              kill -TERM "$pid" 2>/dev/null || true
            done
            exit 1
          fi
        fi
        unset pids[$i]
        unset build_info[$i]
        # Reindex arrays
        pids=("${pids[@]}")
        build_info=("${build_info[@]}")
        break
      fi
    done
    sleep 0.1
  done
}

start_esphome() {
  if [ -n "$requested_target_platform" ] && [ "$requested_target_platform" != "$target_platform_with_version" ]; then
    echo "Skipping $target_platform_with_version"
    return
  fi

  # Apply platform filter if specified
  if [ -n "$platform_filter" ] && [[ ! "$target_platform_with_version" =~ ^$platform_filter ]]; then
    echo "Skipping $target_platform_with_version (filtered)"
    return
  fi

  # create dynamic yaml file in `build` folder.
  component_test_file="./tests/test_build_components/build/$target_component.$test_name.$target_platform_with_version.yaml"

  cp $target_platform_file $component_test_file
  if [[ "$OSTYPE" == "darwin"* ]]; then
    # macOS sed is...different
    sed -i '' "s!\$component_test_file!../../.$f!g" $component_test_file
  else
    sed -i "s!\$component_test_file!../../.$f!g" $component_test_file
  fi

  # Start esphome process in background
  build_count=$((build_count + 1))
  echo "> [$build_count/$total_builds] [$target_component] [$test_name] [$target_platform_with_version]"

  (
    # Add compile process limit for ESPHome internal parallelization
    export ESPHOME_COMPILE_PROCESS_LIMIT=2
    
    # For compilation, add a small random delay to reduce thundering herd effect
    # This helps stagger the package installation requests
    if [ "$esphome_command" = "compile" ]; then
      sleep $((RANDOM % 3))
    fi
    
    python3 -m esphome -s component_name $target_component -s component_dir ../../components/$target_component -s test_name $test_name -s target_platform $target_platform $esphome_command $component_test_file
  ) &

  local pid=$!
  pids+=($pid)
  build_info+=("$target_component/$test_name/$target_platform_with_version")

  # Wait if we've reached the job limit
  wait_for_jobs $num_jobs
}

# First pass: count total builds
echo "Calculating total number of builds..."
for f in ./tests/components/$target_component/*.*.yaml; do
  [ -f "$f" ] || continue
  IFS='/' read -r -a folder_name <<< "$f"
  IFS='.' read -r -a file_name <<< "${folder_name[4]}"
  target_platform="${file_name[1]}"
  file_name_parts=${#file_name[@]}

  if [ "$target_platform" = "all" ] || [ $file_name_parts = 2 ]; then
    for target_platform_file in ./tests/test_build_components/build_components_base.*.yaml; do
      IFS='/' read -r -a folder_name <<< "$target_platform_file"
      IFS='.' read -r -a file_name <<< "${folder_name[3]}"
      target_platform="${file_name[1]}"
      target_platform_with_version=${target_platform_file:52}
      target_platform_with_version=${target_platform_with_version%.*}

      if [ -n "$platform_filter" ] && [[ ! "$target_platform_with_version" =~ ^$platform_filter ]]; then
        continue
      fi
      if [ -n "$requested_target_platform" ] && [ "$requested_target_platform" != "$target_platform_with_version" ]; then
        continue
      fi
      total_builds=$((total_builds + 1))
    done
  else
    target_platform_file="./tests/test_build_components/build_components_base.$target_platform.yaml"
    if [ -f "$target_platform_file" ]; then
      for target_platform_file in ./tests/test_build_components/build_components_base.$target_platform*.yaml; do
        target_platform_with_version=${target_platform_file:52}
        target_platform_with_version=${target_platform_with_version%.*}

        if [ -n "$platform_filter" ] && [[ ! "$target_platform_with_version" =~ ^$platform_filter ]]; then
          continue
        fi
        if [ -n "$requested_target_platform" ] && [ "$requested_target_platform" != "$target_platform_with_version" ]; then
          continue
        fi
        total_builds=$((total_builds + 1))
      done
    fi
  fi
done

echo "Total builds to execute: $total_builds with $num_jobs parallel jobs"
echo

# Second pass: execute builds
for f in ./tests/components/$target_component/*.*.yaml; do
  [ -f "$f" ] || continue
  IFS='/' read -r -a folder_name <<< "$f"
  target_component="${folder_name[3]}"

  IFS='.' read -r -a file_name <<< "${folder_name[4]}"
  test_name="${file_name[0]}"
  target_platform="${file_name[1]}"
  file_name_parts=${#file_name[@]}

  if [ "$target_platform" = "all" ] || [ $file_name_parts = 2 ]; then
    # Test has *not* defined a specific target platform. Need to run tests for all possible target platforms.

    for target_platform_file in ./tests/test_build_components/build_components_base.*.yaml; do
      IFS='/' read -r -a folder_name <<< "$target_platform_file"
      IFS='.' read -r -a file_name <<< "${folder_name[3]}"
      target_platform="${file_name[1]}"
      target_platform_with_version=${target_platform_file:52}
      target_platform_with_version=${target_platform_with_version%.*}

      start_esphome
    done

  else
    # Test has defined a specific target platform.

    # Validate we have a base test yaml for selected platform.
    target_platform_file="./tests/test_build_components/build_components_base.$target_platform.yaml"
    if ! [ -f "$target_platform_file" ]; then
      echo "No base test file [./tests/test_build_components/build_components_base.$target_platform.yaml] for component test [$f] found."
      exit 1
    fi

    for target_platform_file in ./tests/test_build_components/build_components_base.$target_platform*.yaml; do
      # trim off "./tests/test_build_components/build_components_base." prefix
      target_platform_with_version=${target_platform_file:52}
      # ...now remove suffix starting with "." leaving just the test target hardware and software platform (possibly with version)
      # For example: "esp32-s3-idf-50"
      target_platform_with_version=${target_platform_with_version%.*}
      start_esphome
    done
  fi
done

# Wait for all remaining jobs
wait_for_jobs 1

echo
echo "============================================"
echo "Build Summary:"
echo "Total builds: $total_builds"
echo "Failed builds: ${#failed_builds[@]}"

if [ ${#failed_builds[@]} -gt 0 ]; then
  echo
  echo "Failed builds:"
  for build in "${failed_builds[@]}"; do
    echo "  - $build"
  done
  exit 1
else
  echo "All builds completed successfully!"
fi