#!/bin/bash
# Simple wrapper script for running Granite Docling with llama.cpp
# Usage: ./run_granite_docling.sh <image_path> [prompt]

# Check arguments
if [ $# -lt 1 ]; then
    echo "Usage: $0 <image_path> [prompt]"
    echo "Example: $0 document.png"
    echo "Example: $0 document.png 'Convert this table to OTSL.'"
    exit 1
fi

IMAGE_PATH="$1"
PROMPT="${2:-Convert this page to docling.}"

# Check if image exists
if [ ! -f "$IMAGE_PATH" ]; then
    echo "Error: Image file not found: $IMAGE_PATH"
    exit 1
fi

# Get script directory
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# Check for required files
MODEL_FILE="$SCRIPT_DIR/granite-docling-258M-f16.gguf"
MMPROJ_FILE="$SCRIPT_DIR/mmproj-granite-docling-258M-f16.gguf"
TEMPLATE_FILE="$SCRIPT_DIR/chat_template.jinja"

if [ ! -f "$MODEL_FILE" ]; then
    echo "Error: Model file not found: $MODEL_FILE"
    exit 1
fi

if [ ! -f "$MMPROJ_FILE" ]; then
    echo "Error: MMProj file not found: $MMPROJ_FILE"
    exit 1
fi

if [ ! -f "$TEMPLATE_FILE" ]; then
    echo "Error: Chat template file not found: $TEMPLATE_FILE"
    exit 1
fi

# Find llama-mtmd-cli
if command -v llama-mtmd-cli &> /dev/null; then
    LLAMA_CLI="llama-mtmd-cli"
elif [ -f "$HOME/llama.cpp/build/bin/llama-mtmd-cli" ]; then
    LLAMA_CLI="$HOME/llama.cpp/build/bin/llama-mtmd-cli"
else
    echo "Error: llama-mtmd-cli not found. Please install llama.cpp or set it in PATH"
    exit 1
fi

echo "Running Granite Docling on: $IMAGE_PATH"
echo "Prompt: $PROMPT"
echo "---"

# Run inference
$LLAMA_CLI \
  -m "$MODEL_FILE" \
  --mmproj "$MMPROJ_FILE" \
  --image "$IMAGE_PATH" \
  --chat-template "$(cat "$TEMPLATE_FILE")" \
  -p "$PROMPT" \
  -n 512 \
  --temp 0.1 \
  -ngl 99 \
  2>&1 | grep -v "^ggml_cuda_init\|^build:\|^llama_model_load\|^llama_model_loader\|^print_info\|^load:\|^load_tensors\|^llama_context\|^llama_kv_cache\|^common_init\|^mtmd_cli_context\|^clip_model_loader\|^clip_ctx\|^load_hparams\|^warmup\|^alloc_compute\|^main:\|^encoding image\|^image slice\|^decoding image\|^image decoded\|^llama_perf"