#!/usr/bin/env bash
set -Eeuo pipefail

# ============================================================
# Configuration
# ============================================================

# Direct URL to a GGUF file on Hugging Face.
# Override on the command line or export before running.
#MODEL_URL="${MODEL_URL:-}"

## Gemma 4
#MODEL_URL="https://huggingface.co/google/gemma-4-12B-it-qat-q4_0-gguf/resolve/main/gemma-4-12b-it-qat-q4_0.gguf"

MODEL_URL="https://huggingface.co/unsloth/gemma-4-12B-it-qat-GGUF/resolve/main/gemma-4-12B-it-qat-UD-Q4_K_XL.gguf"
MODEL_EXTRA="--mmproj /models/mmproj-F16.gguf --model-draft /models/mtp-gemma-4-12B-it.gguf --spec-type draft-mtp --spec-draft-n-max 2"

#MODEL_URL="https://huggingface.co/unsloth/Qwen3.8-27B-GGUF/resolve/main/Qwen3.8-27B-UD-IQ2_XXS.gguf"   ## 7.27 GB
#MODEL_URL="https://huggingface.co/unsloth/Qwen3.8-27B-GGUF/resolve/main/Qwen3.8-27B-UD-IQ2_S.gguf"     ## 8.37 GB
#MODEL_URL="https://huggingface.co/unsloth/Qwen3.8-27B-GGUF/resolve/main/Qwen3.8-27B-UD-Q2_K_XL.gguf"   ## 9.83 GB
#MODEL_URL="https://huggingface.co/unsloth/Qwen3.8-27B-GGUF/resolve/main/Qwen3.8-27B-UD-IQ3_XXS.gguf"   ## 10.9 GB
#MODEL_URL="https://huggingface.co/unsloth/Qwen3.8-27B-GGUF/resolve/main/Qwen3.8-27B-UD-IQ3_S.gguf"     ## 12.0 GB
#MODEL_URL="https://huggingface.co/unsloth/Qwen3.8-27B-GGUF/resolve/main/Qwen3.8-27B-UD-Q3_K_XL.gguf"   ## 13.1 GB

# Persistent model cache.
MODEL_DIR="${MODEL_DIR:-$HOME/.cache/llama.cpp}"

# Docker images.
#OPENCODE_IMAGE="${OPENCODE_IMAGE:-ghcr.io/anomalyco/opencode:latest}"
OPENCODE_IMAGE="${OPENCODE_IMAGE:-sbx/opencode-image:latest}"
LLAMA_IMAGE="${LLAMA_IMAGE:-ghcr.io/ggml-org/llama.cpp:server-rocm}"

# ROCm compatibility workaround for the AMD Radeon RX 6700 XT.
HSA_OVERRIDE_GFX_VERSION="${HSA_OVERRIDE_GFX_VERSION:-10.3.0}"

# Memory-conscious llama.cpp settings.
CONTEXT_SIZE="${CONTEXT_SIZE:-256000}"
GPU_LAYERS="${GPU_LAYERS:-auto}"
KV_CACHE_TYPE="${KV_CACHE_TYPE:-q4_0}"
PARALLEL="${PARALLEL:-1}"
BATCH_SIZE="${BATCH_SIZE:-512}"
UBATCH_SIZE="${UBATCH_SIZE:-128}"
THREADS="${THREADS:-12}"

# Host port for accessing the llama.cpp API.
LLAMA_PORT="${LLAMA_PORT:-8080}"

# OpenCode configuration.
OPENCODE_CONFIG_DIR="${OPENCODE_CONFIG_DIR:-$HOME/.config/opencode}"

# Readiness timeout in seconds.
STARTUP_TIMEOUT="${STARTUP_TIMEOUT:-360}"

# ============================================================
# Validation
# ============================================================

if [[ -z "$MODEL_URL" ]]; then
    echo "Error: MODEL_URL must point to a GGUF model." >&2
    echo "Example:" >&2
    echo "MODEL_URL='https://huggingface.co/OWNER/REPO/resolve/main/MODEL.gguf' $0" >&2
    exit 1
fi

for cmd in docker curl realpath basename; do
    command -v "$cmd" >/dev/null 2>&1 || {
        echo "Error: Missing required command: $cmd" >&2
        exit 1
    }
done

if ! [[ "$LLAMA_PORT" =~ ^[0-9]+$ ]] ||
   (( LLAMA_PORT < 1 || LLAMA_PORT > 65535 )); then
    echo "Error: LLAMA_PORT must be between 1 and 65535." >&2
    exit 1
fi

WORKDIR="$(pwd -P)"
MODEL_FILENAME="$(basename "${MODEL_URL%%\?*}")"

if [[ "$MODEL_FILENAME" != *.gguf ]]; then
    echo "Error: MODEL_URL must point to a .gguf file." >&2
    exit 1
fi

mkdir -p "$MODEL_DIR" "$OPENCODE_CONFIG_DIR"

MODEL_DIR="$(realpath "$MODEL_DIR")"
OPENCODE_CONFIG_DIR="$(realpath "$OPENCODE_CONFIG_DIR")"
MODEL_PATH="$MODEL_DIR/$MODEL_FILENAME"

SUFFIX="$$"
NETWORK="opencode-llama-$SUFFIX"
LLAMA_CONTAINER="llama-$SUFFIX"
OPENCODE_CONTAINER="opencode-$SUFFIX"

# ============================================================
# Cleanup
# ============================================================

cleanup() {
    local status=$?
    trap - EXIT INT TERM

    echo
    echo "Stopping containers..."

    if docker container inspect "$LLAMA_CONTAINER" >/dev/null 2>&1; then
        docker logs --tail 80 "$LLAMA_CONTAINER" >&2 || true

        docker inspect "$LLAMA_CONTAINER" \
            --format 'Exit={{.State.ExitCode}} OOMKilled={{.State.OOMKilled}} Error={{.State.Error}}' \
            >&2 || true
    fi

    docker rm -f "$OPENCODE_CONTAINER" >/dev/null 2>&1 || true
    docker rm -f "$LLAMA_CONTAINER" >/dev/null 2>&1 || true
    docker network rm "$NETWORK" >/dev/null 2>&1 || true

    exit "$status"
}

trap cleanup EXIT
trap 'exit 130' INT
trap 'exit 143' TERM

# ============================================================
# Download model if needed
# ============================================================

if [[ -s "$MODEL_PATH" ]]; then
    echo "Using cached model: $MODEL_PATH"
else
    echo "Downloading model: $MODEL_FILENAME"

    TEMP_MODEL="${MODEL_PATH}.part"
    CURL_AUTH=()

    if [[ -n "${HF_TOKEN:-}" ]]; then
        CURL_AUTH=(-H "Authorization: Bearer $HF_TOKEN")
    fi

    if ! curl \
        --fail \
        --location \
        --show-error \
        --retry 3 \
        --continue-at - \
        "${CURL_AUTH[@]}" \
        --output "$TEMP_MODEL" \
        "$MODEL_URL"; then
        echo "Error: Model download failed." >&2
        echo "Partial download retained at: $TEMP_MODEL" >&2
        exit 1
    fi

    if [[ ! -s "$TEMP_MODEL" ]]; then
        rm -f "$TEMP_MODEL"
        echo "Error: Downloaded model is empty." >&2
        exit 1
    fi

    mv "$TEMP_MODEL" "$MODEL_PATH"
    echo "Model cached at $MODEL_PATH"
fi

# ============================================================
# Pull images and create network
# ============================================================

echo "Pulling llama.cpp image..."
docker pull "$LLAMA_IMAGE"

echo "Pulling OpenCode image..."
docker pull "$OPENCODE_IMAGE"

docker network create "$NETWORK" >/dev/null

# ============================================================
# Start llama.cpp with ROCm
# ============================================================

echo "Starting llama.cpp..."
echo "Model:          $MODEL_FILENAME"
echo "Context:        $CONTEXT_SIZE"
echo "GPU layers:     $GPU_LAYERS"
echo "KV cache:       $KV_CACHE_TYPE"
echo "Parallel slots: $PARALLEL"
echo "ROCm override:  $HSA_OVERRIDE_GFX_VERSION"
echo "Host API port:  $LLAMA_PORT"

docker run -d \
    --name "$LLAMA_CONTAINER" \
    --network "$NETWORK" \
    --network-alias llama \
    --device=/dev/kfd \
    --device=/dev/dri \
    --group-add video \
    --group-add render \
    --cap-add=SYS_PTRACE \
    --security-opt seccomp=unconfined \
    -e "HSA_OVERRIDE_GFX_VERSION=$HSA_OVERRIDE_GFX_VERSION" \
    --mount "type=bind,src=$MODEL_DIR,dst=/models,readonly" \
    -p "127.0.0.1:${LLAMA_PORT}:8080" \
    "$LLAMA_IMAGE" \
    -m "/models/$MODEL_FILENAME" \
    --host 0.0.0.0 \
    --port 8080 \
    --ctx-size "$CONTEXT_SIZE" \
    --n-gpu-layers "$GPU_LAYERS" \
    --fit on \
    --cache-type-k "$KV_CACHE_TYPE" \
    --cache-type-v "$KV_CACHE_TYPE" \
    --parallel "$PARALLEL" \
    --batch-size "$BATCH_SIZE" \
    --ubatch-size "$UBATCH_SIZE" \
    --threads "$THREADS" \
    --threads-batch "$THREADS" \
    --flash-attn on \
    --load-mode mmap \
    --lazy-mode on \
    --temp 1.0 \
    --top-p 0.95 \
    --top-k 64 \
    --min-p 0.0 \
    --presence-penalty 0.0 \
    --repeat-penalty 1.0 \
    $MODEL_EXTRA

# ============================================================
# Wait for model readiness
# ============================================================

echo "Waiting for llama.cpp..."

READY=false
DEADLINE=$((SECONDS + STARTUP_TIMEOUT))

while (( SECONDS < DEADLINE )); do
    RUNNING="$(docker inspect \
        --format '{{.State.Running}}' \
        "$LLAMA_CONTAINER" 2>/dev/null || echo false)"

    if [[ "$RUNNING" != true ]]; then
        echo "Error: llama.cpp exited during startup." >&2
        docker logs "$LLAMA_CONTAINER" >&2 || true
        exit 1
    fi

    if curl --silent --fail --max-time 2 \
        "http://127.0.0.1:${LLAMA_PORT}/health" >/dev/null; then
        READY=true
        break
    fi

    sleep 2
done

if [[ "$READY" != true ]]; then
    echo "Error: llama.cpp did not become ready within ${STARTUP_TIMEOUT}s." >&2
    docker logs "$LLAMA_CONTAINER" >&2 || true
    exit 1
fi

echo
echo "llama.cpp is ready."
echo "OpenCode API endpoint: http://llama:8080/v1"
echo "Host API endpoint:     http://127.0.0.1:${LLAMA_PORT}/v1"
echo "Workspace:             $WORKDIR"
echo

# ============================================================
# Run OpenCode TUI
# ============================================================

docker run --rm -it \
    --name "$OPENCODE_CONTAINER" \
    --network "$NETWORK" \
    --mount "type=bind,src=$WORKDIR,dst=/workspace" \
    --workdir /workspace \
    --mount "type=bind,src=$OPENCODE_CONFIG_DIR,dst=/home/agent/.config/opencode" \
    "$OPENCODE_IMAGE" \
    opencode /workspace