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

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
RUN_SCRIPT_REL="${RUN_SCRIPT_REL:-scripts/infer_ti2vid_one_stage_1024.sh}"
RUN_SCRIPT="${ROOT_DIR}/${RUN_SCRIPT_REL#./}"
RUN_SCRIPT_BASENAME="$(basename "${RUN_SCRIPT_REL}")"
RUN_SCRIPT_TAG="${RUN_SCRIPT_BASENAME%.sh}"
if [[ "${RUN_SCRIPT_TAG}" == infer_ti2vid_one_stage_* ]]; then
  RUN_SCRIPT_TAG="${RUN_SCRIPT_TAG#infer_ti2vid_one_stage_}"
elif [[ "${RUN_SCRIPT_TAG}" == infer_ti2vid_* ]]; then
  RUN_SCRIPT_TAG="${RUN_SCRIPT_TAG#infer_ti2vid_}"
fi

IMAGE_DIR="${IMAGE_DIR:-${ROOT_DIR}/test_data/image}"
CONTROL_DIR="${CONTROL_DIR:-${ROOT_DIR}/test_data/control}"
PROMPT_DIR="${PROMPT_DIR:-${ROOT_DIR}/test_data/prompt}"
PYTHON_BIN="${PYTHON_BIN:-${ROOT_DIR}/.venv/bin/python}"

DEFAULT_IMAGE_PATH="${DEFAULT_IMAGE_PATH:-${ROOT_DIR}/test_data/image/zelda_castle.png}"
DEFAULT_CONTROL_PT="${DEFAULT_CONTROL_PT:-${ROOT_DIR}/test_data/control/600h_20260214_honghao_fAmoBF9m_20s_game_20s_s0006_f00005345-00005825.control.pt}"
DEFAULT_NUM_FRAMES="${DEFAULT_NUM_FRAMES:-481}"

RUN_TIMESTAMP="$(date +'%Y%m%d_%H%M%S')"
# 中文注释：默认批次目录名跟随实际入口脚本，避免 async/no_anchor 误落到 original 目录。
BATCH_NAME="${BATCH_NAME:-${RUN_SCRIPT_TAG}_batch_${RUN_TIMESTAMP}}"
OUTPUT_ROOT="${OUTPUT_ROOT:-${ROOT_DIR}/results/batch_runs/${BATCH_NAME}}"
CASE_MANIFEST="${CASE_MANIFEST:-${OUTPUT_ROOT}/case_manifest.tsv}"

CASE_GLOB="${CASE_GLOB:-*}"
MAX_CASES="${MAX_CASES:-0}"
DRY_RUN="${DRY_RUN:-false}"
SKIP_EXISTING="${SKIP_EXISTING:-true}"
TIMING_ENABLED="${TIMING_ENABLED:-false}"
TIMING_ROOT="${TIMING_ROOT:-}"

case "${DRY_RUN}" in
  true|false) ;;
  *)
    echo "ERROR: DRY_RUN must be exactly true or false." >&2
    exit 1
    ;;
esac

case "${SKIP_EXISTING}" in
  true|false) ;;
  *)
    echo "ERROR: SKIP_EXISTING must be exactly true or false." >&2
    exit 1
    ;;
esac

case "${TIMING_ENABLED}" in
  true|false) ;;
  *)
    echo "ERROR: TIMING_ENABLED must be exactly true or false." >&2
    exit 1
    ;;
esac

if [[ ! -x "${PYTHON_BIN}" ]]; then
  echo "ERROR: python not found: ${PYTHON_BIN}" >&2
  exit 1
fi

if [[ ! -f "${RUN_SCRIPT}" ]]; then
  echo "ERROR: run script not found: ${RUN_SCRIPT}" >&2
  exit 1
fi

mkdir -p "${OUTPUT_ROOT}" "${OUTPUT_ROOT}/outputs" "${OUTPUT_ROOT}/decoded_chunks" "${OUTPUT_ROOT}/logs"
if [[ "${TIMING_ENABLED}" == "true" ]]; then
  if [[ -z "${TIMING_ROOT}" ]]; then
    TIMING_ROOT="${OUTPUT_ROOT}/timing"
  fi
  mkdir -p "${TIMING_ROOT}"
fi

# 中文注释：先生成 manifest，后续无论切 original / no_anchor / overlap8，都复用同一批 case。
IMAGE_DIR="${IMAGE_DIR}" \
CONTROL_DIR="${CONTROL_DIR}" \
PROMPT_DIR="${PROMPT_DIR}" \
DEFAULT_IMAGE_PATH="${DEFAULT_IMAGE_PATH}" \
DEFAULT_CONTROL_PT="${DEFAULT_CONTROL_PT}" \
DEFAULT_NUM_FRAMES="${DEFAULT_NUM_FRAMES}" \
"${PYTHON_BIN}" - <<'PY' > "${CASE_MANIFEST}"
from __future__ import annotations

import os
import re
import sys
import unicodedata
from pathlib import Path

import torch

SEP = "\x1f"


def list_files(folder: Path, suffixes: tuple[str, ...]) -> list[Path]:
    if not folder.exists():
        return []
    return sorted(
        path for path in folder.iterdir()
        if path.is_file() and path.suffix.lower() in suffixes
    )


def image_stem(path: Path) -> str:
    stem = path.stem
    if stem.endswith(".first_frame"):
        stem = stem[: -len(".first_frame")]
    return stem


def control_stem(path: Path) -> str:
    stem = path.name[:-3] if path.name.endswith(".pt") else path.stem
    if stem.endswith(".control"):
        stem = stem[: -len(".control")]
    return stem


def prompt_stem(path: Path) -> str:
    suffix = ".prompt.txt"
    if path.name.endswith(suffix):
        return path.name[: -len(suffix)]
    return path.stem


def sanitize(text: str) -> str:
    text = unicodedata.normalize("NFKC", text)
    text = re.sub(r"[^0-9A-Za-z._-]+", "_", text).strip("._-")
    return text or "case"


def infer_num_frames(control_path: Path | None, default_frames: int) -> int:
    if control_path is None:
        return default_frames

    obj = torch.load(control_path, map_location="cpu")
    if isinstance(obj, dict):
        for key in ("control_features", "controls", "control", "features", "actions"):
            if key in obj:
                obj = obj[key]
                break

    shape = getattr(obj, "shape", None)
    if not shape:
        raise RuntimeError(f"Cannot infer num_frames from control file: {control_path}")
    return int(shape[0])


image_dir = Path(os.environ["IMAGE_DIR"])
control_dir = Path(os.environ["CONTROL_DIR"])
prompt_dir = Path(os.environ["PROMPT_DIR"])
default_image_path = Path(os.environ["DEFAULT_IMAGE_PATH"]).resolve()
default_control_pt = Path(os.environ["DEFAULT_CONTROL_PT"]).resolve()
default_num_frames = int(os.environ["DEFAULT_NUM_FRAMES"])

images_by_stem = {image_stem(path): path.resolve() for path in list_files(image_dir, (".png", ".jpg", ".jpeg", ".webp"))}
controls_by_stem = {control_stem(path): path.resolve() for path in list_files(control_dir, (".pt",))}
prompts_by_stem = {prompt_stem(path): path.resolve() for path in list_files(prompt_dir, (".txt",))}

all_stems = sorted(set(images_by_stem) | set(controls_by_stem) | set(prompts_by_stem))
if not all_stems:
    raise SystemExit("No test cases discovered under test_data/")

lines: list[str] = []
for idx, stem in enumerate(all_stems, start=1):
    matched_image = images_by_stem.get(stem)
    matched_control = controls_by_stem.get(stem)
    matched_prompt = prompts_by_stem.get(stem)

    image_path = matched_image or default_image_path
    control_path = matched_control or default_control_pt
    prompt_path = matched_prompt

    num_frames = infer_num_frames(matched_control, default_num_frames)

    if matched_image and matched_control:
        case_kind = "paired"
    elif matched_image:
        case_kind = "image_plus_default_control"
    elif matched_control:
        case_kind = "control_plus_default_image"
    else:
        case_kind = "prompt_plus_defaults"

    notes = []
    notes.append("image=matched" if matched_image else "image=default")
    notes.append("control=matched" if matched_control else "control=default")
    notes.append("prompt=matched" if matched_prompt else "prompt=default")
    notes.append(f"raw_stem={stem}")

    case_id = f"{idx:02d}_{case_kind}_{sanitize(stem)}"
    fields = [
        case_id,
        case_kind,
        str(image_path),
        str(control_path),
        str(prompt_path) if prompt_path else "",
        str(num_frames),
        ",".join(notes),
    ]
    lines.append(SEP.join(fields))

sys.stdout.write("\n".join(lines))
if lines:
    sys.stdout.write("\n")
PY

echo "BATCH_NAME      = ${BATCH_NAME}"
echo "RUN_SCRIPT      = ${RUN_SCRIPT}"
echo "RUN_SCRIPT_TAG  = ${RUN_SCRIPT_TAG}"
echo "OUTPUT_ROOT     = ${OUTPUT_ROOT}"
echo "CASE_MANIFEST   = ${CASE_MANIFEST}"
echo "CASE_GLOB       = ${CASE_GLOB}"
echo "MAX_CASES       = ${MAX_CASES}"
echo "DRY_RUN         = ${DRY_RUN}"
echo "SKIP_EXISTING   = ${SKIP_EXISTING}"
echo "TIMING_ENABLED  = ${TIMING_ENABLED}"
echo "TIMING_ROOT     = ${TIMING_ROOT:-<disabled>}"

selected_count=0
planned_count=0

while IFS=$'\x1f' read -r case_id case_kind image_path control_path prompt_path num_frames notes; do
  [[ -n "${case_id}" ]] || continue
  [[ "${case_id}" == ${CASE_GLOB} ]] || continue

  if [[ "${MAX_CASES}" -gt 0 ]] && [[ "${planned_count}" -ge "${MAX_CASES}" ]]; then
    break
  fi
  planned_count=$((planned_count + 1))

  output_path="${OUTPUT_ROOT}/outputs/${case_id}.mp4"
  chunks_dir="${OUTPUT_ROOT}/decoded_chunks/${case_id}"
  log_path="${OUTPUT_ROOT}/logs/${case_id}.log"
  timing_dir=""
  timing_summary_path=""
  if [[ "${TIMING_ENABLED}" == "true" ]]; then
    timing_dir="${TIMING_ROOT}/${case_id}"
    timing_summary_path="${timing_dir}/timing_summary.json"
  fi

  echo
  echo "[${planned_count}] ${case_id}"
  echo "  kind       = ${case_kind}"
  echo "  image      = ${image_path}"
  echo "  control    = ${control_path}"
  echo "  prompt     = ${prompt_path:-<default inline prompt>}"
  echo "  num_frames = ${num_frames}"
  echo "  notes      = ${notes}"
  echo "  output     = ${output_path}"
  echo "  chunks_dir = ${chunks_dir}"
  if [[ "${TIMING_ENABLED}" == "true" ]]; then
    echo "  timing_dir = ${timing_dir}"
  fi

  if [[ "${SKIP_EXISTING}" == "true" ]] && [[ -f "${output_path}" ]]; then
    if [[ "${TIMING_ENABLED}" != "true" ]] || [[ -f "${timing_summary_path}" ]]; then
      echo "  status     = skip_existing"
      continue
    fi
    echo "  status     = rerun_missing_timing"
  fi

  if [[ "${DRY_RUN}" == "true" ]]; then
    echo "  status     = dry_run"
    continue
  fi

  selected_count=$((selected_count + 1))
  if [[ "${TIMING_ENABLED}" == "true" ]]; then
    OUTPUT_PATH="${output_path}" \
    DECODED_CHUNKS_OUTPUT_DIR="${chunks_dir}" \
    IMAGE_PATH="${image_path}" \
    CONTROL_PT="${control_path}" \
    PROMPT_FILE="${prompt_path}" \
    NUM_FRAMES="${num_frames}" \
    TIMING_OUTPUT_DIR="${timing_dir}" \
    TIMING_EVENTS_OUTPUT="" \
    TIMING_SUMMARY_OUTPUT="" \
    TIMING_CASE_ID="${case_id}" \
    TIMING_VARIANT="${RUN_SCRIPT_TAG}" \
    bash "${RUN_SCRIPT}" 2>&1 | tee "${log_path}"
  else
    OUTPUT_PATH="${output_path}" \
    DECODED_CHUNKS_OUTPUT_DIR="${chunks_dir}" \
    IMAGE_PATH="${image_path}" \
    CONTROL_PT="${control_path}" \
    PROMPT_FILE="${prompt_path}" \
    NUM_FRAMES="${num_frames}" \
    TIMING_OUTPUT_DIR="" \
    TIMING_EVENTS_OUTPUT="" \
    TIMING_SUMMARY_OUTPUT="" \
    TIMING_CASE_ID="" \
    TIMING_VARIANT="" \
    bash "${RUN_SCRIPT}" 2>&1 | tee "${log_path}"
  fi
done < "${CASE_MANIFEST}"

echo
echo "Manifest written to: ${CASE_MANIFEST}"
echo "Cases enumerated: ${planned_count}"
echo "Cases launched:   ${selected_count}"