#!/bin/bash
set -e

MODEL_DIR="/data/models"

mkdir -p "$MODEL_DIR"

MODEL_REPO="ornith-ai/Ornith-1.5-9B-GGUF"
MODEL_FILE="Ornith-1.5-9B-Q4_K_M.gguf"

echo "=============================================="
echo " Ornith-1.5-9B"
echo " CPU / Q4_K_M / Long Context"
echo "=============================================="

echo "Downloading model..."

python - <<PY
from huggingface_hub import hf_hub_download, list_repo_files

files = list_repo_files("${MODEL_REPO}")
expected = "${MODEL_FILE}"
if expected not in files:
    print(f"ERROR: {expected} not found in repo. Available GGUF files:")
    for f in sorted(files):
        if f.endswith(".gguf"):
            print(f"  {f}")
    raise SystemExit(1)

path = hf_hub_download(
    repo_id="${MODEL_REPO}",
    filename="${MODEL_FILE}",
    local_dir="${MODEL_DIR}",
)

print("Downloaded:")
print(path)
PY

MODEL="${MODEL_DIR}/${MODEL_FILE}"

if [ ! -f "$MODEL" ]; then
    echo "ERROR: Model not found:"
    echo "$MODEL"
    echo
    echo "Files in model directory:"
    ls -lah "$MODEL_DIR"
    exit 1
fi

echo
echo "Model:"
ls -lh "$MODEL"

echo
echo "Starting llama-server..."

export LD_LIBRARY_PATH="/opt/llama${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"

exec /opt/llama/llama-server \
    --model "$MODEL" \
    --host 0.0.0.0 \
    --port 7860 \
    --ctx-size 32768 \
    --threads 2 \
    --threads-batch 2 \
    --batch-size 512 \
    --ubatch-size 128 \
    --cache-type-k q8_0 \
    --cache-type-v q8_0 \
    --flash-attn on \
    --cont-batching
 2>/dev/null || \
/opt/llama/llama-server \
    --model "$MODEL" \
    --host 0.0.0.0 \
    --port 7860 \
    --ctx-size 32768 \
    --threads 2 \
    --threads-batch 2 \
    --batch-size 512 \
    --ubatch-size 128 \
    --cache-type-k q8_0 \
    --cache-type-v q8_0 \
    --cont-batching