#!/bin/bash

# Model Switching Script for Dictation Service
# Allows easy switching between different speech recognition models

DICTATION_DIR="/mnt/storage/Development/dictation-service"
SHARED_MODELS_DIR="$HOME/.shared/models/vosk-models"
ENHANCED_SCRIPT="$DICTATION_DIR/src/dictation_service/ai_dictation_simple.py"

echo "=== Dictation Model Switcher ==="
echo ""

# Available models
declare -A MODELS=(
    ["small"]="vosk-model-small-en-us-0.15 (40MB) - Fast, Basic Accuracy"
    ["lgraph"]="vosk-model-en-us-0.22-lgraph (128MB) - Good Balance"
    ["full"]="vosk-model-en-us-0.22 (1.8GB) - Best Accuracy"
)

# Show current model
if [ -f "$ENHANCED_SCRIPT" ]; then
    CURRENT_MODEL=$(grep "MODEL_NAME = " "$ENHANCED_SCRIPT" | cut -d'"' -f2)
    echo "Current Model: $CURRENT_MODEL"
    echo ""
fi

# Show available options
echo "Available Models:"
for key in "${!MODELS[@]}"; do
    echo "  $key) ${MODELS[$key]}"
done
echo ""

# Interactive selection
read -p "Select model (small/lgraph/full): " choice

case $choice in
    small|s|S)
        NEW_MODEL="vosk-model-small-en-us-0.15"
        ;;
    lgraph|l|L)
        NEW_MODEL="vosk-model-en-us-0.22-lgraph"
        ;;
    full|f|F)
        NEW_MODEL="vosk-model-en-us-0.22"
        ;;
    *)
        echo "Invalid choice. Current model unchanged."
        exit 1
        ;;
esac

echo ""
echo "Switching to: $NEW_MODEL"

# Check if model directory exists
if [ ! -d "$SHARED_MODELS_DIR/$NEW_MODEL" ]; then
    echo "Error: Model directory $NEW_MODEL not found in $SHARED_MODELS_DIR!"
    echo "Available models:"
    ls -la "$SHARED_MODELS_DIR/"
    exit 1
fi

# Update the script
if [ -f "$ENHANCED_SCRIPT" ]; then
    # Create backup
    cp "$ENHANCED_SCRIPT" "$ENHANCED_SCRIPT.backup"
    echo "✓ Created backup of enhanced_dictation.py"

    # Update model name
    sed -i "s/MODEL_NAME = \".*\"/MODEL_NAME = \"$NEW_MODEL\"/" "$ENHANCED_SCRIPT"
    echo "✓ Updated model in ai_dictation_simple.py"

    # Show model comparison
    echo ""
    echo "Model Comparison:"
    echo "┌─────────────────────────────────────┬──────────┬──────────────┐"
    echo "│ Model                              │ Size     │ WER (lower)   │"
    echo "├─────────────────────────────────────┼──────────┼──────────────┤"
    echo "│ vosk-model-small-en-us-0.15        │ 40MB     │ ~15-20        │"
    echo "│ vosk-model-en-us-0.22-lgraph       │ 128MB    │ 7.82          │"
    echo "│ vosk-model-en-us-0.22              │ 1.8GB    │ 5.69          │"
    echo "└─────────────────────────────────────┴──────────┴──────────────┘"

    echo ""
    echo "Restarting dictation service..."
    systemctl --user restart dictation.service

    # Wait and show status
    sleep 3
    if systemctl --user is-active --quiet dictation.service; then
        echo "✓ Dictation service restarted successfully!"
        echo "✓ Now using: $NEW_MODEL"
        echo ""
        echo "Press Alt+D to test the new model!"
    else
        echo "⚠ Service restart failed. Check logs:"
        echo "  journalctl --user -u dictation.service -f"
    fi

else
    echo "Error: enhanced_dictation.py not found!"
    exit 1
fi

echo ""
echo "To restore backup:"
echo "  cp $ENHANCED_SCRIPT.backup $ENHANCED_SCRIPT"
echo "  systemctl --user restart dictation.service"