#!/bin/bash

# A script to automatically start the server and run an interactive chat session.
# It now allows the user to specify a model for each turn or use the previous one.
#
# 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... ---"
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 use the previous model, just type your message directly."
echo "Type 'exit' or 'quit' to end."
echo "--------------------------------------------------"

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

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

    # Check for model switch input pattern, e.g., "[model_name] <message>"
    if [[ "$user_input" =~ ^\[([a-zA-Z0-9]+)\]\ (.*)$ ]]; 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"
    else
        MODEL_TO_USE="$CURRENT_MODEL"
        PROMPT_TEXT="$user_input"
    fi
    
    # Construct the JSON payload with the model and prompt
    json_payload=$(jq -n \
        --arg prompt "$PROMPT_TEXT" \
        --arg model "$MODEL_TO_USE" \
        '{"prompt": $prompt, "model": $model}')

    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]: $ai_answer"
    fi

done

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