#!/bin/bash
# =============================================================================
# Simple Router — llama.cpp wrapper
# Uses 3 GGUF experts + keyword classifier → picks the right one
# 
# Usage:
#   ./simple_router.sh "Write a Python function to sort a list"
#   ./simple_router.sh "Solve x^2 + 5x + 6 = 0"
#   ./simple_router.sh "What is the capital of Thailand?"
#
# Requirements:
#   - llama.cpp (llama-cli) in PATH
#   - GGUF files in ./gguf/{coding,math,chat}/ 
#   OR set GGUF_CODING, GGUF_MATH, GGUF_CHAT env vars
# =============================================================================

PROMPT="$*"
if [ -z "$PROMPT" ]; then
    echo "Usage: $0 <prompt>"
    exit 1
fi

# ─── GGUF Paths ───
GGUF_DIR="${GGUF_DIR:-./gguf}"
GGUF_CODING="${GGUF_CODING:-$GGUF_DIR/coding/frankenmoe_coding-Q4_K_M.gguf}"
GGUF_MATH="${GGUF_MATH:-$GGUF_DIR/math/frankenmoe_math-Q4_K_M.gguf}"
GGUF_CHAT="${GGUF_CHAT:-$GGUF_DIR/chat/frankenmoe_chat-Q4_K_M.gguf}"

# ─── Classification ───
LOWER=$(echo "$PROMPT" | tr '[:upper:]' '[:lower:]')

CODING_KW=("def " "function" "python" "code" "bug" "debug" "api" "import" "class " "algorithm" "implement" "compile" "syntax" "javascript" "html" "css" "sql" "bash" "git" "docker" "write a" "program" "script" "loop" "array" "list" "dict")

MATH_KW=("solve" "equation" "derivative" "integral" "matrix" "eigen" "theorem" "proof" "sqrt" "log" "sin" "cos" "tan" "sum" "probability" "statistic" "graph" "vector" "polynomial" "x =" "x=" "y =" "calculate" "compute" "find the")

CODING_SCORE=0
MATH_SCORE=0

for kw in "${CODING_KW[@]}"; do
    if [[ "$LOWER" == *"$kw"* ]]; then
        ((CODING_SCORE++))
    fi
done

for kw in "${MATH_KW[@]}"; do
    if [[ "$LOWER" == *"$kw"* ]]; then
        ((MATH_SCORE++))
    fi
done

# ─── Route ───
if [ "$CODING_SCORE" -gt 0 ] && [ "$CODING_SCORE" -ge "$MATH_SCORE" ]; then
    MODEL="$GGUF_CODING"
    DOMAIN="🖥️  CODING"
elif [ "$MATH_SCORE" -gt 0 ]; then
    MODEL="$GGUF_MATH"
    DOMAIN="📐 MATH"
else
    MODEL="$GGUF_CHAT"
    DOMAIN="💬 CHAT"
fi

echo "════════════════════════════════════"
echo "  Simple Router — Domain Router"
echo "  Prompt: ${PROMPT:0:60}..."
echo "  Route:  $DOMAIN"
echo "  Model:  $(basename "$MODEL")"
echo "════════════════════════════════════"
echo ""

# ─── Check model exists ───
if [ ! -f "$MODEL" ]; then
    echo "❌ Model not found: $MODEL"
    echo ""
    echo "Download GGUF models:"
    echo "  wget https://huggingface.co/hotdogs/frankenmoe/resolve/main/coding/frankenmoe_coding-Q4_K_M.gguf"
    echo "  wget https://huggingface.co/hotdogs/frankenmoe/resolve/main/math/frankenmoe_math-Q4_K_M.gguf"
    echo "  wget https://huggingface.co/hotdogs/frankenmoe/resolve/main/chat/frankenmoe_chat-Q4_K_M.gguf"
    exit 1
fi

# ─── Run llama.cpp ───
llama-cli \
    -m "$MODEL" \
    -p "$PROMPT" \
    -n 256 \
    --temp 0.7 \
    --top-p 0.9 \
    --no-display-prompt \
    2>/dev/null

echo ""
echo "════════════════════════════════════"