#!/bin/bash
# One-command deployment for Hebrew legal RAG.
# Clears caches, restarts server, uploads all 3 JSONL files, runs smoke tests.
#
# Usage:
#   ./deploy_hebrew_rag.sh          # full deploy
#   ./deploy_hebrew_rag.sh restart  # only restart
#   ./deploy_hebrew_rag.sh upload   # only re-upload data
#   ./deploy_hebrew_rag.sh test     # only run tests
#   ./deploy_hebrew_rag.sh status   # show current state
#
# Requires: bash, curl, python3

set -e

# Resolve project root — the directory of this script's parent's parent
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
TAU_RAG_DIR="$( cd "$SCRIPT_DIR/.." && pwd )"
PROJECT_ROOT="$( cd "$TAU_RAG_DIR/.." && pwd )"
API="${API:-http://127.0.0.1:8000}"
PRESET="${TAU_RAG_PRESET:-hebrew_legal_prod}"
PORT="${PORT:-8000}"

GREEN='\033[0;32m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
BLUE='\033[0;34m'
NC='\033[0m'

log()  { echo -e "${BLUE}▸${NC} $*"; }
ok()   { echo -e "${GREEN}✅${NC} $*"; }
warn() { echo -e "${YELLOW}⚠️ ${NC} $*"; }
err()  { echo -e "${RED}❌${NC} $*"; }

clear_cache() {
    log "Clearing Python bytecode cache..."
    find "$TAU_RAG_DIR" -name "__pycache__" -type d -exec rm -rf {} + 2>/dev/null || true
    ok "Cache cleared"
}

stop_server() {
    log "Stopping any running uvicorn for tau_rag..."
    pkill -9 -f "uvicorn.*tau_rag" 2>/dev/null || true
    sleep 2
}

start_server() {
    log "Starting uvicorn (preset=$PRESET, port=$PORT, single worker)..."
    cd "$PROJECT_ROOT"
    # v3.x — env vars this script passes through:
    #   TAU_RAG_PRESET              — which preset to load
    #   TAU_RAG_CORS_ORIGINS        — CORS policy (default: *)
    #   TAU_RAG_COLLECT_TRAINING    — opt-in training data collection
    #   TAU_RAG_TAU_CKPT            — custom LLM checkpoint (if preset uses
    #                                 provider=tau_native)
    #   TAU_RAG_TAU_ENCODING_PATH   — dir containing encoding/hebrew_encoder.py
    TAU_RAG_PRESET="$PRESET" \
        TAU_RAG_CORS_ORIGINS="${TAU_RAG_CORS_ORIGINS:-*}" \
        TAU_RAG_COLLECT_TRAINING="${TAU_RAG_COLLECT_TRAINING:-1}" \
        TAU_RAG_TAU_CKPT="${TAU_RAG_TAU_CKPT:-}" \
        TAU_RAG_TAU_ENCODING_PATH="${TAU_RAG_TAU_ENCODING_PATH:-}" \
        TAU_RAG_TAU_DEBUG="${TAU_RAG_TAU_DEBUG:-}" \
        TAU_RAG_GATE_OVERLAP="${TAU_RAG_GATE_OVERLAP:-}" \
        nohup uvicorn tau_rag.api.fastapi_app:app \
            --port "$PORT" --workers 1 \
            >/tmp/tau_rag_server.log 2>&1 &
    local pid=$!
    sleep 5
    # Health check
    if curl -sf "$API/" >/dev/null 2>&1 || curl -sf "$API/v1/index_stats" >/dev/null 2>&1 ; then
        ok "Server up (PID=$pid), logs: /tmp/tau_rag_server.log"
    else
        err "Server failed to start. Check /tmp/tau_rag_server.log"
        tail -20 /tmp/tau_rag_server.log
        exit 1
    fi
}

# Check if the live index has any docs. Returns 0 if populated, 1 if empty.
# Used so `restart` auto-uploads when the server came up with a blank index
# (i.e. no persistence behind the preset).
index_is_empty() {
    local n
    n=$(curl -sf "$API/v1/index_stats" 2>/dev/null \
        | python3 -c "
import sys, json
try:
    d = json.load(sys.stdin)
    print(int(d.get('n_docs') or d.get('n_documents') or 0))
except Exception:
    print(0)
" 2>/dev/null || echo "0")
    [ "$n" -eq 0 ] 2>/dev/null && return 0 || return 1
}

upload_data() {
    local uploads_dir="$TAU_RAG_DIR/runtime/uploads"
    for f in heb_contract_law heb_law_supplement heb_caselaw_supreme; do
        local path="$uploads_dir/${f}.jsonl"
        if [ ! -f "$path" ]; then
            warn "Missing $f.jsonl — skipping"
            continue
        fi
        local size=$(stat -c%s "$path" 2>/dev/null || stat -f%z "$path" 2>/dev/null)
        log "Uploading $f.jsonl ($(($size / 1024))KB)..."
        local result=$(curl -s -X POST "$API/v1/data/upload" -F "file=@$path")
        local n=$(echo "$result" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d.get('n_rows_parsed','?'))" 2>/dev/null || echo "?")
        ok "$f: $n rows indexed"
    done
}

run_tests() {
    local test_script="$SCRIPT_DIR/test_hebrew_rag.sh"
    if [ -x "$test_script" ]; then
        log "Running regression tests..."
        API="$API" "$test_script"
    else
        warn "test_hebrew_rag.sh not found or not executable"
    fi
}

show_status() {
    log "Current state:"
    echo "  PROJECT_ROOT: $PROJECT_ROOT"
    echo "  TAU_RAG_DIR:  $TAU_RAG_DIR"
    echo "  API endpoint: $API"
    echo "  PRESET:       $PRESET"
    echo ""
    log "Server process:"
    pgrep -f "uvicorn.*tau_rag" | head -3 || echo "  (not running)"
    echo ""
    log "Data files:"
    for f in heb_contract_law heb_law_supplement heb_caselaw_supreme; do
        local path="$TAU_RAG_DIR/runtime/uploads/${f}.jsonl"
        if [ -f "$path" ]; then
            local size=$(stat -c%s "$path" 2>/dev/null || stat -f%z "$path" 2>/dev/null)
            local lines=$(wc -l < "$path")
            printf "  %-26s  %s lines, %sKB\n" "${f}.jsonl" "$lines" "$(($size / 1024))"
        fi
    done
    echo ""
    log "Code fingerprints (must be present for full functionality):"
    for kw in "_chunk_is_usable" "heb_ratio" "_query_is_hebrew" "_section_keywords_cache" "אין באמור ייעוץ משפטי" "התעלם|שכח|תתעלם" "_chunk_matches_filter" "_is_statute_N" "public_suggestions"; do
        local file_hit=$(grep -rl "$kw" "$TAU_RAG_DIR" --include="*.py" 2>/dev/null | head -1)
        if [ -n "$file_hit" ]; then
            printf "  ✅  %-30s  found in %s\n" "$kw" "$(basename $file_hit)"
        else
            printf "  ❌  %-30s  MISSING\n" "$kw"
        fi
    done
}

case "${1:-all}" in
    all|deploy)
        clear_cache
        stop_server
        start_server
        upload_data
        run_tests
        ok "Deployment complete."
        ;;
    restart)
        clear_cache
        stop_server
        start_server
        # v2.x — after restart, auto-upload if the index is empty.
        # The Hebrew preset doesn't persist the index across restarts,
        # so without this step every restart leaves the server with 0
        # docs and all regression tests return EMPTY.
        if index_is_empty; then
            warn "Index is empty after restart — running upload to repopulate..."
            upload_data
        else
            ok "Index already populated — skipping upload"
        fi
        ;;
    upload)
        upload_data
        ;;
    test)
        run_tests
        ;;
    status)
        show_status
        ;;
    stop)
        stop_server
        ok "Server stopped."
        ;;
    *)
        echo "Usage: $0 [all|restart|upload|test|status|stop]"
        exit 1
        ;;
esac