#!/usr/bin/env bash
# Merge KV cache settings (q8_0) into a model YAML on /data.
# Usage: ./scripts/apply-kv-profile.sh <model-name-without-.yaml>
# Example: ./scripts/apply-kv-profile.sh gemma-4-12b-it-qat-q4_0
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
STACK_DIR="$(cd "${SCRIPT_DIR}/.." && pwd)"

if [[ ! -f "${STACK_DIR}/.env" ]]; then
  echo "ERROR: ${STACK_DIR}/.env not found"
  exit 1
fi

set -a
# shellcheck disable=SC1091
source "${STACK_DIR}/.env"
set +a

DATA_ROOT="${DATA_ROOT:-/data}"
MODELS_DIR="${DATA_ROOT}/apps/localai/models"

if [[ $# -lt 1 ]]; then
  echo "Usage: $0 <model-name>"
  echo "Example: $0 gemma-4-12b-it-qat-q4_0"
  exit 1
fi

MODEL_NAME="$1"
TARGET="${MODELS_DIR}/${MODEL_NAME}.yaml"

if [[ ! -f "${TARGET}" ]]; then
  echo "ERROR: ${TARGET} not found"
  exit 1
fi

BACKUP="${TARGET}.bak.$(date +%Y%m%d%H%M%S)"
cp "${TARGET}" "${BACKUP}"
echo "Backup: ${BACKUP}"

run_python_merge() {
  python3 - "${1}" <<'PY'
import sys
from pathlib import Path

try:
    import yaml
except ImportError:
    sys.exit("ERROR: python3-yaml required (sudo apt install python3-yaml)")

path = Path(sys.argv[1])
data = yaml.safe_load(path.read_text(encoding="utf-8")) or {}

params = data.setdefault("parameters", {})
kv = {
    "cache_type_k": "q8_0",
    "cache_type_v": "q8_0",
    "flash_attention": True,
}
if "context_size" not in params:
    kv["context_size"] = 8192
params.update(kv)

path.write_text(
    yaml.dump(data, default_flow_style=False, allow_unicode=True, sort_keys=False),
    encoding="utf-8",
)
print(f"Updated KV settings in {path}")
for k, v in kv.items():
    print(f"  {k}: {v}")
if "context_size" in params and "context_size" not in kv:
    print(f"  context_size: {params['context_size']} (unchanged)")
PY
}

if run_python_merge "${TARGET}" 2>/dev/null; then
  :
elif docker ps --format '{{.Names}}' 2>/dev/null | grep -qx localai; then
  echo "Host write failed — applying via docker exec localai (models volume is root-owned)"
  CONTAINER_PATH="/models/${MODEL_NAME}.yaml"
  docker exec localai sh -c "grep -q cache_type_k '${CONTAINER_PATH}' || sed -i '/^    model:/a\\
    cache_type_k: q8_0\\
    cache_type_v: q8_0\\
    flash_attention: true\\
    context_size: 8192' '${CONTAINER_PATH}'"
  grep -E 'cache_type_k|cache_type_v|flash_attention|context_size' "${TARGET}" || true
else
  echo "ERROR: cannot write ${TARGET} (permission denied) and container localai not running"
  echo "Run: sudo $0 ${MODEL_NAME}"
  exit 1
fi

echo ""
echo "Restart LocalAI to apply:"
echo "  cd ${STACK_DIR} && docker compose --profile localai restart localai"