#!/usr/bin/env bash

# Copyright (c) 2026 The Eclipse Foundation
#
# This program and the accompanying materials are made available under the
# terms of the Eclipse Public License 2.0 which is available at
# http://www.eclipse.org/legal/epl-2.0, or the MIT Software License
# which is available at https://opensource.org/license/mit.
#
# SPDX-License-Identifier: EPL-2.0 OR MIT
#
#  Contributors:
#      asgomes - Initial implementation

set -e

# =======================================================
# Validated Models
# =======================================================
MODELS_LOCAL=("devstral-small-2" "qwen3:32b" "qwen3-coder:30b" "Enter manually")
MODELS_OPENAI=("gpt-4o-mini" "gpt-4o" "Enter manually")
MODELS_AZURE=("gpt-5-mini" "o4-mini" "Enter manually")
MODELS_GROQ=("llama-4-scout" "Enter manually")
MODELS_NIM=("moonshotai/kimi-k2-instruct" "Enter manually")
# =======================================================

echo "======================================================="
echo "      MOSAICO Demonstrator - Automated Setup Script    "
echo "======================================================="
echo ""

# Dependency checks & auto-install
echo "Checking system requirements..."

check_cmd() {
    command -v "$1" >/dev/null 2>&1
}

install_pkg() {
    local pkg=$1
    echo "⚠️  Missing '$pkg'. Attempting to install..."
    
    if check_cmd apt-get; then
        sudo apt-get update && sudo apt-get install -y "$pkg"
    elif check_cmd dnf; then
        sudo dnf install -y "$pkg"
    elif check_cmd yum; then
        sudo yum install -y "$pkg"
    elif check_cmd pacman; then
        sudo pacman -Sy --noconfirm "$pkg"
    elif check_cmd zypper; then
        sudo zypper install -y "$pkg"
    elif check_cmd brew; then
        brew install "$pkg"
    else
        echo "❌ Could not determine package manager to install $pkg. Please install it manually."
        exit 1
    fi
}

install_docker() {
    echo "⚠️ Docker is missing."
    if [ "$(uname -s)" = "Darwin" ]; then
        echo "❌ Docker cannot be easily installed via CLI on macOS."
        echo "Please install Docker Desktop from: https://docs.docker.com/desktop/install/mac-install/"
        exit 1
    fi
    
    read -p "Would you like to automatically install Docker using the official get.docker.com script? (y/n) [y]: " INSTALL_DOCKER < /dev/tty
    INSTALL_DOCKER=${INSTALL_DOCKER:-y}
    
    if [[ "$INSTALL_DOCKER" =~ ^[Yy]$ ]]; then
        echo "Downloading and running Docker convenience script..."
        curl -fsSL https://get.docker.com -o get-docker.sh
        sudo sh get-docker.sh
        rm get-docker.sh
        
        echo "Starting Docker service..."
        sudo systemctl start docker || true
        sudo systemctl enable docker || true
        
        echo "Adding current user to the 'docker' group..."
        sudo usermod -aG docker "$USER" || true
        
        echo "✅ Docker installed successfully."
        echo "⚠️ IMPORTANT: You may need to log out and log back in, or run 'newgrp docker' for permission changes to take effect."
    else
        echo "❌ Please install Docker manually and rerun this script."
        exit 1
    fi
}

# Check basic utils
if ! check_cmd curl; then install_pkg curl; fi
if ! check_cmd unzip; then install_pkg unzip; fi

# Check Docker
if ! check_cmd docker; then install_docker; fi

# Check Docker Compose (Modern docker install includes it as a plugin)
if ! (docker compose version >/dev/null 2>&1 || check_cmd docker-compose); then
    echo "❌ Docker is installed, but Docker Compose is missing."
    echo "Please ensure you have the 'docker-compose-plugin' installed."
    exit 1
fi

echo "✅ All dependencies met."
echo ""

# Download and extract repository to temp directory
ZIP_URL="https://gitlab.eclipse.org/eclipse-research-labs/mosaico-project/mosaico-demonstrator/-/archive/main/mosaico-demonstrator-main.zip"
TEMP_DIR=$(mktemp -d)
DIR_NAME="mosaico-demonstrator"

echo "Downloading MOSAICO demonstrator to a temporary directory..."
cd "$TEMP_DIR"
curl -sS -L -o mosaico.zip "$ZIP_URL"

echo "Extracting archive..."
unzip -q mosaico.zip
rm mosaico.zip

mv mosaico-demonstrator-* "$DIR_NAME"
cd "$DIR_NAME"

# Initial environment setup
echo ""
echo "Generating default environment files..."
./00-copy-env.sh

# Interactive LLM configuration
select_model() {
    local options=("$@")
    local default_model="${options[0]}"
    echo ""
    echo "Select a model:"
    
    local i=1
    for opt in "${options[@]}"; do
        echo "$i) $opt"
        ((i++))
    done
    
    read -p "Enter choice [Default: 1]: " REPLY < /dev/tty
    if [[ -z "$REPLY" ]]; then
        MODEL_NAME="$default_model"
    elif [[ "$REPLY" =~ ^[0-9]+$ ]] && [ "$REPLY" -ge 1 ] && [ "$REPLY" -le "${#options[@]}" ]; then
        local idx=$((REPLY - 1))
        local chosen="${options[$idx]}"
        if [[ "$chosen" == "Enter manually" ]]; then
            read -p "Enter custom model name: " MODEL_NAME < /dev/tty
        else
            MODEL_NAME="$chosen"
        fi
    else
        echo "⚠️ Invalid selection. Defaulting to $default_model."
        MODEL_NAME="$default_model"
    fi
}

echo ""
echo "======================================================="
echo "               LLM Provider Configuration              "
echo "======================================================="
echo "1) Local (Ollama / LM Studio)"
echo "2) OpenAI"
echo "3) Azure AI Foundry"
echo "4) Groq"
echo "5) NVIDIA NIM"
echo "6) Custom (OpenAI-Compatible Endpoint)"
echo ""
read -p "Enter your choice (1-6) [Default: 1]: " PROVIDER_CHOICE < /dev/tty
PROVIDER_CHOICE=${PROVIDER_CHOICE:-1}

API_BASE=""
API_KEY=""
MODEL_NAME=""

case $PROVIDER_CHOICE in
    1)
        API_BASE="http://host.docker.internal:11434/v1"
        read -p "Enter API Base URL [$API_BASE]: " INPUT_BASE < /dev/tty
        API_BASE=${INPUT_BASE:-$API_BASE}
        API_KEY="ollama"
        select_model "${MODELS_LOCAL[@]}"
        ;;
    2)
        API_BASE="https://api.openai.com/v1"
        read -s -p "Enter your OpenAI API Key: " API_KEY < /dev/tty; echo ""
        select_model "${MODELS_OPENAI[@]}"
        ;;
    3)
        read -p "Enter your Azure API Base URL: " API_BASE < /dev/tty
        read -s -p "Enter your Azure API Key: " API_KEY < /dev/tty; echo ""
        select_model "${MODELS_AZURE[@]}"
        ;;
    4)
        API_BASE="https://api.groq.com/openai/v1"
        read -s -p "Enter your Groq API Key: " API_KEY < /dev/tty; echo ""
        select_model "${MODELS_GROQ[@]}"
        ;;
    5)
        API_BASE="https://integrate.api.nvidia.com/v1"
        read -s -p "Enter your NVIDIA API Key: " API_KEY < /dev/tty; echo ""
        select_model "${MODELS_NIM[@]}"
        ;;
    6)
        read -p "Enter Custom API Base URL: " API_BASE < /dev/tty
        read -s -p "Enter Custom API Key: " API_KEY < /dev/tty; echo ""
        read -p "Enter Model Name: " MODEL_NAME < /dev/tty
        ;;
    *) echo "❌ Invalid choice."; exit 1 ;;
esac

update_env_var() {
    local key=$1
    local value=$2
    local file=$3
    # If the key exists UNCOMMENTED, replace it.
    if grep -q "^[[:space:]]*${key}=" "$file"; then
        sed "s|^[[:space:]]*${key}=.*|${key}=${value}|" "$file" > "${file}.tmp" && mv "${file}.tmp" "$file"
    else
        # If it doesn't exist or is commented out, append it to the end of the file.
        # Docker Compose uses the last defined value in .env files.
        echo "${key}=${value}" >> "$file"
    fi
}

echo "Updating env/llm.env with your configuration..."
update_env_var "API_BASE" "$API_BASE" "env/llm.env"
update_env_var "API_KEY" "$API_KEY" "env/llm.env"

# Safely determine the prefixed and unprefixed model strings
if [[ "$MODEL_NAME" == openai/* ]]; then
    PREFIXED_MODEL="$MODEL_NAME"
    UNPREFIXED_MODEL="${MODEL_NAME#openai/}"
else
    PREFIXED_MODEL="openai/$MODEL_NAME"
    UNPREFIXED_MODEL="$MODEL_NAME"
fi

# Apply the specific model variable logic
update_env_var "EMFATIC_GENERATOR_MODEL" "$PREFIXED_MODEL" "env/llm.env"
update_env_var "EMFATIC_SEMANTIC_SUPERVISOR_MODEL" "$UNPREFIXED_MODEL" "env/llm.env"
update_env_var "MINI_SWE_MODEL" "$PREFIXED_MODEL" "env/llm.env"
update_env_var "IP_AGENT_MODEL" "$PREFIXED_MODEL" "env/llm.env"

# Specific override for Groq
if [ "$PROVIDER_CHOICE" = "4" ]; then
    echo "Applying specific configurations for Groq..."
    update_env_var "MSWEA_MODEL_CONFIG" "/etc/miniswe/groq-model-config.yaml" "env/llm.env"
fi

# Persistence logic
echo ""
echo "======================================================="
echo "                 Installation Location                 "
echo "======================================================="
read -p "Do you want to make this setup persistent? (If no, it will remain in /tmp and be lost on reboot) (y/n) [y]: " MAKE_PERSISTENT < /dev/tty
MAKE_PERSISTENT=${MAKE_PERSISTENT:-y}

if [[ "$MAKE_PERSISTENT" =~ ^[Yy]$ ]]; then
    DEFAULT_INSTALL_DIR="$HOME/.mosaico-demonstrator"
    read -p "Enter installation directory [$DEFAULT_INSTALL_DIR]: " INSTALL_DIR < /dev/tty
    INSTALL_DIR=${INSTALL_DIR:-$DEFAULT_INSTALL_DIR}
    
    # Expand tilde if present
    INSTALL_DIR="${INSTALL_DIR/#\~/$HOME}"
    
    if [ -d "$INSTALL_DIR" ]; then
        echo "⚠️ Directory $INSTALL_DIR already exists."
        while true; do
            read -p "Do you want to (O)verwrite with new setup, (A)bort, or (P)roceed using existing files? [O/A/P]: " DIR_ACTION < /dev/tty
            case "$DIR_ACTION" in
                [Oo]* )
                    echo "Removing existing directory and replacing with new setup..."
                    rm -rf "$INSTALL_DIR"
                    mv "$TEMP_DIR/$DIR_NAME" "$INSTALL_DIR"
                    TARGET_DIR="$INSTALL_DIR"
                    break
                    ;;
                [Pp]* )
                    echo "Proceeding with existing folder (newly downloaded files and configs will be discarded)..."
                    TARGET_DIR="$INSTALL_DIR"
                    break
                    ;;
                [Aa]* )
                    echo "Setup aborted by user. New files remain in $TEMP_DIR/$DIR_NAME temporarily."
                    exit 0
                    ;;
                * )
                    echo "Invalid choice. Please enter O, A, or P."
                    ;;
            esac
        done
    else
        echo "Moving files to $INSTALL_DIR..."
        mv "$TEMP_DIR/$DIR_NAME" "$INSTALL_DIR"
        TARGET_DIR="$INSTALL_DIR"
    fi
else
    TARGET_DIR="$TEMP_DIR/$DIR_NAME"
    echo "Files will remain in temporary directory: $TARGET_DIR"
fi

cd "$TARGET_DIR"

# MOSAICO Demonstrator CLI setup
echo ""
echo "Configuring the 'mosaico' CLI..."
SHELL_RC=""
if [ -f "$HOME/.zshrc" ]; then
    SHELL_RC="$HOME/.zshrc"
elif [ -f "$HOME/.bashrc" ]; then
    SHELL_RC="$HOME/.bashrc"
fi

if [ -n "$SHELL_RC" ]; then
    if ! grep -q "mosaico()" "$SHELL_RC"; then
        cat << EOF >> "$SHELL_RC"

# MOSAICO Demonstrator CLI
mosaico() {
    case "\$1" in
        start|up)
            cd "$TARGET_DIR" && ./01-compose.sh up -d
            ;;
        stop|down)
            cd "$TARGET_DIR" && ./01-compose.sh stop
            ;;
        status|ps)
            cd "$TARGET_DIR" && ./01-compose.sh ps
            ;;
        logs)
            cd "$TARGET_DIR" && ./01-compose.sh logs "\${@:2}"
            ;;
        pull)
            cd "$TARGET_DIR" && ./01-compose.sh pull "\${@:2}"
            ;;
        dir)
            cd "$TARGET_DIR"
            ;;
        *)
            echo "Usage: mosaico {start|stop|status|logs|pull|dir}"
            echo "  start  - Starts the MOSAICO services"
            echo "  stop   - Stops the MOSAICO services"
            echo "  status - Shows the status of the containers"
            echo "  logs   - Views output from containers (e.g., 'mosaico logs -f')"
            echo "  pull   - Pulls service images from the registry"
            echo "  dir    - Navigates to the installation directory"
            ;;
    esac
}
EOF
        echo "✅ Added 'mosaico' command to $SHELL_RC"
    else
        echo "⚠️ 'mosaico' command already exists in $SHELL_RC. Skipping."
    fi
else
    echo "⚠️ Could not find .bashrc or .zshrc. You can manually run commands from: $TARGET_DIR"
fi

# Start services & check status
echo ""
echo "======================================================="
echo "             Starting MOSAICO Services                 "
echo "======================================================="

# If Docker was just installed, we might need to run the compose script with sudo
if ! docker ps >/dev/null 2>&1; then
    echo "⚠️ Warning: The current user does not have permission to communicate with the Docker daemon."
    echo "Attempting to start services with sudo..."
    sudo ./01-compose.sh up -d
    echo ""
    echo "Checking service status..."
    sudo ./01-compose.sh ps
else
    ./01-compose.sh up -d
    echo ""
    echo "Checking service status..."
    ./01-compose.sh ps
fi

# Extract Langfuse keys for VS Code extension config
LANGFUSE_PUB=$(grep '^LANGFUSE_PUBLIC_KEY=' env/langfuse.env | cut -d '=' -f2)
LANGFUSE_SEC=$(grep '^LANGFUSE_SECRET_KEY=' env/langfuse.env | cut -d '=' -f2)

# Final output
echo ""
echo "======================================================="
echo "                    Setup Complete!                    "
echo "======================================================="
echo "The MOSAICO Demonstrator is now running in the background."
echo ""
echo "1. Install the MOSAICO Reference Client in VS Code or Open VSX-compatible editors (e.g., VSCodium or Eclipse Theia):"
echo "   👉 Download: https://open-vsx.org/extension/MOSAICO/mosaico-reference-client"
echo ""
echo "2. Configure the extension: Replace the contents of your MOSAICO Reference Client config file"
echo "   (usually ~/.continue/config.yaml) with the following:"
echo ""
echo "-------------------------------------------------------"
cat <<EOF
name: Local Config
version: 1.0.0
schema: v1
models:
  - name: Mosaico
    provider: mosaico
    model: mosaico-default
    langfusePublicKey: ${LANGFUSE_PUB}
    langfuseSecretKey: ${LANGFUSE_SEC}
    langfuseBaseUrl: http://localhost:3000
    langfuseLogLevel: warn
    referenceRepositoryUrl: http://localhost:8080/mcp
    modelName: ${MODEL_NAME}
    modelApiBase: ${API_BASE}
    modelApiKey: ${API_KEY}
EOF
echo "-------------------------------------------------------"
echo ""
if [ -n "$SHELL_RC" ]; then
    echo "💡 Note: A 'mosaico' command has been added to your shell."
    echo "   To use it in this terminal, run: source $SHELL_RC"
    echo "   Or simply open a new terminal window."
    echo "   Try 'mosaico stop' to shut down the demonstrator later."
else
    echo "To shut down later, navigate to $TARGET_DIR and run: ./01-compose.sh stop"
fi
echo "======================================================="
