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

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
SCRIPT_PATH="$SCRIPT_DIR/$(basename "${BASH_SOURCE[0]}")"
ROOT_DIR="$(cd "$SCRIPT_DIR/../.." && pwd)"
ANNOTATION_SCRIPT="${ANNOTATION_SCRIPT:-$ROOT_DIR/qwen/generate_multi_event_annotations.py}"
PYTHON_BIN="${PYTHON_BIN:-python}"
INPUT_PATH="${1:-$ROOT_DIR/examples}"
OUTPUT_PATH="${2:-$ROOT_DIR/annotations/qwen/multi_event/generated_multi_event.json}"
MODEL_PATH="${MODEL_PATH:-$ROOT_DIR/checkpoints}"
MAX_PARALLEL_GPUS="${MAX_PARALLEL_GPUS:-8}"
ANNOTATION_WORKERS="${ANNOTATION_WORKERS:-auto}"
KEEP_PARTIAL_OUTPUTS="${KEEP_PARTIAL_OUTPUTS:-0}"
RUN_IN_BACKGROUND="${RUN_IN_BACKGROUND:-1}"
LOG_ROOT="${LOG_ROOT:-$ROOT_DIR/logs/qwen/multi_event}"
RUN_ID="${RUN_ID:-}"
RUN_DIR="${RUN_DIR:-}"
MAIN_LOG="${MAIN_LOG:-}"

EXTRA_ARGS=()
if [ "$#" -gt 2 ]; then
  EXTRA_ARGS=("${@:3}")
fi

contains_arg() {
  local needle="$1"
  shift
  local arg
  for arg in "$@"; do
    if [ "$arg" = "$needle" ]; then
      return 0
    fi
  done
  return 1
}

require_integer() {
  local name="$1"
  local value="$2"
  if [[ ! "$value" =~ ^[0-9]+$ ]]; then
    echo "$name must be a non-negative integer, got: $value" >&2
    exit 1
  fi
}

init_run_dir() {
  mkdir -p "$LOG_ROOT"
  if [ -z "$RUN_ID" ]; then
    RUN_ID="$(date '+%Y%m%d_%H%M%S')_$$"
  fi
  if [ -z "$RUN_DIR" ]; then
    RUN_DIR="$LOG_ROOT/$RUN_ID"
  fi
  mkdir -p "$RUN_DIR"
  if [ -z "$MAIN_LOG" ]; then
    MAIN_LOG="$RUN_DIR/main.log"
  fi
}

if [ -z "${ANNOTATION_BACKGROUND_CHILD:-}" ] \
  && [ "$RUN_IN_BACKGROUND" = "1" ] \
  && ! contains_arg --dry-run "${EXTRA_ARGS[@]}" \
  && ! contains_arg --print-uuids-json "${EXTRA_ARGS[@]}"; then
  init_run_dir
  export ANNOTATION_BACKGROUND_CHILD=1
  export RUN_ID
  export RUN_DIR
  export MAIN_LOG
  nohup bash "$SCRIPT_PATH" "$@" >>"$MAIN_LOG" 2>&1 &
  PID="$!"
  printf '%s\n' "$PID" >"$RUN_DIR/pid"
  printf 'Started background annotation job.\nPID: %s\nRun dir: %s\nMain log: %s\nOutput: %s\n' \
    "$PID" "$RUN_DIR" "$MAIN_LOG" "$OUTPUT_PATH"
  exit 0
fi

init_run_dir

echo "[$(date '+%Y-%m-%d %H:%M:%S')] Starting multi-event annotation run"
echo "run_dir=$RUN_DIR"
echo "output_path=$OUTPUT_PATH"
echo "input_path=$INPUT_PATH"
echo "model_path=$MODEL_PATH"

discover_gpu_ids() {
  if [ -n "${GPU_IDS:-}" ]; then
    printf '%s\n' "$GPU_IDS"
    return 0
  fi
  if [ -n "${CUDA_VISIBLE_DEVICES:-}" ]; then
    printf '%s\n' "$CUDA_VISIBLE_DEVICES"
    return 0
  fi
  if command -v nvidia-smi >/dev/null 2>&1; then
    nvidia-smi --query-gpu=index --format=csv,noheader
    return 0
  fi
}

normalize_gpu_ids() {
  "$PYTHON_BIN" - "$@" <<'PY'
import re
import sys

tokens = []
for raw in sys.argv[1:]:
    parts = re.split(r"[\s,]+", raw.strip())
    for part in parts:
        part = part.strip()
        if part:
            tokens.append(part)
for token in tokens:
    print(token)
PY
}

readarray -t GPU_IDS_ARRAY < <(normalize_gpu_ids "$(discover_gpu_ids)")
GPU_COUNT="${#GPU_IDS_ARRAY[@]}"

require_integer "MAX_PARALLEL_GPUS" "$MAX_PARALLEL_GPUS"
if [ "$ANNOTATION_WORKERS" != "auto" ]; then
  require_integer "ANNOTATION_WORKERS" "$ANNOTATION_WORKERS"
fi

if [ ! -f "$ANNOTATION_SCRIPT" ]; then
  echo "Annotation script not found: $ANNOTATION_SCRIPT" >&2
  exit 1
fi

if contains_arg --dry-run "${EXTRA_ARGS[@]}" || contains_arg --print-uuids-json "${EXTRA_ARGS[@]}"; then
  exec "$PYTHON_BIN" "$ANNOTATION_SCRIPT" \
    --input "$INPUT_PATH" \
    --output "$OUTPUT_PATH" \
    --model-path "$MODEL_PATH" \
    --recursive \
    "${EXTRA_ARGS[@]}"
fi

UUID_JSON="$("$PYTHON_BIN" "$ANNOTATION_SCRIPT" \
  --input "$INPUT_PATH" \
  --output "$OUTPUT_PATH" \
  --model-path "$MODEL_PATH" \
  --recursive \
  --print-uuids-json \
  "${EXTRA_ARGS[@]}")"

UUID_COUNT="$(printf '%s' "$UUID_JSON" | "$PYTHON_BIN" -c 'import json,sys; print(len(json.load(sys.stdin)))')"

if [ "$UUID_COUNT" -eq 0 ]; then
  echo "No pending UUIDs to annotate."
  exit 0
fi

if [ "$ANNOTATION_WORKERS" = "auto" ]; then
  if [ "$GPU_COUNT" -gt 0 ]; then
    WORKER_COUNT="$GPU_COUNT"
  else
    WORKER_COUNT=1
  fi
else
  WORKER_COUNT="$ANNOTATION_WORKERS"
fi

if [ "$WORKER_COUNT" -gt "$MAX_PARALLEL_GPUS" ]; then
  WORKER_COUNT="$MAX_PARALLEL_GPUS"
fi
if [ "$GPU_COUNT" -gt 0 ] && [ "$WORKER_COUNT" -gt "$GPU_COUNT" ]; then
  WORKER_COUNT="$GPU_COUNT"
fi
if [ "$WORKER_COUNT" -gt "$UUID_COUNT" ]; then
  WORKER_COUNT="$UUID_COUNT"
fi
if [ "$WORKER_COUNT" -lt 1 ]; then
  WORKER_COUNT=1
fi

if [ "$WORKER_COUNT" -le 1 ]; then
  exec "$PYTHON_BIN" "$ANNOTATION_SCRIPT" \
    --input "$INPUT_PATH" \
    --output "$OUTPUT_PATH" \
    --model-path "$MODEL_PATH" \
    --recursive \
    "${EXTRA_ARGS[@]}"
fi

HAS_OVERWRITE=0
HAS_RESUME=0
if contains_arg --overwrite "${EXTRA_ARGS[@]}"; then
  HAS_OVERWRITE=1
fi
if contains_arg --resume "${EXTRA_ARGS[@]}"; then
  HAS_RESUME=1
fi

if [ -e "$OUTPUT_PATH" ] && [ "$HAS_OVERWRITE" -eq 0 ] && [ "$HAS_RESUME" -eq 0 ]; then
  echo "Output already exists: $OUTPUT_PATH. Use --overwrite or --resume." >&2
  exit 1
fi

PARTIAL_DIR="$RUN_DIR/partials"
mkdir -p "$PARTIAL_DIR"
cleanup() {
  if [ "${KEEP_PARTIAL_OUTPUTS}" != "1" ] && [ -d "$PARTIAL_DIR" ]; then
    rm -rf "$PARTIAL_DIR"
  fi
}
trap cleanup EXIT

UUID_JSON_PATH="$PARTIAL_DIR/uuids.json"
printf '%s\n' "$UUID_JSON" >"$UUID_JSON_PATH"

"$PYTHON_BIN" - "$UUID_JSON_PATH" "$PARTIAL_DIR" "$WORKER_COUNT" <<'PY'
import json
import sys
from pathlib import Path

uuid_path = Path(sys.argv[1])
out_dir = Path(sys.argv[2])
num_workers = int(sys.argv[3])
uuids = json.loads(uuid_path.read_text(encoding="utf-8"))
buckets = [[] for _ in range(num_workers)]
for idx, uuid in enumerate(uuids):
    buckets[idx % num_workers].append(uuid)
for idx, bucket in enumerate(buckets):
    worker_file = out_dir / f"worker_{idx}.txt"
    worker_file.write_text(
        ("\n".join(bucket) + "\n") if bucket else "",
        encoding="utf-8",
    )
PY

PIDS=()
LOGS=()
PARTIAL_OUTPUTS=()

for ((worker_idx = 0; worker_idx < WORKER_COUNT; worker_idx++)); do
  UUID_FILE="$PARTIAL_DIR/worker_${worker_idx}.txt"
  if [ ! -s "$UUID_FILE" ]; then
    continue
  fi
  PARTIAL_OUTPUT="$PARTIAL_DIR/part_${worker_idx}.json"
  LOG_PATH="$PARTIAL_DIR/worker_${worker_idx}.log"
  PARTIAL_OUTPUTS+=("$PARTIAL_OUTPUT")
  LOGS+=("$LOG_PATH")

  (
    if [ "$GPU_COUNT" -gt 0 ]; then
      export CUDA_VISIBLE_DEVICES="${GPU_IDS_ARRAY[$worker_idx]}"
    fi
    exec "$PYTHON_BIN" "$ANNOTATION_SCRIPT" \
      --input "$INPUT_PATH" \
      --output "$PARTIAL_OUTPUT" \
      --model-path "$MODEL_PATH" \
      --recursive \
      --uuids-file "$UUID_FILE" \
      --overwrite \
      "${EXTRA_ARGS[@]}"
  ) >"$LOG_PATH" 2>&1 &
  PIDS+=("$!")
done

STATUS=0
for pid in "${PIDS[@]}"; do
  if ! wait "$pid"; then
    STATUS=1
  fi
done

if [ "$STATUS" -ne 0 ]; then
  echo "One or more annotation workers failed. Logs are in $RUN_DIR" >&2
  for log_path in "${LOGS[@]}"; do
    if [ -f "$log_path" ]; then
      echo "===== $log_path =====" >&2
      tail -n 40 "$log_path" >&2 || true
    fi
  done
  exit "$STATUS"
fi

"$PYTHON_BIN" - "$OUTPUT_PATH" "$HAS_RESUME" "${PARTIAL_OUTPUTS[@]}" <<'PY'
import json
import sys
from pathlib import Path

output_path = Path(sys.argv[1])
resume = bool(int(sys.argv[2]))
partial_paths = [Path(arg) for arg in sys.argv[3:]]

videos_by_uuid = {}
if resume and output_path.exists():
    payload = json.loads(output_path.read_text(encoding="utf-8"))
    for item in payload.get("videos", []):
        if isinstance(item, dict) and "uuid" in item:
            videos_by_uuid[item["uuid"]] = item

for path in partial_paths:
    if not path.exists():
        continue
    payload = json.loads(path.read_text(encoding="utf-8"))
    for item in payload.get("videos", []):
        if isinstance(item, dict) and "uuid" in item:
            videos_by_uuid[item["uuid"]] = item

ordered = [videos_by_uuid[uuid] for uuid in sorted(videos_by_uuid)]
merged = {
    "total_videos": len(ordered),
    "videos": ordered,
}
output_path.parent.mkdir(parents=True, exist_ok=True)
output_path.write_text(
    json.dumps(merged, ensure_ascii=False, indent=2) + "\n",
    encoding="utf-8",
)
PY

"$PYTHON_BIN" - "$OUTPUT_PATH" "$PARTIAL_DIR" <<'PY'
import sys
from pathlib import Path

output_path = Path(sys.argv[1])
partial_dir = Path(sys.argv[2])
partial_logs = sorted(partial_dir.glob("part_*.manual_camera_desc_mismatches.txt"))
if not partial_logs:
    raise SystemExit(0)

merged_path = output_path.with_name(f"{output_path.stem}.manual_camera_desc_mismatches.txt")
chunks = []
for path in partial_logs:
    text = path.read_text(encoding="utf-8").strip()
    if text:
        chunks.append(text)
if chunks:
    merged_path.write_text("\n\n".join(chunks) + "\n", encoding="utf-8")
    print(f"Wrote manual camera mismatch log to {merged_path}")
PY

if [ "${KEEP_PARTIAL_OUTPUTS}" = "1" ]; then
  echo "Partial worker outputs kept in $PARTIAL_DIR"
fi
echo "[$(date '+%Y-%m-%d %H:%M:%S')] Completed multi-event annotation run"