#!/usr/bin/env bash

set -euo pipefail

usage() {
  cat <<'EOF'
Usage: split.sh <llama_cpp_dir> <input_gguf> <output_dir> [output_prefix] [--remove-input]

Environment overrides:
  SPLIT_MAX_SIZE=50G    Maximum size for each split file.

Examples:
  ./scripts/split.sh ~/code/llama.cpp model.gguf ./Q5_K_S
  ./scripts/split.sh ~/code/llama.cpp model.gguf ./Q5_K_S model-Q5_K_S --remove-input
EOF
}

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

LLAMA_CPP_DIR="$1"
INPUT_GGUF="$2"
OUTPUT_DIR="$3"
OUTPUT_PREFIX_NAME=
REMOVE_INPUT=false
shift 3

while [ $# -gt 0 ]; do
  case "$1" in
    --remove-input)
      REMOVE_INPUT=true
      ;;
    --*)
      usage
      exit 1
      ;;
    *)
      if [ -n "$OUTPUT_PREFIX_NAME" ]; then
        usage
        exit 1
      fi
      OUTPUT_PREFIX_NAME="$1"
      ;;
  esac
  shift
done

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

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

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

if [ -z "$OUTPUT_PREFIX_NAME" ]; then
  input_basename="$(basename "$INPUT_GGUF")"
  OUTPUT_PREFIX_NAME="${input_basename%.gguf}"
fi

SPLIT_MAX_SIZE="${SPLIT_MAX_SIZE:-50G}"
mkdir -p "$OUTPUT_DIR"
OUTPUT_PREFIX="$OUTPUT_DIR/$OUTPUT_PREFIX_NAME"

"$SPLIT_BIN" \
  --split-max-size "$SPLIT_MAX_SIZE" \
  --no-tensor-first-split \
  "$INPUT_GGUF" \
  "$OUTPUT_PREFIX"

if [ "$REMOVE_INPUT" = true ]; then
  rm -f "$INPUT_GGUF"
fi

echo "Split complete. Output saved to: $OUTPUT_DIR"