#!/bin/bash
#
# Fix Ani's Model Configuration
#
# This script helps diagnose and fix Ani's LLM configuration in Letta.
# Common issues: wrong model name, missing reasoning settings, outdated context window.
#
# Usage: ./fix-ani-model.sh [--check] [--apply]
#

COLOR_RESET="\033[0m"
COLOR_RED="\033[31m"
COLOR_GREEN="\033[32m"
COLOR_YELLOW="\033[33m"
COLOR_BLUE="\033[34m"
COLOR_PURPLE="\033[35m"

# Letta container
CONTAINER="${CONTAINER:-aster-0.16.4}"
# Ani's agent ID (verify with --check)
AGENT_ID="${AGENT_ID:-agent-e2b683bf-5b3e-4e0c-ac62-2bbb47ea8351}"

# API endpoint for model verification
MODELS_API="${MODELS_API:-https://api.synthetic.new/openai/v1/models}"

# Current working config for Kimi-K2.5-NVFP4
MODEL_NAME="kimi-k2.5"
CONTEXT_WINDOW=262144
MAX_TOKENS=55000
TEMPERATURE=0.9
ENABLE_REASONER=true
REASONING_EFFORT="high"
MAX_REASONING_TOKENS=75000
PARALLEL_TOOL_CALLS=true

log_info() {
    echo -e "${COLOR_BLUE}[INFO]${COLOR_RESET} $1"
}

log_success() {
    echo -e "${COLOR_GREEN}[OK]${COLOR_RESET} $1"
}

log_warn() {
    echo -e "${COLOR_YELLOW}[WARN]${COLOR_RESET} $1"
}

log_error() {
    echo -e "${COLOR_RED}[ERROR]${COLOR_RESET} $1"
}

log_header() {
    echo -e "${COLOR_PURPLE}$1${COLOR_RESET}"
}

check_container() {
    log_info "Checking Letta container: $CONTAINER"
    if ! docker inspect "$CONTAINER" &>/dev/null; then
        log_error "Container '$CONTAINER' not found"
        log_info "Available containers:"
        docker ps --format "{{.Names}}" | grep -i letta
        return 1
    fi
    log_success "Container found"
    return 0
}

get_current_config() {
    log_header "=== Current Ani Configuration ==="
    docker exec "$CONTAINER" psql -U letta -d letta -c "
        SELECT
            name,
            llm_config->>'model' as model,
            llm_config->>'context_window' as context,
            llm_config->>'temperature' as temp,
            llm_config->>'enable_reasoner' as reasoner,
            llm_config->>'reasoning_effort' as effort,
            llm_config->>'max_reasoning_tokens' as reasoning_tokens,
            llm_config->>'parallel_tool_calls' as parallel,
            llm_config->>'model_endpoint' as endpoint
        FROM agents
        WHERE id = '$AGENT_ID';
    " 2>&1
}

compare_with_api() {
    log_header "=== Comparing with API Specs ==="
    log_info "Fetching models from: $MODELS_API"

    MODELS_JSON=$(curl -s "$MODELS_API" 2>&1)

    if [ $? -ne 0 ]; then
        log_warn "Could not fetch models from API"
        return 1
    fi

    echo "$MODELS_JSON" | jq -r '.data[] | select(.id | contains("Kimi")) | {
        id: .id,
        context: .context_length
    }'
}

check_other_agents() {
    log_header "=== Other Agents for Comparison ==="
    docker exec "$CONTAINER" psql -U letta -d letta -c "
        SELECT
            name,
            llm_config->>'model' as model,
            llm_config->>'enable_reasoner' as reasoner
        FROM agents
        WHERE is_deleted = false
        ORDER BY name;
    " 2>&1
}

diagnose_issues() {
    log_header "=== Diagnosing Configuration Issues ==="

    # Get current model
    CURRENT_MODEL=$(docker exec "$CONTAINER" psql -U letta -d letta -t -c "
        SELECT llm_config->>'model' FROM agents WHERE id = '$AGENT_ID';
    " | tr -d '[:space:]')

    CURRENT_CONTEXT=$(docker exec "$CONTAINER" psql -U letta -d letta -t -c "
        SELECT llm_config->>'context_window' FROM agents WHERE id = '$AGENT_ID';
    " | tr -d '[:space:]')

    CURRENT_REASONER=$(docker exec "$CONTAINER" psql -U letta -d letta -t -c "
        SELECT llm_config->>'enable_reasoner' FROM agents WHERE id = '$AGENT_ID';
    " | tr -d '[:space:]')

    CURRENT_EFFORT=$(docker exec "$CONTAINER" psql -U letta -d letta -t -c "
        SELECT llm_config->>'reasoning_effort' FROM agents WHERE id = '$AGENT_ID';
    " | tr -d '[:space:]')

    CURRENT_RTOKENS=$(docker exec "$CONTAINER" psql -U letta -d letta -t -c "
        SELECT llm_config->>'max_reasoning_tokens' FROM agents WHERE id = '$AGENT_ID';
    " | tr -d '[:space:]')

    CURRENT_PARALLEL=$(docker exec "$CONTAINER" psql -U letta -d letta -t -c "
        SELECT llm_config->>'parallel_tool_calls' FROM agents WHERE id = '$AGENT_ID';
    " | tr -d '[:space:]')

    ISSUES=0

    # Check model name format
    if echo "$CURRENT_MODEL" | grep -q "^hf:"; then
        log_error "Model name has 'hf:' prefix: $CURRENT_MODEL (should be: $MODEL_NAME)"
        ((ISSUES++))
    elif [ "$CURRENT_MODEL" != "$MODEL_NAME" ]; then
        log_warn "Model name: $CURRENT_MODEL (expected: $MODEL_NAME)"
    else
        log_success "Model name: $CURRENT_MODEL"
    fi

    # Check context window
    if [ "$CURRENT_CONTEXT" != "$CONTEXT_WINDOW" ]; then
        log_warn "Context window: $CURRENT_CONTEXT (expected: $CONTEXT_WINDOW)"
    fi

    # Check reasoning settings
    if [ "$CURRENT_REASONER" != "true" ]; then
        log_error "Reasoning disabled: $CURRENT_REASONER (should be: $ENABLE_REASONER)"
        ((ISSUES++))
    fi

    if [ "$CURRENT_EFFORT" != "$REASONING_EFFORT" ]; then
        log_error "Reasoning effort: ${CURRENT_EFFORT:-null} (should be: $REASONING_EFFORT)"
        ((ISSUES++))
    fi

    if [ -z "$CURRENT_RTOKENS" ] || [ "$CURRENT_RTOKENS" = "0" ]; then
        log_error "Max reasoning tokens: ${CURRENT_RTOKENS:-null} (should be: $MAX_REASONING_TOKENS)"
        ((ISSUES++))
    elif [ "$CURRENT_RTOKENS" != "$MAX_REASONING_TOKENS" ]; then
        log_warn "Max reasoning tokens: $CURRENT_RTOKENS (recommended: $MAX_REASONING_TOKENS)"
    fi

    if [ "$CURRENT_PARALLEL" != "true" ]; then
        log_warn "Parallel tool calls: ${CURRENT_PARALLEL:-null} (recommended: $PARALLEL_TOOL_CALLS)"
    fi

    if [ $ISSUES -eq 0 ]; then
        log_success "No critical issues found!"
    else
        log_error "Found $ISSUES issue(s) - run with --apply to fix"
    fi

    return $ISSUES
}

apply_fix() {
    log_header "=== Applying Configuration Fix ==="
    log_warn "This will replace the entire llm_config for Ani"
    read -p "Continue? (y/N): " confirm

    if [ "$confirm" != "y" ] && [ "$confirm" != "Y" ]; then
        log_info "Aborted"
        return 1
    fi

    docker exec "$CONTAINER" psql -U letta -d letta -c "
UPDATE agents SET llm_config = '{
  \"tier\": null,
  \"model\": \"$MODEL_NAME\",
  \"effort\": null,
  \"handle\": \"openai-proxy/hf:nvidia/Kimi-K2.5-NVFP4\",
  \"strict\": false,
  \"verbosity\": null,
  \"max_tokens\": $MAX_TOKENS,
  \"temperature\": $TEMPERATURE,
  \"display_name\": \"hf:nvidia/Kimi-K2.5-NVFP4\",
  \"model_wrapper\": null,
  \"provider_name\": \"letta\",
  \"context_window\": $CONTEXT_WINDOW,
  \"model_endpoint\": \"http://172.17.0.1:4000/v1\",
  \"enable_reasoner\": $ENABLE_REASONER,
  \"response_format\": null,
  \"reasoning_effort\": \"$REASONING_EFFORT\",
  \"frequency_penalty\": null,
  \"provider_category\": \"base\",
  \"compatibility_type\": null,
  \"model_endpoint_type\": \"openai\",
  \"parallel_tool_calls\": $PARALLEL_TOOL_CALLS,
  \"max_reasoning_tokens\": $MAX_REASONING_TOKENS,
  \"put_inner_thoughts_in_kwargs\": false
}'::json
WHERE id = '$AGENT_ID'
RETURNING name, llm_config->>'model' as model, llm_config->>'enable_reasoner' as reasoner;
" 2>&1

    if [ $? -eq 0 ]; then
        log_success "Configuration updated!"
        log_info "Agent may need to reinitialize for changes to take effect"
    else
        log_error "Failed to update configuration"
        return 1
    fi
}

# Main
MODE="${1:-check}"

case "$MODE" in
    --check|-c)
        check_container && get_current_config
        echo ""
        compare_with_api
        echo ""
        check_other_agents
        echo ""
        diagnose_issues
        ;;
    --apply|-a)
        check_container && apply_fix
        ;;
    --model-only|-m)
        # Quick fix: just the model name
        docker exec "$CONTAINER" psql -U letta -d letta -c "
UPDATE agents SET llm_config = (llm_config::jsonb || jsonb_build_object('model', '$MODEL_NAME'))::json
WHERE id = '$AGENT_ID'
RETURNING name, llm_config->>'model' as model;
" 2>&1
        ;;
    --help|-h)
        echo "Fix Ani's Model Configuration"
        echo ""
        echo "Usage: $0 [OPTION]"
        echo ""
        echo "Options:"
        echo "  --check, -c       Show current config and diagnose issues (default)"
        echo "  --apply, -a       Apply recommended fix"
        echo "  --model-only, -m  Quick fix: only update model name"
        echo "  --help, -h        Show this help"
        echo ""
        echo "Environment variables:"
        echo "  CONTAINER      Letta container name (default: aster-0.16.4)"
        echo "  AGENT_ID       Ani's agent ID"
        echo "  MODELS_API     API endpoint for model verification"
        ;;
    *)
        log_error "Unknown option: $MODE"
        echo "Run --help for usage"
        exit 1
        ;;
esac