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

if [ "$#" -ne 1 ]; then
  echo "usage: $0 <fixture-name>" >&2
  echo "example: $0 coder_100" >&2
  exit 2
fi

FIXTURE_NAME="$1"
REPLAYSERVE_ROOT="${REPLAYSERVE_ROOT:-$(cd "$(dirname "$0")/.." && pwd)}"
FRONTIER_ROOT="${FRONTIER_ROOT:-/tmp/toc-llm-sim-research/Frontier}"
if [ -z "${PYTHON_BIN:-}" ] && [ -x "$REPLAYSERVE_ROOT/.venv/bin/python" ]; then
  PYTHON_BIN="$REPLAYSERVE_ROOT/.venv/bin/python"
else
  PYTHON_BIN="${PYTHON_BIN:-python3}"
fi
PYTHON_DEPS_DIR="${PYTHON_DEPS_DIR:-$REPLAYSERVE_ROOT/.deps/python}"
RUN_ROOT="${RUN_ROOT:-$REPLAYSERVE_ROOT/runs/rs1}"
RUN_DIR="${RUN_ROOT}/${FIXTURE_NAME}"
FIXTURE_DIR="${REPLAYSERVE_ROOT}/traces/fixtures/${FIXTURE_NAME}"
TRACE_FILE="${FIXTURE_DIR}/frontier.csv"
SIDECAR_FILE="${FIXTURE_DIR}/sidecar.jsonl"
METRICS_ROOT="${RUN_DIR}/frontier_metrics"
RUN_ID="rs1_${FIXTURE_NAME}"

if [ ! -d "$FRONTIER_ROOT" ]; then
  echo "ERROR: Frontier root does not exist: $FRONTIER_ROOT" >&2
  exit 2
fi
if [ ! -f "$TRACE_FILE" ]; then
  echo "ERROR: fixture trace does not exist: $TRACE_FILE" >&2
  exit 2
fi
if [ ! -f "$SIDECAR_FILE" ]; then
  echo "ERROR: fixture sidecar does not exist: $SIDECAR_FILE" >&2
  exit 2
fi

mkdir -p "$RUN_DIR" "$METRICS_ROOT"

if [ -d "$PYTHON_DEPS_DIR" ]; then
  export PYTHONPATH="$PYTHON_DEPS_DIR:$FRONTIER_ROOT${PYTHONPATH:+:$PYTHONPATH}"
else
  export PYTHONPATH="$FRONTIER_ROOT${PYTHONPATH:+:$PYTHONPATH}"
fi
export WANDB_DISABLED=true
export VIDUR_DISABLE_WANDB=1
export FRONTIER_LOG_LEVEL="${FRONTIER_LOG_LEVEL:-info}"
export PYTHONDONTWRITEBYTECODE=1

CMD=(
  "$PYTHON_BIN" -m frontier.main
  --simulation_mode online
  --sys_arch co-location
  --cc_backend_config_type analytical
  --cluster_config_num_replicas 1
  --cluster_scheduler_config_type sticky_round_robin
  --replica_config_model_name Qwen/Qwen3-32B
  --replica_config_device a800
  --replica_config_network_device a800_dgx
  --replica_config_attn_tensor_parallel_size 2
  --replica_config_num_pipeline_stages 1
  --replica_config_attn_data_parallel_size 1
  --replica_scheduler_config_type vllm_v1
  --decode_cuda_graph_mode full_decode_only
  --vllm_v1_scheduler_config_batch_size_cap 128
  --vllm_v1_scheduler_config_max_tokens_in_batch 32768
  --vllm_v1_scheduler_config_long_prefill_token_threshold 64
  --vllm_v1_scheduler_config_block_size 16
  --vllm_v1_scheduler_config_num_blocks_mode memory_planner
  --vllm_v1_scheduler_config_gpu_memory_utilization 0.9
  --vllm_v1_scheduler_config_non_kv_cache_overhead_bytes 0
  --vllm_v1_scheduler_config_enable_prefix_caching
  --vllm_v1_scheduler_config_enable_chunked_prefill
  --request_generator_config_type trace_replay
  --trace_request_generator_config_trace_file "$TRACE_FILE"
  --trace_request_generator_config_max_tokens 32768
  --random_forrest_execution_time_predictor_config_enable_dummy_mode
  --random_forrest_execution_time_predictor_config_dummy_execution_time_ms 1.0
  --metrics_config_output_dir "$METRICS_ROOT"
  --metrics_config_run_id "$RUN_ID"
  --metrics_config_write_metrics
  --metrics_config_store_request_metrics
  --metrics_config_store_batch_metrics
  --metrics_config_store_token_completion_metrics
  --metrics_config_store_utilization_metrics
  --no-metrics_config_store_plots
  --no-metrics_config_enable_chrome_trace
  --no-metrics_config_write_json_trace
  --no-metrics_config_store_frontier_stage_batch_ledger
)

{
  printf 'cd %q\n' "$FRONTIER_ROOT"
  printf 'export PYTHONPATH=%q\n' "$PYTHONPATH"
  printf 'export WANDB_DISABLED=%q\n' "$WANDB_DISABLED"
  printf 'export VIDUR_DISABLE_WANDB=%q\n' "$VIDUR_DISABLE_WANDB"
  printf 'export FRONTIER_LOG_LEVEL=%q\n' "$FRONTIER_LOG_LEVEL"
  printf 'export PYTHONDONTWRITEBYTECODE=%q\n' "$PYTHONDONTWRITEBYTECODE"
  printf 'command='
  printf '%q ' "${CMD[@]}"
  printf '\n'
} > "$RUN_DIR/command.txt"

{
  printf 'fixture_name=%s\n' "$FIXTURE_NAME"
  printf 'replayserve_root=%s\n' "$REPLAYSERVE_ROOT"
  printf 'frontier_root=%s\n' "$FRONTIER_ROOT"
  printf 'python_deps_dir=%s\n' "$PYTHON_DEPS_DIR"
  printf 'trace_file=%s\n' "$TRACE_FILE"
  printf 'sidecar_file=%s\n' "$SIDECAR_FILE"
  printf 'run_dir=%s\n' "$RUN_DIR"
  printf 'metrics_root=%s\n' "$METRICS_ROOT"
  printf 'run_id=%s\n' "$RUN_ID"
  printf 'frontier_head=%s\n' "$(git -C "$FRONTIER_ROOT" rev-parse HEAD)"
} > "$RUN_DIR/env.txt"

START_EPOCH="$(date +%s)"
printf '%s\n' "$START_EPOCH" > "$RUN_DIR/start_epoch.txt"

set +e
(
  cd "$FRONTIER_ROOT"
  "${CMD[@]}"
) >"$RUN_DIR/stdout.log" 2>"$RUN_DIR/stderr.log"
EXIT_CODE=$?
set -e

END_EPOCH="$(date +%s)"
printf '%s\n' "$END_EPOCH" > "$RUN_DIR/end_epoch.txt"
printf '%s\n' "$EXIT_CODE" > "$RUN_DIR/exit_code.txt"
printf '%s\n' "$((END_EPOCH - START_EPOCH))" > "$RUN_DIR/runtime_seconds.txt"

if [ "$EXIT_CODE" -eq 0 ]; then
  "$PYTHON_BIN" "$REPLAYSERVE_ROOT/tools/postprocess_frontier_smoke.py" \
    --run-dir "$RUN_DIR" \
    --fixture-dir "$FIXTURE_DIR"
fi

exit "$EXIT_CODE"