#!/bin/bash

# A script to automatically start the server and run an interactive chat session.
# It now allows the user to specify a model and a FAISS retriever option for each turn.
#
# REQUIREMENTS:
# - 'jq' must be installed (e.g., sudo apt-get install jq).

BASE_URL="http://127.0.0.1:8000"
DEFAULT_MODEL="deepseek"
CURRENT_MODEL="" # The model used in the last turn

# --- 1. Check for Dependencies ---
if ! command -v jq &> /dev/null
then
    echo "❌ 'jq' is not installed. Please install it to run this script."
    exit 1
fi

# --- 2. Start the FastAPI Server in the Background ---
echo "--- Starting AI Hub Server ---"
uvicorn app.main:app --host 127.0.0.1 --port 8000 &
SERVER_PID=$!

# Define a cleanup function to kill the server on exit
cleanup() {
    echo ""
    echo "--- Shutting Down Server (PID: $SERVER_PID) ---"
    kill $SERVER_PID
}
# Register the cleanup function to run when the script exits (e.g., Ctrl+C or typing 'exit')
trap cleanup EXIT

echo "Server started with PID: $SERVER_PID. Waiting for it to initialize..."
sleep 5 # Wait for the server to be ready

# --- 3. Create a New Conversation Session ---
echo "--- Starting a new conversation session... ---"
# FIX: Added a trailing slash to the /sessions endpoint to avoid a 307 redirect
SESSION_DATA=$(curl -s -X POST "$BASE_URL/sessions/" \
    -H "Content-Type: application/json" \
    -d '{"user_id": "local_user", "model": "'"$DEFAULT_MODEL"'"}' \
    -w '\n%{http_code}') # Add a new line and the status code

# Extract body and status code
HTTP_CODE=$(echo "$SESSION_DATA" | tail -n1)
SESSION_DATA_BODY=$(echo "$SESSION_DATA" | head -n-1)

if [ "$HTTP_CODE" -ne 200 ]; then
    echo "❌ Failed to create a session (HTTP $HTTP_CODE). Server might not have started correctly."
    echo "Response body: $SESSION_DATA_BODY"
    exit 1
fi

SESSION_ID=$(echo "$SESSION_DATA_BODY" | jq -r '.id')
if [ -z "$SESSION_ID" ] || [ "$SESSION_ID" == "null" ]; then
    echo "❌ Failed to create a session. Server response did not contain an ID."
    echo "Response body: $SESSION_DATA_BODY"
    exit 1
fi

# Set the initial model
CURRENT_MODEL="$DEFAULT_MODEL"

echo "✅ Session created with ID: $SESSION_ID. The initial model is '$CURRENT_MODEL'."
echo "--------------------------------------------------"
echo "To switch models, type your message like this: [gemini] <Your message>"
echo "To enable the FAISS retriever, add the flag: [faiss] <Your message>"
echo "You can combine them: [gemini] [faiss] <Your message>"
echo "Type 'exit' or 'quit' to end."
echo "--------------------------------------------------"

# --- 4. Start the Interactive Chat Loop ---
while true; do
    read -p "You: " user_input

    if [[ "$user_input" == "exit" || "$user_input" == "quit" ]]; then
        break
    fi

    # Initialize variables for this turn
    PROMPT_TEXT="$user_input"
    MODEL_TO_USE="$CURRENT_MODEL"
    LOAD_FAISS_RETRIEVER="false"

    # Check for and process the model switch
    # This regex ensures we only match valid models and they must appear first
    if [[ "$PROMPT_TEXT" =~ ^\[(deepseek|gemini)\]\ (.*)$ ]]; then
        MODEL_TO_USE="${BASH_REMATCH[1]}"
        PROMPT_TEXT="${BASH_REMATCH[2]}"
        # Update the current model for the next prompt
        CURRENT_MODEL="$MODEL_TO_USE"
    fi

    # Check for and process the FAISS flag
    # This check is separate to avoid misinterpreting "faiss" as a model
    if [[ "$PROMPT_TEXT" =~ \[faiss\] ]]; then
        LOAD_FAISS_RETRIEVER="true"
        # Use sed to remove the [faiss] flag and trim any leading/trailing spaces
        PROMPT_TEXT=$(echo "$PROMPT_TEXT" | sed 's/\[faiss\]//' | xargs)
    fi
    
    # Check if the prompt is empty after processing flags and models
    if [ -z "$PROMPT_TEXT" ]; then
        echo "Please enter a message."
        continue
    fi
    
    # Construct the JSON payload with the model, prompt, and FAISS flag
    # Note the use of --argjson to pass the boolean value correctly
    json_payload=$(jq -n \
        --arg prompt "$PROMPT_TEXT" \
        --arg model "$MODEL_TO_USE" \
        --argjson faiss_flag "$LOAD_FAISS_RETRIEVER" \
        '{"prompt": $prompt, "model": $model, "load_faiss_retriever": $faiss_flag}')

    echo "Payload: $json_payload" # Optional: for debugging

    ai_response_json=$(curl -s -X POST "$BASE_URL/sessions/$SESSION_ID/chat" \
        -H "Content-Type: application/json" \
        -d "$json_payload")
    
    # Check if the response is valid JSON
    if ! echo "$ai_response_json" | jq -e . >/dev/null; then
        echo "❌ AI: An error occurred or the server returned an invalid response."
        echo "Server response: $ai_response_json"
    else
        ai_answer=$(echo "$ai_response_json" | jq -r '.answer')
        model_used=$(echo "$ai_response_json" | jq -r '.model_used')
        echo "AI [$model_used] [faiss_enabled: $LOAD_FAISS_RETRIEVER]: $ai_answer"
    fi

done

# The 'trap' will automatically call the cleanup function when the loop breaks.