#!/usr/bin/env bash

set -euo pipefail
shopt -s extglob

common_tensor_types="
  "

TYPE_DEFAULT_FALLBACK=Q8_0

recipes=(
  "
  RECIPE_NAME=IQ3_XXS
  TYPE_DEFAULT=q6_k

  ffn_down_exps=iq3_xxs
  ffn_gate_exps=iq3_xxs
  ffn_up_exps=iq3_xxs
  "
)

recipe_line() {
  local line="$1"
  printf '%s' "${line##+([[:space:]])}"
}

recipe_value() {
  local recipe="$1"
  local key="$2"
  local line

  while IFS= read -r line || [ -n "$line" ]; do
    line="$(recipe_line "$line")"
    case "$line" in
      "${key}="*)
        printf '%s' "${line#*=}"
        return 0
        ;;
    esac
  done <<<"$recipe"

  return 1
}

recipe_type_default() {
  local recipe="$1"

  recipe_value "$recipe" TYPE_DEFAULT || printf '%s' "$TYPE_DEFAULT_FALLBACK"
}

print_recipes() {
  local recipe recipe_name type_default

  for recipe in "${recipes[@]}"; do
    recipe_name="$(recipe_value "$recipe" RECIPE_NAME)"
    type_default="$(recipe_type_default "$recipe")"

    printf '  %s:\n' "$recipe_name"
    printf '    TYPE_DEFAULT=%s\n' "$type_default"
    printf '\n'
  done
}

SELECTED_RECIPE_NAME=
SELECTED_TYPE_DEFAULT=
SELECTED_RECIPE=

load_recipe() {
  local requested_recipe="$1"
  local recipe recipe_name

  for recipe in "${recipes[@]}"; do
    recipe_name="$(recipe_value "$recipe" RECIPE_NAME)"
    if [ "$recipe_name" = "$requested_recipe" ]; then
      SELECTED_RECIPE_NAME="$recipe_name"
      SELECTED_TYPE_DEFAULT="$(recipe_type_default "$recipe")"
      SELECTED_RECIPE="$recipe"
      return 0
    fi
  done

  return 1
}

append_tensor_type_lines() {
  local block="$1"
  local output_path="$2"
  local skip_recipe_metadata="${3:-false}"
  local line

  while IFS= read -r line || [ -n "$line" ]; do
    line="$(recipe_line "$line")"
    if [ "$skip_recipe_metadata" = true ]; then
      case "$line" in
        RECIPE_NAME=* | TYPE_DEFAULT=*)
          continue
          ;;
      esac
    fi

    if [ -z "$line" ]; then
      continue
    fi

    printf '%s\n' "$line" >>"$output_path"
  done <<<"$block"
}

write_tensor_type_file() {
  local recipe="$1"
  local output_path="$2"

  : >"$output_path"
  append_tensor_type_lines "$common_tensor_types" "$output_path"
  append_tensor_type_lines "$recipe" "$output_path" true
}

usage() {
  cat <<'EOF'
Usage: quantize.sh <llama_cpp_dir> <recipe_name> [--dry-run] [--split]

Outputs are written relative to the current directory.

Environment overrides:
  INPUT_GGUF=<path>     Input GGUF. Defaults to the first *BF16*.gguf in BF16/.
  IMATRIX_PATH=<path>   Importance matrix path. Defaults to imatrix.gguf.

Available recipes:
EOF
  print_recipes
  cat <<'EOF'
Examples:
  ./scripts/quantize.sh ~/code/llama.cpp IQ3_XXS --dry-run
  ./scripts/quantize.sh ~/code/llama.cpp IQ3_XXS
  ./scripts/quantize.sh ~/code/llama.cpp IQ3_XXS --split
EOF
}

if [ $# -lt 2 ]; then
  usage
  exit 1
fi

LLAMA_CPP_DIR="$1"
REQUESTED_RECIPE="$2"
DRY_RUN=false
SPLIT=false
shift 2

while [ $# -gt 0 ]; do
  case "$1" in
    --dry-run)
      DRY_RUN=true
      ;;
    --split)
      SPLIT=true
      ;;
    *)
      usage
      exit 1
      ;;
  esac
  shift
done

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"
BF16_DIR="$PROJECT_DIR/BF16"
SPLIT_SCRIPT="$SCRIPT_DIR/split.sh"

if ! load_recipe "$REQUESTED_RECIPE"; then
  echo "Error: unknown recipe: $REQUESTED_RECIPE" >&2
  usage >&2
  exit 1
fi

if [ ! -d "$LLAMA_CPP_DIR" ]; then
  echo "Error: llama.cpp directory not found: $LLAMA_CPP_DIR" >&2
  exit 1
fi

QUANTIZE_BIN="$LLAMA_CPP_DIR/build/bin/llama-quantize"
SPLIT_BIN="$LLAMA_CPP_DIR/build/bin/llama-gguf-split"

if [ ! -x "$QUANTIZE_BIN" ]; then
  echo "Error: llama-quantize binary not found: $QUANTIZE_BIN" >&2
  exit 1
fi

if [ "$SPLIT" = true ] && [ ! -x "$SPLIT_SCRIPT" ]; then
  echo "Error: split script not found or not executable: $SPLIT_SCRIPT" >&2
  exit 1
fi

if [ "$SPLIT" = true ] && [ ! -x "$SPLIT_BIN" ]; then
  echo "Error: llama-gguf-split binary not found: $SPLIT_BIN" >&2
  exit 1
fi

if [ -z "${INPUT_GGUF-}" ]; then
  if [ ! -d "$BF16_DIR" ]; then
    echo "Error: BF16 directory not found: $BF16_DIR" >&2
    exit 1
  fi

  INPUT_GGUF="$(find "$BF16_DIR" -maxdepth 1 -name "*BF16*.gguf" -type f | sort | head -n 1)"
fi

if [ -z "$INPUT_GGUF" ] || [ ! -e "$INPUT_GGUF" ]; then
  echo "Error: input GGUF not found: ${INPUT_GGUF:-<none>}" >&2
  exit 1
fi

IMATRIX_PATH="${IMATRIX_PATH:-$PROJECT_DIR/imatrix.gguf}"
if [ ! -e "$IMATRIX_PATH" ]; then
  echo "Error: imatrix file not found: $IMATRIX_PATH" >&2
  exit 1
fi

cd "$LLAMA_CPP_DIR"
if [ -f .venv/bin/activate ]; then
  # shellcheck disable=SC1091
  source .venv/bin/activate
fi
cd - >/dev/null

OUTPUT_BASE_DIR="$(pwd)"
INPUT_BASENAME="$(basename "$INPUT_GGUF")"
MODEL_NAME="$(printf '%s\n' "$INPUT_BASENAME" | sed -E 's/-BF16(-[0-9]+-of-[0-9]+)?\.gguf$//')"
INTERMEDIATE_OUTPUT="$OUTPUT_BASE_DIR/${MODEL_NAME}-${SELECTED_RECIPE_NAME}.gguf"
OUTPUT_DIR="$OUTPUT_BASE_DIR/$SELECTED_RECIPE_NAME"

if [ "$DRY_RUN" = false ] && [ -e "$INTERMEDIATE_OUTPUT" ]; then
  echo "Error: quantized output already exists: $INTERMEDIATE_OUTPUT" >&2
  exit 1
fi

TENSOR_TYPE_FILE="$(mktemp "${TMPDIR:-/tmp}/nex-n2-pro-tensor-types.XXXXXX")"
cleanup() {
  rm -f "$TENSOR_TYPE_FILE"
}
trap cleanup EXIT

write_tensor_type_file "$SELECTED_RECIPE" "$TENSOR_TYPE_FILE"

echo "Input: $INPUT_GGUF"
echo "Recipe name: $SELECTED_RECIPE_NAME"
echo "Tensor type file: $TENSOR_TYPE_FILE"
echo "Default type: $SELECTED_TYPE_DEFAULT"
echo "imatrix: $IMATRIX_PATH"
echo "Output base: $OUTPUT_BASE_DIR"
echo "Split output: $SPLIT"

if [ "$DRY_RUN" = true ]; then
  "$QUANTIZE_BIN" \
    --dry-run \
    --allow-requantize \
    --tensor-type-file "$TENSOR_TYPE_FILE" \
    --imatrix "$IMATRIX_PATH" \
    "$INPUT_GGUF" \
    "$INTERMEDIATE_OUTPUT" \
    "$SELECTED_TYPE_DEFAULT"
else
  "$QUANTIZE_BIN" \
    --allow-requantize \
    --tensor-type-file "$TENSOR_TYPE_FILE" \
    --imatrix "$IMATRIX_PATH" \
    "$INPUT_GGUF" \
    "$INTERMEDIATE_OUTPUT" \
    "$SELECTED_TYPE_DEFAULT"
fi

if [ "$DRY_RUN" = false ]; then
  if [ "$SPLIT" = true ]; then
    "$SPLIT_SCRIPT" \
      "$LLAMA_CPP_DIR" \
      "$INTERMEDIATE_OUTPUT" \
      "$OUTPUT_DIR" \
      "${MODEL_NAME}-${SELECTED_RECIPE_NAME}" \
      --remove-input
  else
    echo "Quantization complete. Output saved to: $INTERMEDIATE_OUTPUT"
  fi
fi