#!/bin/bash

# Dynamic Model Selector
# Intelligently selects models based on task type, availability, and capabilities

# --- Model Capability Database ---

# Model database using simple variables (compatible with older bash)
MODEL_DB_DIR="${LOG_DIR:-/tmp/ai_thinking}/model_db"

# Initialize model database with simple variables
init_model_database() {
    # Create model database directory
    mkdir -p "$MODEL_DB_DIR"

    # Model capabilities by task type (using file-based storage for compatibility)

    # === CONFIGURED MODELS ===
    cat > "$MODEL_DB_DIR/llama3_8b_instruct_q4_K_M" << 'EOF'
coding=0.8
reasoning=0.9
creative=0.7
size=8
speed=0.8
EOF

    cat > "$MODEL_DB_DIR/phi3_3_8b_mini_4k_instruct_q4_K_M" << 'EOF'
coding=0.7
reasoning=0.8
creative=0.6
size=3.8
speed=0.9
EOF

    cat > "$MODEL_DB_DIR/deepseek_r1_1_5b" << 'EOF'
coding=0.6
reasoning=0.9
creative=0.5
size=1.5
speed=0.95
EOF

    cat > "$MODEL_DB_DIR/gemma3n_e2b" << 'EOF'
coding=0.8
reasoning=0.8
creative=0.8
size=2
speed=0.85
EOF

    cat > "$MODEL_DB_DIR/dolphin3_latest" << 'EOF'
coding=0.6
reasoning=0.7
creative=0.8
size=7
speed=0.7
EOF

    # === ADDITIONAL MODELS FROM OLLAMA LIST ===

    # Llama 3.1 - Newer version, should be similar to Llama 3 but potentially better
    cat > "$MODEL_DB_DIR/llama3_1_8b" << 'EOF'
coding=0.82
reasoning=0.92
creative=0.72
size=8
speed=0.82
EOF

    # DeepSeek R1 7B - Larger reasoning model
    cat > "$MODEL_DB_DIR/deepseek_r1_7b" << 'EOF'
coding=0.65
reasoning=0.95
creative=0.55
size=7
speed=0.7
EOF

    # Gemma 3N Latest - Larger version of e2b
    cat > "$MODEL_DB_DIR/gemma3n_latest" << 'EOF'
coding=0.82
reasoning=0.82
creative=0.82
size=7.5
speed=0.8
EOF

    # Gemma 3 4B - Different model family
    cat > "$MODEL_DB_DIR/gemma3_4b" << 'EOF'
coding=0.75
reasoning=0.78
creative=0.75
size=4
speed=0.85
EOF

    # Qwen2.5 7B - Alibaba model, general purpose
    cat > "$MODEL_DB_DIR/qwen2_5_7b" << 'EOF'
coding=0.78
reasoning=0.85
creative=0.7
size=7
speed=0.75
EOF

    # Qwen3 8B - Latest Qwen model
    cat > "$MODEL_DB_DIR/qwen3_8b" << 'EOF'
coding=0.8
reasoning=0.88
creative=0.72
size=8
speed=0.78
EOF

    # Qwen3 4B - Smaller Qwen model
    cat > "$MODEL_DB_DIR/qwen3_4b" << 'EOF'
coding=0.75
reasoning=0.82
creative=0.68
size=4
speed=0.85
EOF

    # Qwen3 1.7B - Smallest Qwen model
    cat > "$MODEL_DB_DIR/qwen3_1_7b" << 'EOF'
coding=0.65
reasoning=0.7
creative=0.6
size=1.7
speed=0.95
EOF

    # DeepScaler - Performance optimization focus
    cat > "$MODEL_DB_DIR/deepscaler_latest" << 'EOF'
coding=0.7
reasoning=0.8
creative=0.65
size=3.6
speed=0.88
EOF

    # Yasser Qwen2.5 - Fine-tuned variant
    cat > "$MODEL_DB_DIR/yasserrmd_Qwen2_5_7B_Instruct_1M_latest" << 'EOF'
coding=0.82
reasoning=0.9
creative=0.75
size=7
speed=0.75
EOF

    # Nomic Embed Text - Specialized for embeddings, not general tasks
    cat > "$MODEL_DB_DIR/nomic_embed_text_latest" << 'EOF'
coding=0.1
reasoning=0.1
creative=0.1
size=0.274
speed=0.95
EOF
}

# Get model capability score
get_model_capability() {
    local model_key="$1"
    local task_type="$2"

    # Convert model name to filename-friendly format
    local safe_name=$(echo "$model_key" | tr ':' '_' | tr '.' '_')
    local db_file="$MODEL_DB_DIR/$safe_name"

    if [ -f "$db_file" ]; then
        grep "^${task_type}=" "$db_file" | cut -d'=' -f2
    else
        echo "0.5"  # Default capability score
    fi
}

# Get model size
get_model_size() {
    local model_key="$1"
    local safe_name=$(echo "$model_key" | tr ':' '_' | tr '.' '_')
    local db_file="$MODEL_DB_DIR/$safe_name"

    if [ -f "$db_file" ]; then
        grep "^size=" "$db_file" | cut -d'=' -f2
    else
        echo "5"  # Default size
    fi
}

# Get model speed
get_model_speed() {
    local model_key="$1"
    local safe_name=$(echo "$model_key" | tr ':' '_' | tr '.' '_')
    local db_file="$MODEL_DB_DIR/$safe_name"

    if [ -f "$db_file" ]; then
        grep "^speed=" "$db_file" | cut -d'=' -f2
    else
        echo "0.5"  # Default speed
    fi
}

# --- Model Discovery ---

# Get list of available models
get_available_models() {
    ollama list 2>/dev/null | tail -n +2 | awk '{print $1}' | sort
}

# Check if a model is available
is_model_available() {
    local model="$1"
    ollama list 2>/dev/null | grep -q "^${model}\s"
}

# --- Task Type Classification ---

# Classify task type from prompt and mechanism
classify_task_type() {
    local prompt="$1"
    local mechanism="$2"

    # Task type classification based on mechanism
    case "$mechanism" in
        "puzzle")
            echo "coding"
            ;;
        "socratic")
            echo "reasoning"
            ;;
        "exploration")
            echo "reasoning"
            ;;
        "consensus")
            echo "reasoning"
            ;;
        "critique")
            echo "reasoning"
            ;;
        "synthesis")
            echo "reasoning"
            ;;
        "peer-review")
            echo "reasoning"
            ;;
        *)
            # Fallback to keyword-based classification
            if echo "$prompt" | grep -q -i "code\|algorithm\|function\|program\|implement"; then
                echo "coding"
            elif echo "$prompt" | grep -q -i "write\|story\|creative\|poem\|essay"; then
                echo "creative"
            else
                echo "reasoning"
            fi
            ;;
    esac
}

# --- Model Selection Logic ---

# Select best model for task
select_best_model() {
    local task_type="$1"
    local available_models="$2"
    local preferred_models="$3"

    local best_model=""
    local best_score=0

    # First, try preferred models if available
    if [ -n "$preferred_models" ]; then
        for model in $preferred_models; do
            if echo "$available_models" | grep -q "^${model}$" && is_model_available "$model"; then
                local capability_score=$(get_model_capability "$model" "$task_type")
                local speed_score=$(get_model_speed "$model")
                local model_size=$(get_model_size "$model")
                local size_score=$(echo "scale=2; $model_size / 10" | bc -l 2>/dev/null || echo "0.5")

                # Calculate weighted score (capability is most important)
                local total_score=$(echo "scale=2; ($capability_score * 0.6) + ($speed_score * 0.3) + ($size_score * 0.1)" | bc -l 2>/dev/null || echo "0.5")

                if (( $(echo "$total_score > $best_score" | bc -l 2>/dev/null || echo "0") )); then
                    best_score=$total_score
                    best_model=$model
                fi
            fi
        done
    fi

    # If no preferred model is good, find best available model
    if [ -z "$best_model" ]; then
        for model in $available_models; do
            if is_model_available "$model"; then
                local capability_score=$(get_model_capability "$model" "$task_type")
                local speed_score=$(get_model_speed "$model")
                local model_size=$(get_model_size "$model")
                local size_score=$(echo "scale=2; $model_size / 10" | bc -l 2>/dev/null || echo "0.5")

                local total_score=$(echo "scale=2; ($capability_score * 0.6) + ($speed_score * 0.3) + ($size_score * 0.1)" | bc -l 2>/dev/null || echo "0.5")

                if (( $(echo "$total_score > $best_score" | bc -l 2>/dev/null || echo "0") )); then
                    best_score=$total_score
                    best_model=$model
                fi
            fi
        done
    fi

    if [ -n "$best_model" ]; then
        echo "Selected model: $best_model (score: $best_score, task: $task_type)" 1>&2
        echo "$best_model"
    else
        echo "No suitable model found" >&2
        echo ""
    fi
}

# --- Main Selection Function ---

# Smart model selection
select_model_for_task() {
    local prompt="$1"
    local mechanism="$2"
    local preferred_models="$3"

    # Initialize database
    init_model_database

    # Get available models
    local available_models
    available_models=$(get_available_models)

    if [ -z "$available_models" ]; then
        echo "No models available via Ollama" >&2
        echo ""
        return 1
    fi

    # Classify task type
    local task_type
    task_type=$(classify_task_type "$prompt" "$mechanism")

    # Select best model
    local selected_model
    selected_model=$(select_best_model "$task_type" "$available_models" "$preferred_models")

    if [ -n "$selected_model" ]; then
        echo "$selected_model"
        return 0
    else
        echo ""
        return 1
    fi
}

# --- Utility Functions ---

# Get model info
get_model_info() {
    local model="$1"

    echo "Model: $model"
    echo "Size: $(get_model_size "$model")B"
    echo "Speed: $(get_model_speed "$model")"
    echo "Coding: $(get_model_capability "$model" "coding")"
    echo "Reasoning: $(get_model_capability "$model" "reasoning")"
    echo "Creative: $(get_model_capability "$model" "creative")"
}

# Export functions
export -f init_model_database
export -f get_available_models
export -f is_model_available
export -f classify_task_type
export -f select_best_model
export -f select_model_for_task
export -f get_model_info
export -f get_model_capability
export -f get_model_size
export -f get_model_speed