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

action="${1:?action is required}"
root="${STEM_REPO:-$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)}"
profile="${STEM_PROFILE:-native-amd64}"
target="${STEM_TARGET:-native}"
selector="${STEM_TASK:-task3}"
state="${STEM_STATE_DIR:-$root/.stem/instances/local}"

tasks=(
  task01_first_fit_split
  task02_address_order_coalesce
  task03_freertos_heap4
)

resolve_task() {
  local value="${1,,}" digits number stem
  for stem in "${tasks[@]}"; do
    if [[ "$value" == "$stem" ]]; then
      printf '%s\n' "$stem"
      return
    fi
  done
  if [[ "$value" =~ ^(task|t)?[-_]?0*([1-3])($|[-_].*) ]]; then
    digits="${BASH_REMATCH[2]}"
    number=$((10#$digits))
    printf '%s\n' "${tasks[number - 1]}"
    return
  fi
  printf 'Unknown task selector: %s\n' "$1" >&2
  exit 2
}

task_stem="$(resolve_task "$selector")"
task_digits="${task_stem#task}"
task_digits="${task_digits%%_*}"
task_number=$((10#$task_digits))
task_key="$(printf 'task%02d' "$task_number")"
source_file="$root/src/tasks/$task_stem.c"
artifact_dir="$root/.stem/artifacts/$profile/$target/$task_key"
mkdir -p "$state" "$artifact_dir"

freertos_kernel_path() {
  if [[ -n "${FREERTOS_KERNEL_PATH:-}" ]]; then
    printf '%s\n' "$FREERTOS_KERNEL_PATH"
  elif [[ -f /opt/FreeRTOS-Kernel/portable/MemMang/heap_4.c ]]; then
    printf '%s\n' /opt/FreeRTOS-Kernel
  else
    printf '%s\n' "$root/../lab-rv32i-freertos-task-stack-vector/vendor/FreeRTOS-Kernel"
  fi
}

verify_kernel() {
  local kernel="$1"
  test -f "$kernel/portable/MemMang/heap_4.c"
  grep -F 'FreeRTOS Kernel V11.3.0' "$kernel/include/FreeRTOS.h" >/dev/null
}

freertos_rp2350_path() {
  if [[ -n "${FREERTOS_RP2350_PATH:-}" ]]; then
    printf '%s\n' "$FREERTOS_RP2350_PATH"
  else
    printf '%s\n' /opt/FreeRTOS-RP2350
  fi
}

verify_rp2350_kernel() {
  local kernel="$1"
  test -f "$kernel/portable/ThirdParty/GCC/RP2350_RISC-V/portASM.S"
  grep -F '<DEVELOPMENT BRANCH>' "$kernel/include/FreeRTOS.h" >/dev/null
}

copy_artifact() {
  local path="$1"
  [[ -f "$path" ]] && cp -f "$path" "$artifact_dir/"
}

build_native() {
  local mode="${1:-normal}" kernel build_dir output cc="${CC:-cc}"
  local flags=(-std=c11 -O0 -g3 -fno-omit-frame-pointer -fno-inline
               -fno-optimize-sibling-calls -Wall -Wextra -Wpedantic
               -DHOST_PRINTF=1 -I"$root/src/tasks")
  build_dir="$state/build/native/$task_stem"
  output="$build_dir/prog"
  if [[ "$mode" == sanitizer ]]; then
    flags+=(-fsanitize=address,undefined)
  fi
  mkdir -p "$build_dir"
  if [[ "$task_number" -eq 3 ]]; then
    kernel="$(freertos_kernel_path)"
    verify_kernel "$kernel"
    flags+=(-DHEAP4_HOST_ADAPTER=1 -I"$root/include/host-shim")
    "$cc" "${flags[@]}" "$source_file" \
      "$kernel/portable/MemMang/heap_4.c" -o "$output"
  else
    "$cc" "${flags[@]}" "$source_file" -o "$output"
  fi
  objdump -dS "$output" > "$build_dir/prog.lst"
  cp -f "$output" "$artifact_dir/prog"
  copy_artifact "$build_dir/prog.lst"
  printf '%s\n' "$output"
}

rv_env_root() {
  if [[ -n "${RV_ENV_ROOT:-}" ]]; then
    printf '%s\n' "$RV_ENV_ROOT"
  elif [[ -d /opt/rv-env ]]; then
    printf '%s\n' /opt/rv-env
  else
    printf '%s\n' /opt/stem
  fi
}

build_hazard3() {
  local env_root kernel build_dir elf
  env_root="$(rv_env_root)"
  kernel="$(freertos_kernel_path)"
  verify_kernel "$kernel"
  build_dir="$state/build/hazard3"
  make -B -C "$root" "task$task_number" \
    "BUILD_ROOT=$build_dir" "RV_ENV_ROOT=$env_root" \
    "FREERTOS_KERNEL_PATH=$kernel" PROFILE=observe >&2 || return
  elf="$build_dir/$task_stem/prog.elf"
  test -s "$elf"
  copy_artifact "$elf"
  copy_artifact "${elf%.elf}.bin"
  copy_artifact "${elf%.elf}.lst"
  copy_artifact "${elf%.elf}.map"
  printf '%s\n' "$elf"
}

hazard3_tb() {
  local env_root tb
  env_root="$(rv_env_root)"
  tb="$env_root/vendor/Hazard3/test/sim/tb_verilator/tb"
  if [[ ! -x "$tb" ]]; then
    make -C "$(dirname -- "$tb")" tb >&2
  fi
  printf '%s\n' "$tb"
}

build_rp2350() {
  local binary_type="${1:-no_flash}" build_kind build_dir elf kernel rp_kernel
  local cmake_args
  case "$binary_type" in
    no_flash) build_kind=ram ;;
    default) build_kind=flash ;;
    *) printf 'Unsupported RP2350 binary type: %s\n' "$binary_type" >&2; return 2 ;;
  esac
  kernel="$(freertos_kernel_path)"
  verify_kernel "$kernel"
  rp_kernel="$(freertos_rp2350_path)"
  verify_rp2350_kernel "$rp_kernel"
  build_dir="$state/build/rp2350-rv-$build_kind"
  cmake_args=(
    -S "$root" -B "$build_dir"
    -DPICO_PLATFORM=rp2350-riscv
    -DPICO_BOARD="${PICO_BOARD:-pico2_w}"
    -DPICO_GCC_TRIPLE=riscv32-pico-elf
    -DHEAP4_UPSTREAM_ROOT="$kernel"
    -DFREERTOS_RP2350_PATH="$rp_kernel"
    -DHEAP4_PROFILE=observe
    -DHEAP4_BINARY_TYPE="$binary_type"
    -DCMAKE_BUILD_TYPE=Debug
    -DCMAKE_EXPORT_COMPILE_COMMANDS=ON
  )
  if command -v ninja >/dev/null 2>&1; then
    cmake_args+=(-G Ninja)
  else
    cmake_args+=(-G "Unix Makefiles")
  fi
  if [[ -d /opt/picotool/picotool ]]; then
    cmake_args+=(-Dpicotool_DIR=/opt/picotool/picotool)
  fi
  cmake "${cmake_args[@]}" >&2
  cmake --build "$build_dir" --target "$task_stem" >&2
  elf="$build_dir/$task_stem.elf"
  test -s "$elf"
  riscv64-unknown-elf-objdump -dS -M no-aliases,numeric "$elf" \
    > "$build_dir/$task_stem.lst"
  copy_artifact "$elf"
  copy_artifact "$build_dir/$task_stem.bin"
  copy_artifact "$build_dir/$task_stem.hex"
  copy_artifact "$build_dir/$task_stem.uf2"
  copy_artifact "$build_dir/$task_stem.dis"
  copy_artifact "$build_dir/$task_stem.lst"
  copy_artifact "$build_dir/$task_stem.elf.map"
  printf '%s\n' "$elf"
}

openocd_target="${STEM_OPENOCD_TARGET:-target/rp2350.cfg}"
gdb_port="${GDB_PORT:-3333}"
gdb_target="${GDB_TARGET:-127.0.0.1:$gdb_port}"
openocd_pid=""

stop_owned_openocd() {
  if [[ -n "$openocd_pid" ]]; then
    kill "$openocd_pid" 2>/dev/null || true
    wait "$openocd_pid" 2>/dev/null || true
  fi
}

start_openocd() {
  local log="$state/openocd-rp2350.log" attempt
  if nc -z 127.0.0.1 "$gdb_port" 2>/dev/null; then
    return
  fi
  openocd -f interface/cmsis-dap.cfg \
    -c "adapter speed ${STEM_ADAPTER_KHZ:-5000}" \
    -c 'set USE_CORE { rv0 }' -c 'set USE_SMP 0' \
    -c "set FLASHSIZE ${STEM_RP2350_FLASH_SIZE:-0x400000}" \
    -f "$openocd_target" >"$log" 2>&1 &
  openocd_pid=$!
  for attempt in $(seq 1 80); do
    if nc -z 127.0.0.1 "$gdb_port" 2>/dev/null; then
      return
    fi
    if ! kill -0 "$openocd_pid" 2>/dev/null; then
      sed -n '1,240p' "$log" >&2
      return 1
    fi
    sleep 0.1
  done
  printf 'OpenOCD did not open GDB port %s\n' "$gdb_port" >&2
  return 1
}

flash_rp2350() {
  local elf="$1" log="$state/openocd-rp2350-flash.log"
  if nc -z 127.0.0.1 "$gdb_port" 2>/dev/null; then
    printf 'GDB port %s is busy; stop the current debug session before deploy.\n' "$gdb_port" >&2
    return 1
  fi
  if ! openocd -f interface/cmsis-dap.cfg \
      -c "adapter speed ${STEM_ADAPTER_KHZ:-5000}" \
      -c 'set USE_CORE { cm0 rv0 }' -c 'set USE_SMP 0' \
      -c "set FLASHSIZE ${STEM_RP2350_FLASH_SIZE:-0x400000}" \
      -f "$openocd_target" -c init \
      -c 'if {[target_present rp2350.rv0]} { targets rp2350.rv0 } else { targets rp2350.cm0 }' \
      -c 'reset halt' -c "flash write_image erase $elf" \
      -c "verify_image $elf" -c 'reset run' -c shutdown >"$log" 2>&1; then
    sed -n '1,260p' "$log" >&2
    return 1
  fi
  grep -E 'wrote [0-9]+ bytes|verified [0-9]+ bytes' "$log" || true
  printf 'rp2350-flash-ok image=%s verify=passed architecture=RISC-V\n' "$elf"
}

hardware_test_task03() {
  local elf="$1" output
  start_openocd
  trap stop_owned_openocd EXIT
  output="$(gdb-multiarch -nx -q -batch "$elf" \
    -ex 'set architecture riscv:rv32' \
    -ex "target extended-remote $gdb_target" \
    -ex 'monitor reset halt' \
    -ex 'hbreak heap4_fragmented_checkpoint' \
    -ex 'hbreak task03_debug_checkpoint' \
    -ex continue \
    -ex 'printf "fragmented blocks=%lu free=%lu largest=%lu\n", g_fragmented_stats.xNumberOfFreeBlocks, g_fragmented_stats.xAvailableHeapSpaceInBytes, g_fragmented_stats.xSizeOfLargestFreeBlockInBytes' \
    -ex continue \
    -ex 'printf "task03 initial=%lu after=%lu final=%lu min=%lu aligned=%d oom=%d hooks=%d asserts=%d pass=%d\n", g_initial_free, g_after_allocations, g_final_free, g_minimum_ever_free, g_allocations_aligned, g_oom_is_null, g_malloc_failed_hooks, g_assert_failures, g_task03_pass' \
    -ex 'printf "frame sp=0x%08lx fp=0x%08lx pc=0x%08lx align=%lu\n", (unsigned long)$sp, (unsigned long)$s0, (unsigned long)$pc, ((unsigned long)$sp & 15)' \
    -ex 'x/96bx ucHeap' \
    -ex 'delete breakpoints' -ex 'monitor halt' -ex disconnect)"
  printf '%s\n' "$output"
  grep -E '^fragmented blocks=2 ' <<<"$output" >/dev/null
  grep -E '^task03 .* aligned=1 oom=1 hooks=1 asserts=0 pass=1$' <<<"$output" >/dev/null
  grep -E '^frame .* align=0$' <<<"$output" >/dev/null
  printf 'rp2350-hardware-ok task=task03 heap_4=upstream load=flash\n'
}

configure_debug_view() {
  export STEM_DEBUG_WINDOW_NAME=heap4
  case "$task_number" in
    1)
      export STEM_DEBUG_INITIAL_BREAK=task01_debug_checkpoint
      export STEM_DEBUG_MEMORY_EXPR='&g_split_arena' STEM_DEBUG_MEMORY_LENGTH=288
      ;;
    2)
      export STEM_DEBUG_INITIAL_BREAK=task02_fragmented_checkpoint
      export STEM_DEBUG_MEMORY_EXPR='&g_coalesce_arena' STEM_DEBUG_MEMORY_LENGTH=160
      ;;
    3)
      export STEM_DEBUG_INITIAL_BREAK=heap4_fragmented_checkpoint
      export STEM_DEBUG_MEMORY_EXPR='ucHeap' STEM_DEBUG_MEMORY_LENGTH=160
      ;;
  esac
}

case "$profile:$target:$action" in
  native-amd64:native:build)
    build_native normal >/dev/null
    ;;
  native-amd64:native:test)
    output="$(build_native sanitizer)"
    ASAN_OPTIONS=detect_leaks=1 UBSAN_OPTIONS=halt_on_error=1 "$output"
    ;;
  native-amd64:native:run)
    output="$(build_native normal)"
    exec "$output"
    ;;
  native-amd64:native:debug)
    output="$(build_native normal)"
    export RV_REPO="$root" RV_EDITOR="${STEM_EDITOR:-nvim}"
    export RV_HOST_ELF="$output" RV_HOST_SRC="$source_file"
    export RV_HOST_LST="$(dirname -- "$output")/prog.lst"
    exec /opt/stem/scripts/host-entry.sh debug "task$task_number"
    ;;
  hazard3-sim:hazard3-baremetal:build)
    build_hazard3 >/dev/null
    ;;
  hazard3-sim:hazard3-baremetal:test|hazard3-sim:hazard3-baremetal:run)
    elf="$(build_hazard3)"
    tb="$(hazard3_tb)"
    exec "$tb" --bin "${elf%.elf}.bin" --cpuret
    ;;
  hazard3-sim:hazard3-baremetal:debug)
    elf="$(build_hazard3)"
    configure_debug_view
    export STEM_DEBUG_BACKEND=hazard3 STEM_DEBUG_ELF="$elf"
    export STEM_DEBUG_BIN="${elf%.elf}.bin" STEM_DEBUG_SRC="$source_file"
    export STEM_DEBUG_LST="${elf%.elf}.lst"
    exec /opt/stem/scripts/card-debug-ui.sh
    ;;
  rp2350:rp2350-rv:build)
    build_rp2350 >/dev/null
    ;;
  rp2350:rp2350-rv:test)
    elf="$(build_rp2350)"
    riscv64-unknown-elf-readelf -h "$elf" | grep -F 'Machine:' | grep -F 'RISC-V'
    if [[ "$task_number" -eq 3 ]]; then
      "$root/tests/check_task03_elf.sh" "$elf"
    fi
    printf 'rp2350-build-ok task=%s binary_type=no_flash\n' "$task_key"
    ;;
  rp2350:rp2350-rv:deploy)
    elf="$(build_rp2350 default)"
    flash_rp2350 "$elf"
    if [[ "$task_number" -eq 3 ]]; then
      hardware_test_task03 "$elf"
    else
      printf 'rp2350-deploy-ok task=%s load=flash persistent=1\n' "$task_key"
    fi
    ;;
  rp2350:rp2350-rv:debug)
    elf="$(build_rp2350)"
    configure_debug_view
    export STEM_DEBUG_BACKEND=rp2350 STEM_DEBUG_ELF="$elf"
    export STEM_DEBUG_SRC="$source_file" STEM_DEBUG_LST="${elf%.elf}.lst"
    exec /opt/stem/scripts/card-debug-ui.sh
    ;;
  *)
    printf 'Unsupported profile/target/action: %s/%s/%s\n' "$profile" "$target" "$action" >&2
    exit 2
    ;;
esac