oliveowl

#!/bin/bash

# oliveowl - Terminal AI Chat Assistant

# --- Configuration ---
CONFIG_DIR="$HOME/.config/oliveowl"
ENV_FILE="$CONFIG_DIR/.env"
HISTORY_DIR="$CONFIG_DIR/history"
CONFIG_FILE="$CONFIG_DIR/config"
PROMPT_HISTORY_FILE="$CONFIG_DIR/prompt_history.txt"
MAX_PROMPT_HISTORY=50

# --- Dependency Checks ---
check_dependency() {
  if ! command -v "$1" &>/dev/null; then
    echo "Error: Required command '$1' not found. Please install it." >&2
    exit 1
  fi
}

check_clipboard_tool() {
  # Try wl-copy first: check if command exists AND if it runs without error (e.g., connects to Wayland)
  if command -v wl-copy &>/dev/null && printf '' | wl-copy &>/dev/null; then
    CLIPBOARD_TOOL="wl-copy"
    echo "Using clipboard tool: wl-copy" # Optional: for debugging
  # If wl-copy fails or isn't present, try xclip
  elif command -v xclip &>/dev/null; then
    CLIPBOARD_TOOL="xclip -selection clipboard"
    echo "Using clipboard tool: xclip" # Optional: for debugging
  else
    # Error: Neither working wl-copy nor xclip found
    echo "Error: No functional clipboard tool found (need working 'wl-copy' for Wayland or 'xclip' for X11)." >&2
    exit 1
  fi
  # The debug echo is now inside the if/elif blocks
}

check_dependency "curl"
check_dependency "jq"
check_dependency "fzf"
check_dependency "bat"
check_dependency "gum" # Added gum dependency check
check_clipboard_tool   # Sets $CLIPBOARD_TOOL

# --- Initial Setup ---
setup_config() {
  mkdir -p "$CONFIG_DIR"
  mkdir -p "$HISTORY_DIR"
  
  # Create prompt history file if it doesn't exist
  if [ ! -f "$PROMPT_HISTORY_FILE" ]; then
    touch "$PROMPT_HISTORY_FILE"
  fi

  if [ ! -f "$ENV_FILE" ]; then
    echo "Creating initial config file: $ENV_FILE"
    echo "# Add your API keys here" >"$ENV_FILE"
    echo "GEMINI_API_KEY=" >>"$ENV_FILE"
    echo "OPENROUTER_API_KEY=" >>"$ENV_FILE"
    echo "OPENAI_API_KEY=" >>"$ENV_FILE"
    echo "CEREBRAS_API_KEY=" >>"$ENV_FILE"
    echo "# Ollama base URL (e.g., http://localhost:11434)" >>"$ENV_FILE"
    echo "OLLAMA_BASE_URL=" >>"$ENV_FILE"
    echo "# Custom OpenAI-compatible providers" >>"$ENV_FILE"
    echo "# Format: CUSTOM_PROVIDER_<PROVIDER_NAME>_URL=https://api.example.com/v1" >>"$ENV_FILE"
    echo "# Format: CUSTOM_PROVIDER_<PROVIDER_NAME>_KEY=your_api_key" >>"$ENV_FILE"
    echo "# Example:" >>"$ENV_FILE"
    echo "# CUSTOM_PROVIDER_MYPROVIDER_URL=https://api.openai.com/v1" >>"$ENV_FILE"
    echo "# CUSTOM_PROVIDER_MYPROVIDER_KEY=your_api_key" >>"$ENV_FILE"
    chmod 600 "$ENV_FILE" # Secure the file
    echo "Please edit $ENV_FILE and add your API keys and/or Ollama URL."
    # Optionally exit here or prompt user
  fi

  # Load environment variables
  set -a # Automatically export all variables
  # Check if ENV_FILE exists before sourcing
  if [ -f "$ENV_FILE" ]; then
    source "$ENV_FILE"
  else
    echo "Warning: Environment file $ENV_FILE not found." >&2
  fi
  set +a

  # Check if keys are set (basic check)
  if [ -z "$GEMINI_API_KEY" ] && [ -z "$OPENROUTER_API_KEY" ] && [ -z "$OPENAI_API_KEY" ] && [ -z "$CEREBRAS_API_KEY" ] && [ -z "$OLLAMA_BASE_URL" ] && ! grep -q "^CUSTOM_PROVIDER_" "$ENV_FILE"; then
    echo "Warning: No API keys (Gemini, OpenRouter, OpenAI, Cerebras) or OLLAMA_BASE_URL or custom providers found in $ENV_FILE. Please configure at least one." >&2
    # Decide if we should exit or continue
  fi

  # Load or create config file
  if [ ! -f "$CONFIG_FILE" ]; then
    echo "Creating default settings file: $CONFIG_FILE"
    echo "API_PROVIDER=" >"$CONFIG_FILE" # Will be 'gemini' or 'openrouter'
    echo "MODEL=" >>"$CONFIG_FILE"
    echo "EDITOR=\"vi\"" >>"$CONFIG_FILE" # Add default editor
  fi

  # Load config settings
  # Check if CONFIG_FILE exists before sourcing
  if [ -f "$CONFIG_FILE" ]; then
    source "$CONFIG_FILE"
  else
    echo "Warning: Config file $CONFIG_FILE not found. Using defaults." >&2
  fi

  # Set default editor if not found in config
  EDITOR="${EDITOR:-vi}"

  # NOTE: Configuration completeness check moved to main()
  # Ensure EDITOR command exists - Moved to after setup_config call before main()
}

# --- Configuration Function ---
fetch_openrouter_models() {
  local models_json
  echo "Fetching models from OpenRouter..." >&2
  models_json=$(curl -s https://openrouter.ai/api/v1/models)
  if [ $? -ne 0 ] || [ -z "$models_json" ]; then
    echo "Error: Failed to fetch models from OpenRouter." >&2
    return 1
  fi
  # Extract model IDs, filter out duplicates, sort
  echo "$models_json" | jq -r '.data[].id' | sort -u
}

fetch_openai_models() {
  if [ -z "$OPENAI_API_KEY" ]; then
    echo "Error: OPENAI_API_KEY not set in $ENV_FILE. Cannot fetch models." >&2
    return 1
  fi

  local models_json
  echo "Fetching models from OpenAI..." >&2
  models_json=$(curl -s -H "Authorization: Bearer $OPENAI_API_KEY" https://api.openai.com/v1/models)

  if [ $? -ne 0 ] || [ -z "$models_json" ]; then
    echo "Error: Failed to fetch models from OpenAI." >&2
    return 1
  fi

  # Check for API errors
  local openai_error=$(echo "$models_json" | jq -r '.error.message // empty')
  if [ -n "$openai_error" ]; then
    echo "Warning: OpenAI API Error when fetching models: $openai_error" >&2
    return 1
  fi

  # Check if data array exists (like OpenAI)
  if ! echo "$models_json" | jq -e '.data | type == "array"' >/dev/null 2>&1; then
    echo "Warning: Invalid JSON or missing 'data' array in OpenAI models response." >&2
    # echo "OpenAI response: $models_json" >&2 # For debugging
    return 1
  fi

  # Extract model IDs, filter out duplicates, sort
  local extracted_models
  extracted_models=$(echo "$models_json" | jq -r '.data[].id // empty' | sort -u | grep .)

  if [ -z "$extracted_models" ]; then
    echo "Warning: No models found in OpenAI response or failed to parse names." >&2
    return 1
  fi

  echo "$extracted_models"
  return 0 # Success
}

fetch_cerebras_models() {
  if [ -z "$CEREBRAS_API_KEY" ]; then
    echo "Error: CEREBRAS_API_KEY not set in $ENV_FILE. Cannot fetch models." >&2
    return 1
  fi

  local models_json
  echo "Fetching models from Cerebras..." >&2
  models_json=$(curl -s -H "Authorization: Bearer $CEREBRAS_API_KEY" https://api.cerebras.ai/v1/models)

  if [ $? -ne 0 ] || [ -z "$models_json" ]; then
    echo "Error: Failed to fetch models from Cerebras." >&2
    return 1
  fi

  # Check for API errors
  local cerebras_error=$(echo "$models_json" | jq -r '.error.message // empty')
  if [ -n "$cerebras_error" ]; then
    echo "Warning: Cerebras API Error when fetching models: $cerebras_error" >&2
    return 1
  fi

  # Check if data array exists (like OpenAI)
  if ! echo "$models_json" | jq -e '.data | type == "array"' >/dev/null 2>&1; then
    echo "Warning: Invalid JSON or missing 'data' array in Cerebras models response." >&2
    # echo "Cerebras response: $models_json" >&2 # For debugging
    return 1
  fi

  # Extract model IDs, filter out duplicates, sort
  local extracted_models
  extracted_models=$(echo "$models_json" | jq -r '.data[].id // empty' | sort -u | grep .)

  if [ -z "$extracted_models" ]; then
    echo "Warning: No models found in Cerebras response or failed to parse names." >&2
    return 1
  fi

  echo "$extracted_models"
  return 0 # Success
}

# Function to fetch models from a local Ollama instance
fetch_ollama_models() {
  local ollama_url_base="${OLLAMA_BASE_URL:-http://localhost:11434}"
  # Remove trailing slash if present
  ollama_url_base=$(echo "$ollama_url_base" | sed 's:/*$::')
  local models_api_url="${ollama_url_base}/api/tags"
  local models_json

  echo "Fetching models from Ollama at $ollama_url_base..." >&2
  # Timeout curl after 5 seconds for local requests to avoid long hangs
  models_json=$(curl --connect-timeout 5 -s "$models_api_url")

  if [ $? -ne 0 ]; then
    echo "Warning: curl command failed for Ollama. Is Ollama running and accessible at $ollama_url_base?" >&2
    return 1
  fi

  if [ -z "$models_json" ]; then
    echo "Warning: Received empty response from Ollama models API." >&2
    return 1
  fi

  # Check for Ollama API errors (e.g., if the endpoint exists but returns an error object)
  local ollama_error=$(echo "$models_json" | jq -r '.error // empty')
  if [ -n "$ollama_error" ]; then
    echo "Warning: Ollama API Error when fetching models: $ollama_error" >&2
    return 1
  fi

  # Check if the 'models' array exists and is an array
  if ! echo "$models_json" | jq -e '.models | type == "array"' >/dev/null 2>&1; then
    echo "Warning: Invalid JSON or missing 'models' array in Ollama response." >&2
    # echo "Ollama response: $models_json" >&2 # For debugging
    return 1
  fi

  # Extract model names, filter out duplicates, sort
  local extracted_models
  extracted_models=$(echo "$models_json" | jq -r '.models[].name // empty' | sort -u | grep .) # grep . to remove empty lines if any

  if [ -z "$extracted_models" ]; then
    echo "Warning: No models found in Ollama response or failed to parse names." >&2
    return 1 # Indicate no models found or parse error
  fi

  echo "$extracted_models"
  return 0 # Success
}

configure_settings() {
  echo "--- OliveOwl Configuration ---"

  # 1. Select Configuration Option
  local config_options="Change Model\nAdd Provider\nChange Editor\nEdit Instructions"
  local selected_option=$(echo -e "$config_options" | fzf --prompt="Select Configuration Option: " --height=5 --layout=reverse)

  if [ -z "$selected_option" ]; then
    echo "Configuration cancelled."
    return 1
  fi

  # Handle Add Provider option
  if [[ "$selected_option" == "Add Provider" ]]; then
    add_provider
    # After adding provider, return to continue with previous flow
    return 0
  elif [[ "$selected_option" == "Change Editor" ]]; then
    # Handle Change Editor option
    local current_editor=${EDITOR:-vi}
    local new_editor
    new_editor=$(gum input --header "Configure Editor" --placeholder "Enter preferred editor command (current: $current_editor)..." --value "$current_editor")
    if [ $? -eq 0 ] && [ -n "$new_editor" ]; then
      EDITOR="$new_editor"
    fi
    # Save editor configuration
    echo "Saving editor configuration..."
    if [ -f "$CONFIG_FILE" ]; then
      # Update existing config file with new editor
      sed -i "s|^EDITOR=.*|EDITOR=\"$EDITOR\"|" "$CONFIG_FILE"
    else
      # Create new config file with editor setting
      echo "API_PROVIDER=\"$API_PROVIDER\"" >"$CONFIG_FILE"
      echo "MODEL=\"$MODEL\"" >>"$CONFIG_FILE"
      echo "EDITOR=\"$EDITOR\"" >>"$CONFIG_FILE"
    fi
    echo "Editor configuration saved."
    return 0
  elif [[ "$selected_option" == "Edit Instructions" ]]; then
    # Handle Edit Instructions option
    local instructions_file="$CONFIG_DIR/extra_instructions.txt"
    "$EDITOR" "$instructions_file"
    echo "Instructions saved."
    return 0
  fi

  # 2. Select API Provider (for Change Model or Add Provider paths)
  local providers="Gemini\nOpenRouter\nOpenAI\nOllama\nCerebras"
  # Add custom providers to the list
  while IFS= read -r line; do
    if [[ "$line" =~ ^CUSTOM_PROVIDER_([A-Z0-9_]+)_URL= ]]; then
      provider_name="${BASH_REMATCH[1]}"
      providers="$providers\n$provider_name"
    fi
  done < "$ENV_FILE"
  
  API_PROVIDER=$(echo -e "$providers" | fzf --prompt="Select API Provider: " --height=7 --layout=reverse)

  if [ -z "$API_PROVIDER" ]; then
    echo "Configuration cancelled."
    return 1
  fi

  # Check if API key exists for selected provider and prompt to add if missing
  local api_key_exists=true
  case "$API_PROVIDER" in
    "Gemini")
      if [ -z "$GEMINI_API_KEY" ]; then
        api_key_exists=false
      fi
      ;;
    "OpenRouter")
      if [ -z "$OPENROUTER_API_KEY" ]; then
        api_key_exists=false
      fi
      ;;
    "OpenAI")
      if [ -z "$OPENAI_API_KEY" ]; then
        api_key_exists=false
      fi
      ;;
    "Cerebras")
      if [ -z "$CEREBRAS_API_KEY" ]; then
        api_key_exists=false
      fi
      ;;
"Ollama")
       # For Ollama, check if URL is set
       if [ -z "$OLLAMA_BASE_URL" ]; then
         api_key_exists=false
       fi
       ;;
    *)
       # Check if it's a custom provider
       if [[ "$API_PROVIDER" =~ ^[A-Z0-9_]+$ ]]; then
         local provider_url_var="CUSTOM_PROVIDER_${API_PROVIDER}_URL"
         local provider_key_var="CUSTOM_PROVIDER_${API_PROVIDER}_KEY"
         if [ -z "${!provider_url_var}" ] || [ -z "${!provider_key_var}" ]; then
           api_key_exists=false
         fi
       fi
       ;;
   esac

  if [ "$api_key_exists" = false ]; then
    echo "API key/URL not found for $API_PROVIDER."
    local add_key_choice
    add_key_choice=$(gum choose --header "Add $API_PROVIDER configuration?" "Yes" "No")
    
    if [ "$add_key_choice" = "Yes" ]; then
      if [[ "$API_PROVIDER" =~ ^[A-Z0-9_]+$ ]]; then
        # It's a custom provider, need to add it
        add_custom_provider
      else
        add_provider_key "$API_PROVIDER"
      fi
      # Reload environment variables
      set -a
      if [ -f "$ENV_FILE" ]; then
        source "$ENV_FILE"
      fi
      set +a
    else
      # If user chooses not to add key, continue but warn them
      echo "Warning: You may not be able to use $API_PROVIDER without configuration."
    fi
  fi

  # 2. Select Model
  local model_list=""
  local fetched_models=""
  local fzf_input=""
  local manual_entry_option="-- Manual Model Entry --"

  case "$API_PROVIDER" in
  "Gemini")
    # Check for Gemini Key first
    if [ -z "$GEMINI_API_KEY" ]; then
      echo "Warning: GEMINI_API_KEY not set in $ENV_FILE. Cannot fetch models." >&2
      # Proceed with only manual entry option
    else
      echo "Fetching Gemini models..."
      local gemini_url="https://generativelanguage.googleapis.com/v1beta/models?key=$GEMINI_API_KEY"
      local gemini_response=$(curl -s "$gemini_url")
      if [ $? -ne 0 ]; then
        echo "Warning: curl command failed for Gemini. Using manual entry only." >&2
      else
        local gemini_error=$(echo "$gemini_response" | jq -r '.error.message // empty')
        if [ -n "$gemini_error" ]; then
          echo "Warning: Gemini API Error: $gemini_error. Using manual entry only." >&2
        elif echo "$gemini_response" | jq -e '.models' >/dev/null 2>&1; then
          fetched_models=$(echo "$gemini_response" | jq -r '.models[].name')
          if [ -z "$fetched_models" ]; then
            echo "Warning: No models found in Gemini response. Using manual entry only." >&2
          fi
        else
          echo "Warning: Invalid JSON or missing 'models' in Gemini response. Using manual entry only." >&2
        fi
      fi
    fi
    ;;
  "OpenRouter")
    # OpenRouter key not strictly needed for model listing, but good practice to have it set for actual use later
    if [ -z "$OPENROUTER_API_KEY" ]; then
      echo "Warning: OPENROUTER_API_KEY not set in $ENV_FILE. You will need it to use OpenRouter models." >&2
      # Continue fetching anyway, as the endpoint is public
    fi
    # The fetch_openrouter_models function handles its own "Fetching..." message and errors.
    fetched_models=$(fetch_openrouter_models)
    if [ $? -ne 0 ] || [ -z "$fetched_models" ]; then
      # fetch_openrouter_models already prints detailed errors to stderr.
      # This message just confirms fallback to manual entry.
      echo "Warning: Using manual entry for OpenRouter models due to fetch error or no models found." >&2
      fetched_models="" # Ensure it's empty so the manual entry path is taken later
    fi
    ;;
  "OpenAI")
    # The fetch_openai_models function handles its own "Fetching..." message and errors.
    fetched_models=$(fetch_openai_models)
    if [ $? -ne 0 ] || [ -z "$fetched_models" ]; then
      # fetch_openai_models already prints detailed errors to stderr.
      echo "Warning: Using manual entry for OpenAI models due to fetch error or no models found." >&2
      fetched_models="" # Ensure it's empty so the manual entry path is taken later
    else
      echo "Successfully fetched OpenAI models." >&2
    fi
    ;;
  "Cerebras")
    # The fetch_cerebras_models function handles its own "Fetching..." message and errors.
    fetched_models=$(fetch_cerebras_models)
    if [ $? -ne 0 ] || [ -z "$fetched_models" ]; then
      # fetch_cerebras_models already prints detailed errors to stderr.
      echo "Warning: Using manual entry for Cerebras models due to fetch error or no models found." >&2
      fetched_models="" # Ensure it's empty so the manual entry path is taken later
    else
      echo "Successfully fetched Cerebras models." >&2
    fi
    ;;
"Ollama")
     # OLLAMA_BASE_URL from .env will be used by call_api and fetch_ollama_models.
     # Check if OLLAMA_BASE_URL is set, if not, fetching will use default.
     if [ -z "$OLLAMA_BASE_URL" ]; then
       echo "Info: OLLAMA_BASE_URL not set in $ENV_FILE. Will attempt to fetch models from default http://localhost:11434." >&2
     fi
     fetched_ollama_models=$(fetch_ollama_models)
     if [ $? -eq 0 ] && [ -n "$fetched_ollama_models" ]; then
       fetched_models="$fetched_ollama_models"
       echo "Successfully fetched Ollama models." >&2
     else
       echo "Warning: Could not fetch Ollama models. Using manual entry only." >&2
       fetched_models="" # Ensure it's empty on failure
     fi
     ;;
   *)
     # Check if it's a custom provider
     if [[ "$API_PROVIDER" =~ ^[A-Z0-9_]+$ ]]; then
       local provider_url_var="CUSTOM_PROVIDER_${API_PROVIDER}_URL"
       local provider_key_var="CUSTOM_PROVIDER_${API_PROVIDER}_KEY"
       if [ -n "${!provider_url_var}" ] && [ -n "${!provider_key_var}" ]; then
         fetched_custom_models=$(fetch_custom_models "$API_PROVIDER")
         if [ $? -eq 0 ] && [ -n "$fetched_custom_models" ]; then
           fetched_models="$fetched_custom_models"
           echo "Successfully fetched models from $API_PROVIDER." >&2
         else
           echo "Warning: Could not fetch models from $API_PROVIDER. Using manual entry only." >&2
           fetched_models="" # Ensure it's empty on failure
         fi
       else
         echo "Warning: Custom provider $API_PROVIDER not properly configured. Using manual entry only." >&2
         fetched_models=""
       fi
     else
       echo "Invalid provider selected."
       exit 1
     fi
     ;;
   esac

  # Prepare input for fzf: Manual option first, then fetched models
  fzf_input="$manual_entry_option"
  if [ -n "$fetched_models" ]; then
    # Sort models alphabetically for better usability
    fetched_models=$(echo "$fetched_models" | sort)
    fzf_input="$fzf_input\n$fetched_models"
  fi

  # Use fzf for selection
  MODEL=$(echo -e "$fzf_input" | fzf --prompt="Select Model for $API_PROVIDER: " --height=15 --layout=reverse)

  # Handle fzf cancellation
  if [ -z "$MODEL" ]; then
    echo "Configuration cancelled."
    return 1
  fi

  # Check if manual entry was selected
  if [[ "$MODEL" == "$manual_entry_option" ]]; then
    echo "Manual entry selected."
    MODEL=$(gum input --placeholder "Enter exact model name/ID for $API_PROVIDER...")
    # Handle cancellation of gum input
    if [ $? -ne 0 ] || [ -z "$MODEL" ]; then
      echo "Manual entry cancelled or empty. Configuration cancelled."
      return 1
    fi
  fi
  # MODEL variable now holds either the selected fetched model or the manually entered one

  # 3. Configure Editor
  local current_editor=${EDITOR:-vi} # Get current or default from sourced config or default
  local new_editor
  new_editor=$(gum input --header "Configure Editor" --placeholder "Enter preferred editor command (current: $current_editor)..." --value "$current_editor")
  # Update EDITOR only if gum input succeeded and value is not empty
  if [ $? -eq 0 ] && [ -n "$new_editor" ]; then
    EDITOR="$new_editor"
  fi

  # 4. Save Configuration
  echo "Saving configuration..."
  if [ -f "$CONFIG_FILE" ]; then
    # Update existing config file
    sed -i "s|^API_PROVIDER=.*|API_PROVIDER=\"$API_PROVIDER\"|" "$CONFIG_FILE"
    sed -i "s|^MODEL=.*|MODEL=\"$MODEL\"|" "$CONFIG_FILE"
    # Only update editor if it doesn't exist or is empty
    if ! grep -q "^EDITOR=" "$CONFIG_FILE" || [ -z "$(grep "^EDITOR=" "$CONFIG_FILE" | cut -d'=' -f2)" ]; then
      echo "EDITOR=\"$EDITOR\"" >>"$CONFIG_FILE"
    fi
  else
    # Create new config file
    echo "API_PROVIDER=\"$API_PROVIDER\"" >"$CONFIG_FILE"
    echo "MODEL=\"$MODEL\"" >>"$CONFIG_FILE"
    echo "EDITOR=\"$EDITOR\"" >>"$CONFIG_FILE"
  fi

  echo "Configuration saved to $CONFIG_FILE:"
  echo "  Provider: $API_PROVIDER"
  echo "  Model: $MODEL"
  echo "  Editor: $EDITOR"
  echo "Setup complete. You can now run '$0'." # Use $0 for the current script name
}

# --- Prompt History Functions ---
# Save a prompt to history
save_prompt_to_history() {
  local prompt="$1"
  if [ -n "$prompt" ]; then
    # Add to history file
    echo "$prompt" >> "$PROMPT_HISTORY_FILE"
    
    # Keep only the last MAX_PROMPT_HISTORY entries
    if [ -f "$PROMPT_HISTORY_FILE" ]; then
      tail -n "$MAX_PROMPT_HISTORY" "$PROMPT_HISTORY_FILE" > "$PROMPT_HISTORY_FILE.tmp"
      mv "$PROMPT_HISTORY_FILE.tmp" "$PROMPT_HISTORY_FILE"
    fi
  fi
}

# --- Provider Management Functions ---
# Add a new API provider and key
add_provider() {
  while true; do
    echo "Add Provider selected."
    
    # Select provider
    local providers="Gemini\nOpenRouter\nOpenAI\nCerebras\nOllama\nCustom OpenAI-Compatible"
    local selected_provider=$(echo -e "$providers" | fzf --prompt="Select Provider to Add: " --height=7 --layout=reverse)
    
    if [ -z "$selected_provider" ]; then
      echo "Provider selection cancelled."
      break
    fi
    
    if [[ "$selected_provider" == "Custom OpenAI-Compatible" ]]; then
      add_custom_provider
    else
      add_provider_key "$selected_provider"
    fi
    
    # Ask if user wants to add another provider
    echo ""
    local add_another
    add_another=$(gum choose --header "Add another provider?" "Yes" "No")
    
    if [ "$add_another" != "Yes" ]; then
      break
    fi
  done
  
  # Reload environment variables
  set -a
  if [ -f "$ENV_FILE" ]; then
    source "$ENV_FILE"
  fi
  set +a
  
  return 0
}

# Add API key or URL for a specific provider
add_provider_key() {
  local selected_provider="$1"
  local api_key=""
  local env_var_name=""
  
  # Handle Ollama differently (URL instead of API key)
  if [[ "$selected_provider" == "Ollama" ]]; then
    local ollama_url
    ollama_url=$(gum input --header "Ollama Configuration" --placeholder "Enter Ollama base URL (e.g., http://localhost:11434)...")
    if [ $? -eq 0 ] && [ -n "$ollama_url" ]; then
      # Update .env file with OLLAMA_BASE_URL
      if [ -f "$ENV_FILE" ]; then
        # Check if OLLAMA_BASE_URL already exists
        if grep -q "^OLLAMA_BASE_URL=" "$ENV_FILE"; then
          # Update existing OLLAMA_BASE_URL
          sed -i "s|^OLLAMA_BASE_URL=.*|OLLAMA_BASE_URL=$ollama_url|" "$ENV_FILE"
        else
          # Add new OLLAMA_BASE_URL
          echo "OLLAMA_BASE_URL=$ollama_url" >>"$ENV_FILE"
        fi
      else
        # Create new .env file
        echo "# Add your API keys here" >"$ENV_FILE"
        echo "GEMINI_API_KEY=" >>"$ENV_FILE"
        echo "OPENROUTER_API_KEY=" >>"$ENV_FILE"
        echo "OPENAI_API_KEY=" >>"$ENV_FILE"
        echo "CEREBRAS_API_KEY=" >>"$ENV_FILE"
        echo "OLLAMA_BASE_URL=$ollama_url" >>"$ENV_FILE"
        chmod 600 "$ENV_FILE"
      fi
      echo "Ollama URL saved to $ENV_FILE"
      return 0
    else
      echo "Ollama configuration cancelled."
      return 1
    fi
  else
    # Handle other providers (API keys)
    case "$selected_provider" in
      "Gemini") env_var_name="GEMINI_API_KEY" ;;
      "OpenRouter") env_var_name="OPENROUTER_API_KEY" ;;
      "OpenAI") env_var_name="OPENAI_API_KEY" ;;
      "Cerebras") env_var_name="CEREBRAS_API_KEY" ;;
    esac
    
    # Input API key securely
    api_key=$(gum input --header "$selected_provider API Key" --placeholder "Enter your $selected_provider API key..." --password)
    if [ $? -eq 0 ] && [ -n "$api_key" ]; then
      # Update .env file with API key
      if [ -f "$ENV_FILE" ]; then
        # Check if the API key variable already exists
        if grep -q "^$env_var_name=" "$ENV_FILE"; then
          # Update existing API key
          sed -i "s|^$env_var_name=.*|$env_var_name=$api_key|" "$ENV_FILE"
        else
          # Add new API key
          echo "$env_var_name=$api_key" >>"$ENV_FILE"
        fi
      else
        # Create new .env file
        echo "# Add your API keys here" >"$ENV_FILE"
        echo "GEMINI_API_KEY=" >>"$ENV_FILE"
        echo "OPENROUTER_API_KEY=" >>"$ENV_FILE"
        echo "OPENAI_API_KEY=" >>"$ENV_FILE"
        echo "CEREBRAS_API_KEY=" >>"$ENV_FILE"
        echo "OLLAMA_BASE_URL=" >>"$ENV_FILE"
        echo "$env_var_name=$api_key" >>"$ENV_FILE"
        chmod 600 "$ENV_FILE"
      fi
      echo "$selected_provider API key saved to $ENV_FILE"
      return 0
    else
      echo "API key input cancelled."
      return 1
    fi
  fi
}

# Add a custom OpenAI-compatible provider
add_custom_provider() {
  echo "Add Custom OpenAI-Compatible Provider selected."
  
  # Get provider name
  local provider_name
  provider_name=$(gum input --header "Provider Configuration" --placeholder "Enter provider name (e.g., ZAI, MYPROVIDER)...")
  if [ $? -ne 0 ] || [ -z "$provider_name" ]; then
    echo "Provider name input cancelled."
    return 1
  fi
  
  # Convert to uppercase for consistency
  provider_name=$(echo "$provider_name" | tr '[:lower:]' '[:upper:]')
  
  # Check if provider already exists
  local provider_url_var="CUSTOM_PROVIDER_${provider_name}_URL"
  local provider_key_var="CUSTOM_PROVIDER_${provider_name}_KEY"
  
  if [ -f "$ENV_FILE" ] && (grep -q "^$provider_url_var=" "$ENV_FILE" || grep -q "^$provider_key_var=" "$ENV_FILE"); then
    echo "Provider $provider_name already exists. Use /config to change model or update manually."
    return 1
  fi
  
  # Get base URL
  local base_url
  base_url=$(gum input --header "Provider Configuration" --placeholder "Enter base URL (e.g., https://api.openai.com/v1)...")
  if [ $? -ne 0 ] || [ -z "$base_url" ]; then
    echo "Base URL input cancelled."
    return 1
  fi
  
  # Remove trailing slash if present
  base_url=$(echo "$base_url" | sed 's:/*$::')
  
  # Get API key
  local api_key
  api_key=$(gum input --header "Provider Configuration" --placeholder "Enter API key..." --password)
  if [ $? -ne 0 ] || [ -z "$api_key" ]; then
    echo "API key input cancelled."
    return 1
  fi
  
  # Update .env file
  if [ -f "$ENV_FILE" ]; then
    echo "CUSTOM_PROVIDER_${provider_name}_URL=$base_url" >>"$ENV_FILE"
    echo "CUSTOM_PROVIDER_${provider_name}_KEY=$api_key" >>"$ENV_FILE"
  else
    echo "# Add your API keys here" >"$ENV_FILE"
    echo "GEMINI_API_KEY=" >>"$ENV_FILE"
    echo "OPENROUTER_API_KEY=" >>"$ENV_FILE"
    echo "OPENAI_API_KEY=" >>"$ENV_FILE"
    echo "CEREBRAS_API_KEY=" >>"$ENV_FILE"
    echo "OLLAMA_BASE_URL=" >>"$ENV_FILE"
    echo "CUSTOM_PROVIDER_${provider_name}_URL=$base_url" >>"$ENV_FILE"
    echo "CUSTOM_PROVIDER_${provider_name}_KEY=$api_key" >>"$ENV_FILE"
    chmod 600 "$ENV_FILE"
  fi
  
  echo "Custom provider $provider_name saved to $ENV_FILE"
  return 0
}

# Fetch models from a custom OpenAI-compatible provider
fetch_custom_models() {
  local provider_name="$1"
  local provider_url_var="CUSTOM_PROVIDER_${provider_name}_URL"
  local provider_key_var="CUSTOM_PROVIDER_${provider_name}_KEY"
  
  local base_url="${!provider_url_var}"
  local api_key="${!provider_key_var}"
  
  if [ -z "$base_url" ] || [ -z "$api_key" ]; then
    echo "Error: Provider $provider_name not properly configured." >&2
    return 1
  fi
  
  local models_api_url="${base_url}/models"
  local models_json
  
  echo "Fetching models from $provider_name at $base_url..." >&2
  models_json=$(curl -s -H "Authorization: Bearer $api_key" "$models_api_url")
  
  if [ $? -ne 0 ] || [ -z "$models_json" ]; then
    echo "Error: Failed to fetch models from $provider_name." >&2
    return 1
  fi
  
  # Check for API errors
  local custom_error=$(echo "$models_json" | jq -r '.error.message // empty')
  if [ -n "$custom_error" ]; then
    echo "Warning: $provider_name API Error when fetching models: $custom_error" >&2
    return 1
  fi
  
  # Check if data array exists
  if ! echo "$models_json" | jq -e '.data | type == "array"' >/dev/null 2>&1; then
    echo "Warning: Invalid JSON or missing 'data' array in $provider_name models response." >&2
    return 1
  fi
  
  # Extract model IDs, filter out duplicates, sort
  local extracted_models
  extracted_models=$(echo "$models_json" | jq -r '.data[].id // empty' | sort -u | grep .)
  
  if [ -z "$extracted_models" ]; then
    echo "Warning: No models found in $provider_name response or failed to parse names." >&2
    return 1
  fi
  
  echo "$extracted_models"
  return 0
}

# Get previous prompt by index (1 = most recent)
get_previous_prompt() {
  local index=$1
  if [ -f "$PROMPT_HISTORY_FILE" ] && [ "$index" -gt 0 ]; then
    tail -n "$index" "$PROMPT_HISTORY_FILE" | head -n 1
  fi
}

# Display prompt history and let user select
select_from_prompt_history() {
  if [ ! -f "$PROMPT_HISTORY_FILE" ] || [ ! -s "$PROMPT_HISTORY_FILE" ]; then
    echo "No prompt history available."
    return 1
  fi
  
  # Create a numbered list of recent prompts
  local history_items=()
  local count=0
  while IFS= read -r line; do
    if [ -n "$line" ]; then
      history_items+=("$line")
      count=$((count + 1))
    fi
  done < <(tail -n "$MAX_PROMPT_HISTORY" "$PROMPT_HISTORY_FILE" | tac)
  
  if [ ${#history_items[@]} -eq 0 ]; then
    echo "No prompt history available."
    return 1
  fi
  
  # Use gum choose to let user select from history
  local selected
  selected=$(printf '%s\n' "${history_items[@]}" | gum choose --header "Select a previous prompt (or press Esc to cancel):")
  local gum_status=$?
  
  if [ $gum_status -eq 0 ] && [ -n "$selected" ]; then
    echo "$selected"
    return 0
  else
    return 1
  fi
}

# Display help information
show_help() {
  echo "OliveOwl Commands:"
  echo "  /exit, /quit, /q    - Exit the application"
  echo "  /history, /h        - Load previous chat sessions"
  echo "  /config, /c         - Complete configuration (API + editor + other settings)"
  echo "  /view, /v           - View current chat in editor"
  echo "  /new, /n            - Start a new chat session"
  echo "  /prompt-history, /ph, /p - Select from previous prompts"
  echo "  /help, /?           - Show this help message"
  echo ""
  echo "All commands must be the only text in your input to be recognized."
}

# --- System Prompt ---
# Instructs the AI on formatting and behavior
load_system_prompt() {
  local base_prompt=$(cat <<'EOF'
You are a helpful AI assistant. Answer user questions clearly and concisely.

For regular questions, respond in natural language without using code blocks.

ONLY use Markdown code blocks when providing code, commands, scripts, or any content that users might need to copy and paste. This includes:
- Shell commands (bash, zsh, fish)
- Programming code (Python, JavaScript, etc.)
- Configuration files
- Docker commands
- Git commands
- Ollama commands
- Game commands (Minecraft, etc.)
- HTML/CSS/SQL or other markup/query languages

When using code blocks, format them properly so they can be copied and used directly:
```bash
echo "example command"
```

For all other responses, use plain text without code blocks.
EOF
)

  local instructions_file="$CONFIG_DIR/extra_instructions.txt"
  if [ -f "$instructions_file" ] && [ -s "$instructions_file" ]; then
    local extra_instructions=$(cat "$instructions_file")
    echo -e "$base_prompt\n\nAdditional Instructions:\n$extra_instructions"
  else
    echo "$base_prompt"
  fi
}

SYSTEM_PROMPT=$(load_system_prompt)

# --- History Management ---
CHAT_HISTORY=() # In-memory array of JSON objects for the current session
CURRENT_HISTORY_FILE=""
CURRENT_SESSION_NAME=""
USER_PROMPT_COUNT=0
IS_MANUALLY_NAMED=false

# Function to create the API-specific JSON object string for a message
# Usage: create_api_message_json "role" "content"
# Returns a compact JSON string suitable for the configured API provider
create_api_message_json() {
  local role="$1"
  local content="$2"
  # Content is passed directly, assuming it doesn't need pre-escaping for jq --arg
  # json_content=$(echo "$content" | jq -Rsa .) # REMOVED

  if [[ "$API_PROVIDER" == "Gemini" ]]; then
    # Gemini uses "user" and "model" roles, and a "parts" array
    local gemini_role="$role"
    # Map OpenRouter's 'assistant' role if needed (though we primarily use 'model' for Gemini AI responses)
    if [[ "$role" == "assistant" ]]; then
      gemini_role="model"
    fi
    # Simple text part for now - Pass content directly to --arg text_content
    # jq needs the value for "text" to be a valid JSON string. Let jq handle the encoding.
    jq -n --arg role "$gemini_role" --arg text_content "$content" \
      '{role: $role, parts: [{"text": $text_content}]}' | jq -c .
  else # OpenRouter uses "user" and "assistant"
    local openrouter_role="$role"
    # Map Gemini's 'model' role if needed (though we primarily use 'assistant' for OpenRouter AI responses)
    if [[ "$role" == "model" ]]; then
      openrouter_role="assistant"
    fi
    # Pass content directly to --arg content
    jq -n --arg role "$openrouter_role" --arg content "$content" \
      '{role: $role, content: $content}' | jq -c .
  fi
}

# Function to load history from a file into the CHAT_HISTORY array
# Usage: load_history "path/to/history.json"
load_history() {
  local file_path="$1"
  if [ ! -f "$file_path" ]; then
    echo "History file not found: $file_path" >&2
    return 1
  fi

  # Check if the file is valid JSON and contains the expected structure
  if ! jq -e '(.session_name != null) and (.messages | type == "array")' "$file_path" >/dev/null 2>&1; then
    echo "Error: History file is not in the new format or is corrupt: $file_path" >&2
    # TODO: Offer to migrate old history file format? For now, we fail.
    return 1
  fi

  # Load data from the new structure
  CURRENT_SESSION_NAME=$(jq -r '.session_name' "$file_path")
  mapfile -t CHAT_HISTORY < <(jq -c '.messages[]' "$file_path")
  if [ $? -ne 0 ]; then
    echo "Error: Failed to parse messages from history file: $file_path" >&2
    CHAT_HISTORY=() # Clear history on parse error
    return 1
  fi

  # Set session properties
  CURRENT_HISTORY_FILE="$file_path"
  # Heuristic: if filename doesn't start with chat_, it was manually named.
  if [[ ! "$(basename "$file_path")" =~ ^chat_ ]]; then
    IS_MANUALLY_NAMED=true
  else
    IS_MANUALLY_NAMED=false
  fi

  # Recalculate user prompt count from loaded history
  USER_PROMPT_COUNT=0
  for item_json in "${CHAT_HISTORY[@]}"; do
    if [[ $(echo "$item_json" | jq -r '.message_json.role') == "user" ]]; then
      ((USER_PROMPT_COUNT++))
    fi
  done

  echo "Loaded Session: $CURRENT_SESSION_NAME"
  echo "History file: $CURRENT_HISTORY_FILE"
}

# Function to display the current CHAT_HISTORY
display_history() {
  echo "--- Chat History for: $CURRENT_SESSION_NAME ---"
  # Iterate through the CHAT_HISTORY array (each element is a JSON string object)
  for history_item_json in "${CHAT_HISTORY[@]}"; do
    # Parse the outer object to get message_json and model_used
    local message_json model_used role content display_model_name
    message_json=$(echo "$history_item_json" | jq -r '.message_json // empty')
    model_used=$(echo "$history_item_json" | jq -r '.model_used // empty') # Will be empty for user messages

    if [ -z "$message_json" ]; then
      echo "Warning: Skipping invalid history item." >&2
      continue
    fi

    # Parse the inner message_json to get role and content
    role=$(echo "$message_json" | jq -r '.role')
    # Extract the raw content string
    content=$(echo "$message_json" | jq -r 'if .parts then .parts[0].text else .content end')

    if [[ "$role" == "user" ]]; then
      # Display directly, assuming jq -r decoded correctly
      echo -e "\n\e[1m\e[34mYou:\e[0m" # Print prompt (Bold Blue)
      echo "$content"
    elif [[ "$role" == "model" || "$role" == "assistant" ]]; then
      # Use the specific model_used if available, otherwise fall back to the current $MODEL
      display_model_name="${model_used:-$MODEL}"
      echo -e "\n\e[1m\e[35mAI ($display_model_name):\e[0m"                   # Bold Purple for AI, showing specific model, with explicit reset
      echo -e "\e[1m\e[35m=============================================\e[0m" # Bold Purple for AI, showing specific model, with explicit reset
      # Pipe the extracted content directly to bat, assuming jq -r decoded correctly
      echo "$content" | bat --language md --paging=never --style=plain --color=always
    # Handle potential 'system' role if it ever gets stored/displayed (currently shouldn't)
    # elif [[ "$role" == "system" ]]; then
    #   echo -e "\n\e[35mSystem:\e[0m $content" # Magenta for system
    fi
  done
  echo "--------------------"
}

# Function to update cache for a single file (targeted update)
# Usage: update_cache_for_file "/path/to/file.json" "Session Name"
update_cache_for_file() {
   local file_path="$1"
   local session_name="$2"
   local cache_file="$HOME/.oliveowl_history_cache"
   
   if [ ! -f "$file_path" ] || [ -z "$session_name" ]; then
     return 1
   fi
   
   # Create cache directory if it doesn't exist
   mkdir -p "$(dirname "$cache_file")"
   
   # Build display name
   local display_name="$session_name"
   local basename=$(basename "$file_path" .json)
   if [[ $basename =~ chat_([0-9]{8}_[0-9]{6}) ]]; then
     local timestamp="${BASH_REMATCH[1]}"
     display_name="$session_name ($timestamp)"
   else
     display_name="$session_name ($basename)"
   fi
   
   # Add/update entry in cache file - NEW CHAT ON TOP
   local temp_cache=$(mktemp)
   
   # Add new/updated entry FIRST (to appear on top)
   echo -e "$display_name\t$file_path" > "$temp_cache"
   
   # Keep existing entries (except this file if it exists)
   if [ -f "$cache_file" ]; then
     while IFS=$'\t' read -r cached_display_name cached_file_path; do
       # Skip old entry for this file
       if [ "$cached_file_path" != "$file_path" ]; then
         echo -e "$cached_display_name\t$cached_file_path" >> "$temp_cache"
       fi
     done < "$cache_file"
   fi
   
   # Replace cache file
   mv "$temp_cache" "$cache_file"
 }

# Function to select and load a history file using fzf
# Returns 0 on success, 1 on failure/cancel
# Build or update the history cache file in background
build_history_cache() {
   local cache_file="$HOME/.oliveowl_history_cache"
   
   # Get cache file timestamp if it exists
   local cache_time=0
   if [ -f "$cache_file" ]; then
     cache_time=$(stat -c %Y "$cache_file" 2>/dev/null || echo "0")
   fi
   
   # Find newest history file timestamp
   local newest_file_time=0
   newest_file_time=$(find "$HISTORY_DIR" -maxdepth 1 -name "*.json" -exec stat -c %Y {} \; 2>/dev/null | sort -n | tail -1 || echo "0")
   
   # If cache is newer than all files AND cache exists AND cache is not empty, skip rebuild (cache is fresh)
   # Always rebuild if cache doesn't exist, cache is empty, or no history files exist
   if [ "$cache_time" -gt "$newest_file_time" ] && [ "$newest_file_time" -gt 0 ] && [ -f "$cache_file" ] && [ -s "$cache_file" ]; then
     return 0  # Cache is fresh, no rebuild needed
   fi
  
  # Find all .json files sorted by modification time (newest first)
  local -a history_files
  mapfile -t history_files < <(find "$HISTORY_DIR" -maxdepth 1 -name "*.json" -print0 2>/dev/null | xargs -0 ls -t 2>/dev/null)

  if [ ${#history_files[@]} -eq 0 ]; then
    # Create empty cache file if no history exists
    > "$cache_file" 2>/dev/null
    return 0
  fi

  # Build cache incrementally: read existing cache and update only changed files
  local cache_entries=""
  local processed_files=()  # Track which files we've processed
  
  # If cache exists, read it and keep valid entries
  if [ -f "$cache_file" ] && [ -s "$cache_file" ]; then
    while IFS=$'\t' read -r cached_display_name cached_file_path; do
      # Skip if file no longer exists or isn't readable
      if [ ! -f "$cached_file_path" ] || [ ! -r "$cached_file_path" ]; then
        continue
      fi
      
      # Check if file has been modified since last cache build
      local file_time=$(stat -c %Y "$cached_file_path" 2>/dev/null || echo "0")
      if [ "$file_time" -le "$cache_time" ]; then
        # File hasn't changed, keep cached entry
        cache_entries+="$cached_display_name\t$cached_file_path"$'\n'
        processed_files+=("$cached_file_path")
      fi
    done < "$cache_file"
  fi
  
  # Process new or modified files
  for file_path in "${history_files[@]}"; do
    # Skip if already processed (from cache) and unchanged
    local already_processed=false
    for processed_path in "${processed_files[@]}"; do
      if [ "$file_path" = "$processed_path" ]; then
        already_processed=true
        break
      fi
    done
    if $already_processed; then
      continue
    fi
    
    # Skip if file is not readable
    if [ ! -r "$file_path" ]; then
      continue
    fi
    
    # Extract session name using fast text processing
    local session_name
    if ! session_name=$(sed -n '2s/.*"session_name"[[:space:]]*:[[:space:]]*"\([^"]*\)".*/\1/p' "$file_path" 2>/dev/null); then
      continue
    fi
    
    # Skip if we didn't find a session name
    if [ -z "$session_name" ]; then
      continue
    fi
    
    # Create display name with timestamp for uniqueness
    local display_name="$session_name"
    local basename=$(basename "$file_path" .json)
    if [[ $basename =~ chat_([0-9]{8}_[0-9]{6}) ]]; then
      local timestamp="${BASH_REMATCH[1]}"
      display_name="$session_name ($timestamp)"
    else
      display_name="$session_name ($basename)"
    fi
    
    # Add to cache entries
    cache_entries+="$display_name\t$file_path"$'\n'
  done

  # Write the updated cache
  printf "%s" "$cache_entries" > "$cache_file" 2>/dev/null
}

select_and_load_history() {
   local cache_file="$HOME/.oliveowl_history_cache"
   
   # Check if there are any history files first
   local history_files_count
   history_files_count=$(find "$HISTORY_DIR" -maxdepth 1 -name "*.json" 2>/dev/null | wc -l)
   
   if [ "$history_files_count" -eq 0 ]; then
     echo "No history files found in $HISTORY_DIR."
     return 1 # Indicate failure
   fi
   
   # Rebuild cache only if it doesn't exist or is empty
   # Files are added to cache individually on save, so we only need full rebuild for existing files
   if [ ! -f "$cache_file" ] || [ ! -s "$cache_file" ]; then
     build_history_cache >/dev/null 2>&1
   fi
   
   # Read the cache file
   local fzf_input
   if [ -f "$cache_file" ] && [ -s "$cache_file" ]; then
     fzf_input=$(cat "$cache_file" 2>/dev/null)
   else
     echo "No history files found in $HISTORY_DIR."
     return 1 # Indicate failure
   fi

  if [ -z "$fzf_input" ]; then
    echo "No valid history files found to load."
    return 1
  fi

# Use fzf to select a session.
   local selected_line
   selected_line=$(echo -e "$fzf_input" |
     fzf --prompt="Select Chat Session to Load: " \
         --height=100% --layout=reverse \
         --delimiter='\t' --with-nth=1 \
         --preview 'echo {} | cut -f2 | xargs -r bat --style=numbers --color=always --language=json --wrap=character' \
         --preview-window=right:60%:wrap)

  # Check fzf exit status and if a selection was made
  if [ $? -eq 0 ] && [ -n "$selected_line" ]; then
    local chosen_file_path
    chosen_file_path=$(echo "$selected_line" | cut -f2) # Cut the second field (the path)

    if [ -n "$chosen_file_path" ] && [ -f "$chosen_file_path" ]; then
      clear # Clear the terminal screen
      load_history "$chosen_file_path"
      display_history
      return 0 # Success
    else
      echo "Error: Could not find the file for the selected session."
      return 1 # Failure
    fi
  else
    echo "History loading cancelled."
    return 1 # Cancelled
  fi
}

# Function to save the current CHAT_HISTORY array to CURRENT_HISTORY_FILE
save_history() {
  if [ -z "$CURRENT_HISTORY_FILE" ]; then
    echo "Error: No history file set for saving." >&2
    return 1
  fi

  # Check if this is the first save (file doesn't exist yet)
  local is_first_save=false
  if [ ! -f "$CURRENT_HISTORY_FILE" ]; then
    is_first_save=true
  fi

  # Create a JSON array from the CHAT_HISTORY bash array
  local messages_json_array
  messages_json_array=$(printf "%s\n" "${CHAT_HISTORY[@]}" | jq -s '.')
  if [ $? -ne 0 ]; then
    echo "Error: Failed to serialize messages for saving." >&2
    return 1
  fi

  # Create the final JSON object and write to the history file
  # This will create the file if it doesn't exist (lazy creation)
  if jq -n \
    --arg name "$CURRENT_SESSION_NAME" \
    --argjson messages "$messages_json_array" \
    '{session_name: $name, messages: $messages}' >"$CURRENT_HISTORY_FILE"; then
    # Show file creation message only on first save
    if [ "$is_first_save" = true ]; then
      echo "Session file created: $CURRENT_HISTORY_FILE" >&2
    fi
    # Update cache with new file immediately (targeted update)
    update_cache_for_file "$CURRENT_HISTORY_FILE" "$CURRENT_SESSION_NAME"
  else
    echo "Error: Failed to save history to $CURRENT_HISTORY_FILE" >&2
    return 1 # Indicate save failure
  fi
}

# Function to format and view history in editor
view_history() {
  if [ ${#CHAT_HISTORY[@]} -eq 0 ]; then
    echo "Chat history is empty. Nothing to view."
    return
  fi

  local temp_history_file
  temp_history_file=$(mktemp --suffix=.md) # Create temp markdown file
  # Ensure temp file is deleted on exit/interrupt
  trap 'rm -f "$temp_history_file" 2>/dev/null' EXIT INT TERM HUP

  echo "# Chat History ($(date))" >"$temp_history_file"
  echo "" >>"$temp_history_file"

  # Iterate through history and format as Markdown
  for history_item_json in "${CHAT_HISTORY[@]}"; do
    local message_json model_used role content display_model_name
    message_json=$(echo "$history_item_json" | jq -r '.message_json // empty')
    model_used=$(echo "$history_item_json" | jq -r '.model_used // empty')

    if [ -z "$message_json" ]; then continue; fi # Skip invalid items

    role=$(echo "$message_json" | jq -r '.role')
    content=$(echo "$message_json" | jq -r 'if .parts then .parts[0].text else .content end')

    if [[ "$role" == "user" ]]; then
      echo "### You:" >>"$temp_history_file"
      echo "" >>"$temp_history_file"
      echo "$content" >>"$temp_history_file"
    elif [[ "$role" == "model" || "$role" == "assistant" ]]; then
      display_model_name="${model_used:-$MODEL}"
      echo "### AI ($display_model_name):" >>"$temp_history_file"
      echo "" >>"$temp_history_file"
      echo "$content" >>"$temp_history_file" # Raw content, editor handles wrapping
    fi
    echo "" >>"$temp_history_file"
    echo "---" >>"$temp_history_file" # Separator
    echo "" >>"$temp_history_file"
  done

  # Open the temp file in the configured editor
  "$EDITOR" "$temp_history_file"

  # Remove trap and delete file explicitly after editor closes
  trap - EXIT INT TERM HUP
  rm -f "$temp_history_file" 2>/dev/null
  echo "Resuming chat..." # Indicate return to chat
}

# Function to prepare a new session (but don't create file yet)
# Takes an optional session name as argument
create_new_session_file() {
  local session_name="$1" # Optional name passed from start_new_session
  local filename timestamp sanitized_name

  # Reset session state variables
  CHAT_HISTORY=()
  USER_PROMPT_COUNT=0
  IS_MANUALLY_NAMED=false

  timestamp=$(date +"%Y%m%d_%H%M%S")

  if [ -z "$session_name" ]; then
    # Nameless session, will be auto-named by AI
    IS_MANUALLY_NAMED=false
    CURRENT_SESSION_NAME="New Chat"
    filename="chat_${timestamp}.json"
    echo "Starting new chat session (will be named automatically)..."
  else
    # User-provided name
    IS_MANUALLY_NAMED=true
    CURRENT_SESSION_NAME="$session_name"
    # Sanitize the name for the filename
    sanitized_name=$(echo "$session_name" | sed -e 's/ /_/g' -e 's/[^a-zA-Z0-9_-]//g')
    if [ -z "$sanitized_name" ]; then # Handle case where sanitization removes everything
      filename="chat_${timestamp}.json"
      echo "Invalid name provided, using timestamp for filename..."
    else
      # Append timestamp to user-provided name to avoid collisions easily
      filename="${sanitized_name}_${timestamp}.json"
    fi
    echo "Starting new named session: $CURRENT_SESSION_NAME"
  fi

  CURRENT_HISTORY_FILE="$HISTORY_DIR/$filename"
  # NOTE: File is NOT created here anymore - it will be created on first save
  echo "Session prepared: $CURRENT_SESSION_NAME"
  
  # Add new session to cache immediately for /h visibility
  update_cache_for_file "$CURRENT_HISTORY_FILE" "$CURRENT_SESSION_NAME"
}

# Function to handle the initial session prompt (new, load history, or exit)
# Returns 0 if a session was started (new or loaded), 1 if history loading failed/cancelled
start_new_session() {
  local user_choice

  # Prompt for session name or command
  # Prompt for session name or command using gum input
  user_choice=$(gum input --placeholder "Enter name for new chat session (or /history, /h, /config, /c, /exit, /quit, /q)")

  # Check if gum input was cancelled (e.g., Ctrl+C or Esc)
  if [ $? -ne 0 ]; then
    echo "Input cancelled. Exiting."
    exit 1
  fi

  case "$user_choice" in
  "/exit"|"/quit"|"/q" | "exit")
    gum style --foreground "#FF0000" --bold "Exiting OliveOwl."
    exit 0
    ;;
  "/history"|"/h")
    select_and_load_history # This function handles messages and returns status
    return $?               # Return status (0 for success, 1 for fail/cancel)
    ;;
  "/config"|"/c")
    echo "Switching to config..."
    configure_settings # Re-run config
    # Re-source config in case it changed
    if [ -f "$CONFIG_FILE" ]; then
      source "$CONFIG_FILE"
    fi
    echo "Config updated. Provider: $API_PROVIDER, Model: $MODEL"
    # After config, prompt again for session action
    start_new_session # Recursive call
    return $?         # Return status from the recursive call
    ;;
  *)
    # Treat anything else (name or blank) as a request for a new session
    create_new_session_file "$user_choice"
    return 0 # Indicate success (new session created)
    ;;
  esac
}

# --- API Call Implementation ---
# Takes user input, sends history + input to API, returns AI response text
call_api() {
  local user_input="$1"
  local api_url api_key payload response response_text error_message

  # Reload the system prompt to include any updated instructions
  SYSTEM_PROMPT=$(load_system_prompt)

  # 1. Prepare payload: Convert history messages to the format required by the CURRENT API provider
  local converted_messages=()
  for history_item_json in "${CHAT_HISTORY[@]}"; do
    # Extract the raw message_json string from the wrapper object string
    local original_msg_json=$(echo "$history_item_json" | jq -r '.message_json // empty')
    if [ -z "$original_msg_json" ]; then
      echo "Warning: Skipping empty message_json in history item." >&2
      continue
    fi

    # Parse the original message to get role and content, regardless of format
    local original_role original_content
    original_role=$(echo "$original_msg_json" | jq -r '.role')
    # Try extracting from Gemini format first, then OpenRouter format
    original_content=$(echo "$original_msg_json" | jq -r 'if .parts then .parts[0].text else .content end // empty')

    if [ -z "$original_role" ] || [ -z "$original_content" ]; then
      echo "Warning: Skipping history item with missing role or content: $original_msg_json" >&2
      continue
    fi

    # Convert the extracted role/content to the format needed by the CURRENT provider
    local converted_msg_json=$(create_api_message_json "$original_role" "$original_content")
    if [ -n "$converted_msg_json" ]; then
      converted_messages+=("$converted_msg_json")
    else
      echo "Warning: Failed to convert history message: $original_msg_json" >&2
    fi
  done
  # Add the current user's prompt to the list of messages to be sent
  local user_api_json
  user_api_json=$(create_api_message_json "user" "$user_input")
  if [ -n "$user_api_json" ]; then
    converted_messages+=("$user_api_json")
  fi

  # Combine the *converted* message JSON strings into a single JSON array string for the API
  local history_json_array=$(printf "%s\n" "${converted_messages[@]}" | jq -s '.')

  # 2. Set API specifics based on provider
  if [[ "$API_PROVIDER" == "Gemini" ]]; then
    api_key="$GEMINI_API_KEY"
    # Note: Gemini API URL structure might vary slightly based on region or specific model version. Adjust if needed.
    # Using v1beta for generative models as it's common.
    # The MODEL variable fetched from the API already contains the "models/" prefix.
    api_url="https://generativelanguage.googleapis.com/v1beta/${MODEL}:generateContent?key=${api_key}"

    # Gemini payload structure: { "contents": [history array], "systemInstruction": { "parts": [{"text": "..."}] } }
    # Construct the payload JSON using jq for proper escaping
    payload=$(jq -n --argjson history "$history_json_array" --arg system_prompt "$SYSTEM_PROMPT" \
      '{contents: $history, systemInstruction: {parts: [{"text": $system_prompt}]}}')

    # 3. Make the curl request for Gemini
    response=$(curl -s -X POST "$api_url" \
      -H "Content-Type: application/json" \
      -d "$payload")

    # 4. Parse Gemini response and handle errors
    if [ $? -ne 0 ]; then
      echo "Error: curl command failed for Gemini." >&2
      return 1
    fi

    # Check for API errors in the response JSON
    error_message=$(echo "$response" | jq -r '.error.message // empty')
    if [ -n "$error_message" ]; then
      echo "Error: Gemini API Error: $error_message" >&2
      # You might want to see the full error: echo "$response" >&2
      return 1
    fi

    # Extract the text content. Gemini nests it under candidates -> content -> parts -> text
    # It might return multiple candidates, usually the first is sufficient.
    # It might also return multiple parts, concatenate them.
    response_text=$(echo "$response" | jq -r '.candidates[0].content.parts[]?.text // empty' | paste -sd '\n')

    if [ -z "$response_text" ]; then
      # Handle cases like safety blocks or empty responses
      local finish_reason=$(echo "$response" | jq -r '.candidates[0].finishReason // "UNKNOWN"')
      if [[ "$finish_reason" == "SAFETY" ]]; then
        echo "Warning: Gemini response blocked due to safety settings." >&2
        response_text="[Response blocked by safety settings]"
      elif [[ "$finish_reason" == "RECITATION" ]]; then
        echo "Warning: Gemini response blocked due to recitation policy." >&2
        response_text="[Response blocked by recitation policy]"
      else
        echo "Error: Could not extract text from Gemini response. Finish Reason: $finish_reason" >&2
        echo "Full Gemini Response: $response" >&2 # Temporarily uncommented for debugging
        return 1
      fi
    fi

  elif [[ "$API_PROVIDER" == "OpenRouter" ]]; then
    api_key="$OPENROUTER_API_KEY"
    api_url="https://openrouter.ai/api/v1/chat/completions"

    # OpenRouter payload structure: { "model": "...", "messages": [history array including system prompt] }
    # Create the system message JSON separately
    local system_message_json=$(create_api_message_json "system" "$SYSTEM_PROMPT")
    # Prepend the system message JSON string to the extracted history array string
    # Note: history_json_array already contains the user/assistant messages
    local messages_json_array_with_system=$(echo "$history_json_array" | jq --argjson sys_msg "$system_message_json" '[$sys_msg] + .')

    payload=$(jq -n --arg model "$MODEL" --argjson messages "$messages_json_array_with_system" \
      '{model: $model, messages: $messages}')
    # Add other parameters like temperature, max_tokens if desired:
    # '{model: $model, messages: $messages, temperature: 0.7, max_tokens: 1024}'

    # 3. Make the curl request for OpenRouter
    response=$(curl -s -X POST "$api_url" \
      -H "Content-Type: application/json" \
      -H "Authorization: Bearer $api_key" \
      -H "HTTP-Referer: http://localhost" \
      -H "X-Title: OliveOwl CLI" \
      -d "$payload")

    # 4. Parse OpenRouter response and handle errors
    if [ $? -ne 0 ]; then
      echo "Error: curl command failed for OpenRouter." >&2
      return 1
    fi

    # Check for API errors
    error_message=$(echo "$response" | jq -r '.error.message // empty')
    if [ -n "$error_message" ]; then
      echo "Error: OpenRouter API Error: $error_message" >&2
      # echo "Full OpenRouter Response: $response" >&2 # for debugging
      return 1
    fi

    # Extract the text content. OpenRouter nests it under choices -> message -> content
    response_text=$(echo "$response" | jq -r '.choices[0].message.content // empty')

    if [ -z "$response_text" ]; then
      echo "Error: Could not extract text from OpenRouter response." >&2
      # echo "Full OpenRouter Response: $response" >&2 # for debugging
      return 1
    fi

  elif [[ "$API_PROVIDER" == "OpenAI" ]]; then
    api_key="$OPENAI_API_KEY"
    api_url="https://api.openai.com/v1/chat/completions"

    # OpenAI payload is like OpenRouter: { "model": "...", "messages": [...] }
    local system_message_json=$(create_api_message_json "system" "$SYSTEM_PROMPT")
    local messages_json_array_with_system=$(echo "$history_json_array" | jq --argjson sys_msg "$system_message_json" '[$sys_msg] + .')

    payload=$(jq -n --arg model "$MODEL" --argjson messages "$messages_json_array_with_system" \
      '{model: $model, messages: $messages}')

    # Make the curl request for OpenAI
    response=$(curl -s -X POST "$api_url" \
      -H "Content-Type: application/json" \
      -H "Authorization: Bearer $api_key" \
      -d "$payload")

    # Parse OpenAI response and handle errors
    if [ $? -ne 0 ]; then
      echo "Error: curl command failed for OpenAI." >&2
      return 1
    fi

    # Check for API errors
    error_message=$(echo "$response" | jq -r '.error.message // empty')
    if [ -n "$error_message" ]; then
      echo "Error: OpenAI API Error: $error_message" >&2
      # echo "Full OpenAI Response: $response" >&2 # for debugging
      return 1
    fi

    # Extract the text content. OpenAI nests it under choices -> message -> content
    response_text=$(echo "$response" | jq -r '.choices[0].message.content // empty')

    if [ -z "$response_text" ]; then
      echo "Error: Could not extract text from OpenAI response." >&2
      # echo "Full OpenAI Response: $response" >&2 # for debugging
      return 1
    fi

  elif [[ "$API_PROVIDER" == "Cerebras" ]]; then
    api_key="$CEREBRAS_API_KEY"
    api_url="https://api.cerebras.ai/v1/chat/completions"

    # Cerebras payload is like OpenRouter: { "model": "...", "messages": [...] }
    local system_message_json=$(create_api_message_json "system" "$SYSTEM_PROMPT")
    local messages_json_array_with_system=$(echo "$history_json_array" | jq --argjson sys_msg "$system_message_json" '[$sys_msg] + .')

    payload=$(jq -n --arg model "$MODEL" --argjson messages "$messages_json_array_with_system" \
      '{model: $model, messages: $messages, stream: false}')

    # Make the curl request for Cerebras
    response=$(curl -s -X POST "$api_url" \
      -H "Content-Type: application/json" \
      -H "Authorization: Bearer $api_key" \
      -d "$payload")

    # Parse Cerebras response and handle errors
    if [ $? -ne 0 ]; then
      echo "Error: curl command failed for Cerebras." >&2
      return 1
    fi

    # Check for API errors
    error_message=$(echo "$response" | jq -r '.error.message // empty')
    if [ -n "$error_message" ]; then
      echo "Error: Cerebras API Error: $error_message" >&2
      # echo "Full Cerebras Response: $response" >&2 # for debugging
      return 1
    fi

    # Extract the text content. Cerebras nests it under choices -> message -> content
    response_text=$(echo "$response" | jq -r '.choices[0].message.content // empty')

    if [ -z "$response_text" ]; then
      echo "Error: Could not extract text from Cerebras response." >&2
      # echo "Full Cerebras Response: $response" >&2 # for debugging
      return 1
    fi

  elif [[ "$API_PROVIDER" == "Ollama" ]]; then
    # Use OLLAMA_BASE_URL from .env, default if not set or empty
    local ollama_url_base="${OLLAMA_BASE_URL:-http://localhost:11434}"
    # Remove trailing slash if present
    ollama_url_base=$(echo "$ollama_url_base" | sed 's:/*$::')
    api_url="${ollama_url_base}/api/chat"

    # Ollama payload structure: { "model": "...", "messages": [...], "stream": false }
    # Messages array should include system prompt, then history.
    # create_api_message_json for "system" role:
    local system_message_json=$(create_api_message_json "system" "$SYSTEM_PROMPT")

    # history_json_array is already prepared with user/assistant messages
    # in the correct format (like OpenRouter) by the common code at the start of call_api.
    # Prepend system message to history_json_array
    local messages_json_array_with_system=$(echo "$history_json_array" | jq --argjson sys_msg "$system_message_json" '[$sys_msg] + .')

    payload=$(jq -n --arg model "$MODEL" --argjson messages "$messages_json_array_with_system" \
      '{model: $model, messages: $messages, stream: false}')

    # Make the curl request for Ollama
    # No API key typically, but Content-Type is needed.
    response=$(curl -s -X POST "$api_url" \
      -H "Content-Type: application/json" \
      -d "$payload")

    # Parse Ollama response and handle errors
    if [ $? -ne 0 ]; then
      echo "Error: curl command failed for Ollama. Is Ollama running at $ollama_url_base?" >&2
      return 1
    fi

    # Check for API errors in the response JSON (Ollama might return error in 'error' field)
    error_message=$(echo "$response" | jq -r '.error // empty')
    if [ -n "$error_message" ]; then
      echo "Error: Ollama API Error: $error_message" >&2
      # echo "Full Ollama Response: $response" >&2 # for debugging
      return 1
    fi

    # Extract the text content. Ollama nests it under message -> content
    response_text=$(echo "$response" | jq -r '.message.content // empty')

    if [ -z "$response_text" ]; then
      # Check if 'done' is false and 'message' is missing, could be an issue
      local done_status=$(echo "$response" | jq -r '.done // "true"')
      if [[ "$done_status" != "true" ]] && ! echo "$response" | jq -e '.message' > /dev/null; then
         echo "Error: Ollama response incomplete or malformed (done: $done_status, no message content)." >&2
      else
         echo "Error: Could not extract text from Ollama response. It might be empty." >&2
      fi
      # echo "Full Ollama Response: $response" >&2 # for debugging
      return 1
    fi
  elif [[ "$API_PROVIDER" =~ ^[A-Z0-9_]+$ ]]; then
    # Handle custom OpenAI-compatible providers
    local provider_url_var="CUSTOM_PROVIDER_${API_PROVIDER}_URL"
    local provider_key_var="CUSTOM_PROVIDER_${API_PROVIDER}_KEY"
    
    local base_url="${!provider_url_var}"
    local api_key="${!provider_key_var}"
    
    if [ -z "$base_url" ] || [ -z "$api_key" ]; then
      echo "Error: Custom provider $API_PROVIDER not properly configured." >&2
      return 1
    fi
    
    # Remove trailing slash if present
    base_url=$(echo "$base_url" | sed 's:/*$::')
    api_url="${base_url}/chat/completions"

    # Custom provider payload structure (same as OpenAI/OpenRouter)
    local system_message_json=$(create_api_message_json "system" "$SYSTEM_PROMPT")
    local messages_json_array_with_system=$(echo "$history_json_array" | jq --argjson sys_msg "$system_message_json" '[$sys_msg] + .')

    payload=$(jq -n --arg model "$MODEL" --argjson messages "$messages_json_array_with_system" \
      '{model: $model, messages: $messages}')

    # Make the curl request for custom provider
    response=$(curl -s -X POST "$api_url" \
      -H "Content-Type: application/json" \
      -H "Authorization: Bearer $api_key" \
      -d "$payload")

    # Parse custom provider response and handle errors
    if [ $? -ne 0 ]; then
      echo "Error: curl command failed for $API_PROVIDER." >&2
      return 1
    fi

    # Check for API errors
    error_message=$(echo "$response" | jq -r '.error.message // empty')
    if [ -n "$error_message" ]; then
      echo "Error: $API_PROVIDER API Error: $error_message" >&2
      # echo "Full $API_PROVIDER Response: $response" >&2 # for debugging
      return 1
    fi

    # Extract the text content. Custom providers should follow OpenAI format.
    response_text=$(echo "$response" | jq -r '.choices[0].message.content // empty')

    if [ -z "$response_text" ]; then
      echo "Error: Could not extract text from $API_PROVIDER response." >&2
      # echo "Full $API_PROVIDER Response: $response" >&2 # for debugging
      return 1
    fi
  else
    echo "Error: Unknown API Provider '$API_PROVIDER'." >&2
    return 1
  fi

  # 5. Return the extracted text
  echo "$response_text"
  return 0 # Indicate success
}

# --- Command Copying ---
# --- Command Copying ---
handle_command_copying() {
  local ai_response="$1"
  local line
  local -a current_copyable_items=() # Stores the full content of each block
  local -a current_preview_items=()  # Stores the preview text for gum choose
  local in_block=0
  local current_block=""

  # Use process substitution to read line by line to extract blocks
  while IFS= read -r line || [[ -n "$line" ]]; do
    # Check for closing delimiter first (only if already in a block)
    if [[ $in_block -eq 1 ]] && [[ "$line" =~ ^[[:space:]]*\`\`\`[[:space:]]*$ ]]; then
      # Exiting a block
      in_block=0
      # Store the accumulated block content if non-empty
      if [ -n "$current_block" ]; then
        # Use printf '%s' to avoid adding an extra newline at the end
        local full_block_content
        full_block_content=$(printf '%s' "$current_block")
        current_copyable_items+=("$full_block_content")

        # --- Create the preview for gum choose ---
        local first_line
        # Get first line for display, add ellipsis if multi-line
        first_line=$(echo "$full_block_content" | head -n 1)
        # Check if the block has more than one line
        if [[ $(echo "$full_block_content" | wc -l) -gt 1 ]]; then
          # Check if the first line is shorter than, say, 60 chars before adding ellipsis
          if [[ ${#first_line} -lt 60 ]]; then
            first_line="$first_line [...]"
          else
            # If the first line is long, show the first 60 chars + ellipsis
            first_line="${first_line:0:60} [...]"
          fi
        fi
        # Add the formatted first line to the preview list
        current_preview_items+=("$first_line")
        # --- End preview creation ---

      fi
      current_block="" # Reset for next potential block
    # Check for opening delimiter (only if not already in a block)
    # Allows optional language specifier (alphanumeric, _, -)
    elif [[ $in_block -eq 0 ]] && [[ "$line" =~ ^[[:space:]]*\`\`\`[a-zA-Z0-9_-]*[[:space:]]*$ ]]; then
      # Entering a block
      in_block=1
      current_block="" # Start accumulating (exclude delimiter line)
    # If inside a block, accumulate content
    elif [[ $in_block -eq 1 ]]; then
      # Append line to current block, handling the first line vs subsequent lines
      if [ -z "$current_block" ]; then
        current_block="$line"
      else
        # Append with a newline separator
        current_block=$(printf '%s\n%s' "$current_block" "$line")
      fi
    fi
  done < <(printf '%s\n' "$ai_response") # Feed the AI response line by line

  # Note: Unterminated blocks (if the AI response ends mid-block) are ignored.

  if [ ${#current_copyable_items[@]} -eq 0 ]; then
    # No code blocks found in the AI response, nothing to do.
    # You might want to add a small message or just return silently.
    # For now, returning silently is fine.
    return
  fi

  # --- Start of the new copy loop ---
  while true; do
    if [ ${#current_copyable_items[@]} -eq 0 ]; then
      echo ""
      gum style "All available blocks copied from this response." --foreground "#00FF00" --bold
      echo ""
      break # Exit the loop, all items have been copied
    fi

    local gum_options_string=""
    local -a gum_display_previews=() # To hold previews for current gum prompt

    # Prepare options for gum choose from *currently available* previews
    for i in "${!current_preview_items[@]}"; do
      # The user sees 0-indexed options for the *current* list
      gum_display_previews+=("$i: ${current_preview_items[$i]}")
    done

    local cancel_loop_option_text="Stop Copy loop"
    # Add the cancel option to the display previews array as the last item
    gum_display_previews+=("$cancel_loop_option_text")

    # Convert array to newline-separated string for gum choose
    gum_options_string=$(printf "%s\n" "${gum_display_previews[@]}")

    local selected_option
    local item_count=${#current_preview_items[@]}
    # Dynamic height for gum choose: items + cancel_option + header_space
    local gum_height=$((item_count + 2)) # +1 for cancel option, +1 for header
    # Set min/max height for gum prompt (optional, but good for usability)
    if [[ "$gum_height" -lt 3 ]]; then gum_height=3; fi
    if [[ "$gum_height" -gt 15 ]]; then gum_height=15; fi


    selected_option=$(printf "%b" "$gum_options_string" | gum choose --header "Select code block to copy (or cancel):" --height "$gum_height")
    local gum_status=$? # Capture gum exit status

    # Check if the user selected the cancel option or if gum choose was cancelled (e.g., Esc/Ctrl+C)
    if [[ "$selected_option" == "$cancel_loop_option_text" ]] || [[ -z "$selected_option" ]]; then
      echo ""
      gum style "Exiting copy mode." --foreground "#FFFF00" --bold # Yellow for exiting
      echo ""
      break # Exit the main copy loop
    else
      # Extract the index chosen by the user. This index is relative to the current list.
      local selected_index_str=$(echo "$selected_option" | cut -d':' -f1)
      local selected_index=-1

      if [[ "$selected_index_str" =~ ^[0-9]+$ ]]; then
        selected_index=$selected_index_str
      fi

      # Validate selected_index (it should be valid if gum choose worked correctly)
      if [[ $selected_index -ne -1 ]] && [[ $selected_index -lt ${#current_copyable_items[@]} ]]; then
        local content_to_copy="${current_copyable_items[$selected_index]}"
        # Get the preview text of the item being copied *before* removal for the message
        local copied_preview_text="${current_preview_items[$selected_index]}"


        printf '%s' "$content_to_copy" | $CLIPBOARD_TOOL
        echo ""
        # Use the stored preview for the message
        gum style "Copied block: '${copied_preview_text}' to clipboard!" --foreground "#00FF00" --bold
        echo ""

        # --- Remove the copied item from BOTH arrays ---
        local temp_copyable=()
        for i in "${!current_copyable_items[@]}"; do
          if [[ $i -ne $selected_index ]]; then
            temp_copyable+=("${current_copyable_items[$i]}")
          fi
        done
        current_copyable_items=("${temp_copyable[@]}")
        unset temp_copyable

        local temp_previews=()
        for i in "${!current_preview_items[@]}"; do
          if [[ $i -ne $selected_index ]]; then
            temp_previews+=("${current_preview_items[$i]}")
          fi
        done
        current_preview_items=("${temp_previews[@]}")
        unset temp_previews
        # --- End of removal ---
      else
        echo "Error: Invalid selection. Returning to copy options." >&2
        # Loop will continue, offering the (now potentially smaller) list again.
      fi
    fi
  done
  # --- End of the new copy loop ---
}

# --- Main Chat Loop ---
main() {
  # Display ASCII Art Welcome
  echo -e "\e[1;35m" # Set color to bold purple
  echo "  ,-.   ,-."
  echo " ( O ) (o.o)"
  echo "  \`-’   |_)  oliveowl"
  echo "    “who?”"
  echo -e "\e[0m" # Reset color and bold

  # Check if configuration is complete before starting chat
  if [ -z "$API_PROVIDER" ] || [ -z "$MODEL" ]; then
    echo "Configuration incomplete." >&2
    echo "Please run '$0 --config' first to select API provider and model." >&2
    exit 1
  fi

  # Display Welcome and Instructions using echo with ANSI colors
  echo -e "\e[0;94mWelcome to \e[1;35mOliveOwl\e[0;94m! Provider: \e[1;35m$API_PROVIDER\e[0;94m, Model: \e[1;35m$MODEL\e[0m"
  echo -e "\e[0;94mType \e[1;35m/exit\e[0;94m, \e[1;35m/quit\e[0;94m, or \e[1;35m/q\e[0;94m to quit, \e[1;35m/history\e[0;94m or \e[1;35m/h\e[0;94m to load previous chat, \e[1;35m/new\e[0;94m or \e[1;35m/n\e[0;94m for new chat, \e[1;35m/config\e[0;94m or \e[1;35m/c\e[0;94m to reconfigure.\e[0m"
  echo -e "\e[0;94mType \e[1;35m/view\e[0;94m or \e[1;35m/v\e[0;94m to open current history in \e[1;35m$EDITOR\e[0m.\e[0m" # Add /view instruction

  # Initial session setup
  local start_status
  start_new_session # Prompt user for initial action (new, history, exit)
  start_status=$?
  if [ $start_status -ne 0 ]; then
    # If /history was chosen but failed or was cancelled, start a default new session
    echo "Falling back to new timestamped session."
    create_new_session_file "" # Create default timestamped session
  fi
  # If start_new_session exited or loaded history successfully, we proceed

# Main loop
# Variable to hold prompt history selection
local prompt_history_selected=""

while true; do
  # --- Read user input with history support ---
  local user_input=""    
    # Display prompt manually before calling gum
    printf '\n\e[1;34mYou:\e[0m\n' # Print "You:" prompt on its own line (Bold Blue)

    # Use gum write to capture multi-line input.
    # Set a placeholder to encourage input.
    # Capture output into user_input variable.
    # Check exit status to handle cancellation (e.g., Ctrl+C in gum write).
    if ! user_input=$(gum write --value "$prompt_history_selected" --placeholder "Enter your prompt (Ctrl+D or Esc to finish)..."); then
      # gum write was cancelled (e.g., Ctrl+C)
      echo "Input cancelled."
      # Decide how to handle cancellation - exit or continue loop? Let's continue.
      continue
    fi
    
    # Clear the prompt history selection after use
    prompt_history_selected=""
    # --- End of gum write input ---

    # Trim leading/trailing whitespace (optional, but often helpful)
    # Also remove placeholder if user exits without typing (gum write might return the placeholder)
    local placeholder="Enter your prompt (Ctrl+D or Esc to finish)..."
    if [[ "$user_input" == "$placeholder" ]]; then
      user_input=""
    fi
    user_input=$(echo "$user_input" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//')

    # Handle commands (check the input captured from gum)
    # Commands must be the entire input now, not just the first line.
    local is_command_handled=0 # Flag to check if a command was processed
    case "$user_input" in
    "/exit"|"/quit"|"/q")
      is_command_handled=1
      echo "Exiting OliveOwl."
      break
      ;;
    "/new"|"/n")
      is_command_handled=1
      # Clear the terminal screen like /history does
      clear
      # Prompt again for new session name, history load, or exit
      local new_status
      start_new_session
      new_status=$?
      if [ $new_status -ne 0 ]; then
        # If /history was chosen but failed or was cancelled, start a default new session
        echo "Falling back to new timestamped session."
        create_new_session_file "" # Create default timestamped session
      fi
      continue # Skip API call for this turn
      ;;
    "/history"|"/h"|"/hist")
      is_command_handled=1
      # Just attempt to load history, don't change session if cancelled
      select_and_load_history
      ;;
    "/prompt-history"|"/ph"|"/p")
      is_command_handled=1
      # Let user select from prompt history
      local selected_prompt
      if selected_prompt=$(select_from_prompt_history); then
        # Pre-fill the next input with the selected prompt
        prompt_history_selected="$selected_prompt"
        # Continue to show the prompt again with pre-filled value
        continue
      else
        echo "Prompt history selection cancelled."
        continue  # Skip to next iteration
      fi
      ;;
    "/config"|"/c"|"/conf")
      is_command_handled=1
      echo "Switching to config..."
      if configure_settings; then
        # Re-source config in case it changed
        # Check if CONFIG_FILE exists before sourcing
        if [ -f "$CONFIG_FILE" ]; then
          source "$CONFIG_FILE"
        fi
        echo "Config updated. Provider: $API_PROVIDER, Model: $MODEL"
      else
        echo "Configuration cancelled or failed. Continuing with current settings."
      fi
      # After config, continue the current session with the new settings or current settings
      continue
      ;;
    "/view"|"/v")
      is_command_handled=1
      view_history # Call the new function
      continue     # Skip API call for this turn
      ;;
    "/help"|"/?")
      is_command_handled=1
      show_help
      continue
      ;;

    "")                    # Handle empty input
      is_command_handled=1 # Treat empty input like a command (skip API call)
      continue
      ;;
    esac

    # If input was not a command and not empty, print it back before API call
    if [[ $is_command_handled -eq 0 ]] && [ -n "$user_input" ]; then
      # The "You:" prompt was already printed before gum write.
      # Just print the captured input. Use printf to handle potential special chars.
      printf '%s\n' "$user_input"
      
      # Save non-empty prompts to history
      save_prompt_to_history "$user_input"
    fi

    # --- Prepare for API call ---
    local final_user_input="$user_input"
    local rename_request=false

    # Only add non-empty, non-command input to history and process for API call
    if [[ $is_command_handled -eq 0 ]] && [ -n "$user_input" ]; then
      ((USER_PROMPT_COUNT++)) # Increment user prompt counter

      # Check if we should request a session rename
      # Trigger on 1st, 6th, 11th, etc. prompt for non-manual sessions
      if [[ "$IS_MANUALLY_NAMED" == "false" ]] && { [[ $USER_PROMPT_COUNT -eq 1 ]] || [[ $(( (USER_PROMPT_COUNT - 1) % 5 )) -eq 0 && $USER_PROMPT_COUNT -gt 1 ]]; }; then
        rename_request=true
        # Append special instruction to the user's prompt
        final_user_input=$(printf '%s\n\n(Side task: Suggest a short title for our conversation. Please start your entire response with the title in the format <| The Title |> and then give your main answer.)' "$user_input")
      fi

    # Skip API call if it was a command or empty input
    else
      continue
    fi

    # Call the API with gum spinner (no retry logic)
    local ai_response=""      # Initialize ai_response
    local api_call_status=1   # Initialize to failure status (1 means error)
    local spinner_pid
    local start_time end_time # Declare variables for timing

    # Start spinner in the background
    gum spin --spinner dot --title "Waiting for AI response..." -- sleep infinity &
    spinner_pid=$!
    trap 'kill $spinner_pid 2>/dev/null; trap - EXIT' EXIT

    # Record start time for token speed calculation
    start_time=$(date +%s%N) # Nanoseconds for precision

    # Make the actual API call
    ai_response=$(call_api "$final_user_input")
    api_call_status=$? # Capture exit status immediately

    # Record end time
    end_time=$(date +%s%N)

    # Stop spinner and clean up trap
    kill $spinner_pid 2>/dev/null || true
    trap - EXIT

    # Check if API call failed
    if [ $api_call_status -ne 0 ] || [ -z "$ai_response" ]; then
      # API call failed
      # The error message from call_api should have already been printed
      echo "Error: API call failed. Returning to prompt."
      # Continue to the next iteration of the main chat loop to get new user input.
      continue
    fi

    # --- Parse for session rename if requested ---
    if [[ "$rename_request" == "true" ]]; then
      # Check if the response starts with the special tag <|...|>
      # We check only the first line of the response for the tag.
      if echo "$ai_response" | head -n 1 | grep -q "^<| .* |>"; then
        # Extract the name from between <| and |> from the first line.
        local new_session_name
        new_session_name=$(echo "$ai_response" | head -n 1 | sed -n 's/^<| \(.*\) |>.*$/\1/p')

        if [ -n "$new_session_name" ]; then
          # Trim leading/trailing whitespace from the extracted name
          CURRENT_SESSION_NAME=$(echo "$new_session_name" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//')
          # 1. Remove the <|...|> tag part from the first line.
          # 2. If the first line is now empty (because the tag was on its own line), delete it.
          ai_response=$(echo "$ai_response" | sed -e '1s/^<|.*|>[[:space:]]*//' -e '1{/^$/d;}')
          # Save history immediately to persist the new name
          save_history
        else
          # Tag was detected but name extraction failed.
          echo "Info: AI response format for title was incorrect. Continuing with current name." >&2
        fi
      else
        # AI didn't provide the title tag. The full ai_response is used as is.
        echo "Info: AI did not provide a session title. Continuing with current name." >&2
      fi
    fi

    # --- Calculate and Display Token Speed ---
    # This block runs only if the API call was ultimately successful (api_call_status == 0)
    # and ai_response is populated. The start_time and end_time were set
    # during the *last successful attempt* within the retry loop.

    local time_taken_ns=$((end_time - start_time))
    # Convert nanoseconds to seconds (float) for calculation. Use bc for precision.
    # Ensure scale is sufficient for small time differences.
    local time_taken_s=$(echo "scale=3; $time_taken_ns / 1000000000" | bc)

    # Count words in the AI response
    # `wc -w` counts words. Use process substitution to feed the response.
    local word_count=$(echo -n "$ai_response" | wc -w)

    local words_per_second="N/A" # Default if time_taken_s is zero or word_count is zero

    # Avoid division by zero if time_taken_s is 0 or very close to 0, or if word_count is 0
    # bc returns 0 if the result of division is less than its scale allows without truncation.
    # We should check if time_taken_s is greater than a very small threshold like 0.001.
    # A simpler check is if word_count > 0 and time_taken_s is not "0.000" (or similar for your scale)
    if [ "$word_count" -gt 0 ] && [ "$(echo "$time_taken_s > 0.000" | bc -l)" -eq 1 ]; then
        words_per_second=$(echo "scale=2; $word_count / $time_taken_s" | bc)
    elif [ "$word_count" -gt 0 ] && [ "$time_taken_s" == "0.000" ]; then # Handle case where time is too short for bc with current scale
        words_per_second="Effectively Instant (words: $word_count)"
    fi
    # --- End of Token Speed Calculation ---


    # Add User and AI responses to history
    local user_api_json=$(create_api_message_json "user" "$user_input") # History stores original prompt
    local user_history_item=$(jq -n --argjson msg "$user_api_json" --arg model "" \
      '{message_json: $msg, model_used: $model}' | jq -c .)
    CHAT_HISTORY+=("$user_history_item")

    local response_role="assistant" # Default for OpenRouter
    if [[ "$API_PROVIDER" == "Gemini" ]]; then
      response_role="model"
    fi
    local ai_api_json=$(create_api_message_json "$response_role" "$ai_response")
    local ai_history_item=$(jq -n --argjson msg "$ai_api_json" --arg model "$MODEL" \
      '{message_json: $msg, model_used: $model}' | jq -c .)
    CHAT_HISTORY+=("$ai_history_item")

    # Save history after adding both messages
    if ! save_history; then
      echo "Warning: Failed to save history after AI response." >&2
    fi

    # Interpret the raw AI response first
    local interpreted_response
    interpreted_response=$(printf '%b' "$ai_response")

    # Display interpreted AI response using bat
    echo -e "\n\e[1;35mAI ($MODEL):\e[0m"                                        # Bold Purple for AI
    echo -e "\e[1m\e[35m=============================================\e[0m" # Bold Purple for AI, showing specific model, with explicit reset
    echo "$interpreted_response" | bat --language md --paging=never --style=plain --color=always
    # Display token speed (words per second)
    # Use a different color for the speed, perhaps a subtle grey or cyan.
    # \e[0;90m is dark grey. \e[0;36m is cyan.
    echo -e "\e[0;90mSpeed: $words_per_second words/sec (Time: ${time_taken_s}s, Words: $word_count)\e[0m"
    echo -e "\e[0;90mSession: $CURRENT_SESSION_NAME\e[0m" # Display current session name
    
    # Show hint about prompt history commands in grey
    if [ -f "$PROMPT_HISTORY_FILE" ] && [ -s "$PROMPT_HISTORY_FILE" ]; then
      echo -e "\e[0;90m(Hint: Type /prompt-history or /ph or /p to select from previous prompts)\e[0m"
    fi

    # Handle command copying using the interpreted response
    handle_command_copying "$interpreted_response"

  done
}

# --- Argument Parsing ---
if [[ "$1" == "--config" ]]; then
  setup_config # Ensure .env is loaded first for API key checks
  configure_settings
  exit 0
fi

# Cleanup function to remove temporary files
cleanup_temp_files() {
  # Clean up any temporary files (cache is now persistent, don't delete)
  return 0
}

# Set up trap to cleanup on exit
trap cleanup_temp_files EXIT

# --- Run Setup and Main ---
# Setup needs to run first to load config for main
setup_config
# Check EDITOR dependency after setup_config sets/loads it
# Extract just the command from the EDITOR variable (e.g., "nvim" from "nvim -R")
read -r editor_command _ <<< "$EDITOR"
check_dependency "$editor_command"

# Build history cache in background for faster /h performance
build_history_cache >/dev/null 2>&1 &

main

exit 0