#!/bin/bash

# AgentGraph Quick Setup Script
# This script automates the setup process for AgentGraph

set -e  # Exit on any error

echo "🕸️ AgentGraph Quick Setup"
echo "=========================="
echo

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Helper functions
print_success() {
    echo -e "${GREEN}✅ $1${NC}"
}

print_info() {
    echo -e "${BLUE}ℹ️  $1${NC}"
}

print_warning() {
    echo -e "${YELLOW}⚠️  $1${NC}"
}

print_error() {
    echo -e "${RED}❌ $1${NC}"
}

# Check if running in correct directory
if [ ! -f "main.py" ] || [ ! -f "pyproject.toml" ]; then
    print_error "Please run this script from the AgentGraph root directory"
    exit 1
fi

# Function to check if command exists
command_exists() {
    command -v "$1" >/dev/null 2>&1
}

# Function to prompt for OpenAI API key
setup_env_file() {
    print_info "Setting up environment configuration..."
    
    if [ ! -f ".env" ]; then
        if [ -f ".env.example" ]; then
            cp .env.example .env
            print_success "Created .env file from .env.example"
        else
            print_error ".env.example file not found"
            exit 1
        fi
    else
        print_info ".env file already exists, checking configuration..."
    fi
    
    # Check if OPENAI_API_KEY is set
    if grep -q "^OPENAI_API_KEY=your_openai_api_key_here" .env || grep -q "^OPENAI_API_KEY=$" .env; then
        print_warning "OpenAI API key is not configured"
        echo
        echo "Please enter your OpenAI API key (starts with 'sk-'):"
        echo "You can get one at: https://platform.openai.com/account/api-keys"
        read -p "OpenAI API Key: " openai_key
        
        if [ -n "$openai_key" ]; then
            # Update the API key in .env file
            if [[ "$OSTYPE" == "darwin"* ]]; then
                # macOS
                sed -i '' "s/^OPENAI_API_KEY=.*/OPENAI_API_KEY=$openai_key/" .env
            else
                # Linux
                sed -i "s/^OPENAI_API_KEY=.*/OPENAI_API_KEY=$openai_key/" .env
            fi
            print_success "OpenAI API key configured"
        else
            print_warning "No API key provided. You can set it later in the .env file"
        fi
    else
        print_success "OpenAI API key is already configured"
    fi
}

# Function to setup and run with Docker
setup_docker() {
    print_info "Setting up AgentGraph with Docker..."
    
    # Check if Docker is installed
    if ! command_exists docker; then
        print_error "Docker is not installed. Please install Docker first:"
        echo "  - macOS: https://docs.docker.com/desktop/mac/"
        echo "  - Linux: https://docs.docker.com/engine/install/"
        echo "  - Windows: https://docs.docker.com/desktop/windows/"
        exit 1
    fi
    
    # Check if Docker is running
    if ! docker info > /dev/null 2>&1; then
        print_error "Docker is not running. Please start Docker and try again."
        exit 1
    fi
    
    print_success "Docker is available"
    
    # Setup environment file
    setup_env_file
    
    print_info "Building Docker image (this may take a few minutes)..."
    if docker build -t agentgraph . > /dev/null 2>&1; then
        print_success "Docker image built successfully"
    else
        print_error "Failed to build Docker image"
        exit 1
    fi
    
    # Stop existing container if running
    if docker ps -q -f name=agentgraph-app > /dev/null 2>&1; then
        print_info "Stopping existing AgentGraph container..."
        docker stop agentgraph-app > /dev/null 2>&1
        docker rm agentgraph-app > /dev/null 2>&1
        print_success "Existing container stopped"
    fi
    
    print_info "Starting AgentGraph container..."
    if docker run -d --name agentgraph-app -p 7860:7860 --env-file .env agentgraph > /dev/null 2>&1; then
        print_success "AgentGraph container started successfully"
    else
        print_error "Failed to start container"
        exit 1
    fi
    
    # Wait a moment for startup
    sleep 5
    
    # Check if the service is healthy
    print_info "Checking service health..."
    for i in {1..30}; do
        if curl -s http://localhost:7860/api/observability/health-check > /dev/null 2>&1; then
            print_success "AgentGraph is running and healthy!"
            break
        fi
        if [ $i -eq 30 ]; then
            print_warning "Service health check timed out, but container is running"
            break
        fi
        sleep 2
    done
    
    echo
    print_success "🎉 AgentGraph Setup Complete!"
    echo
    echo "🌐 Access AgentGraph at: http://localhost:7860"
    echo "📚 API Documentation: http://localhost:7860/docs"
    echo
    echo "💡 Useful commands:"
    echo "  - View logs: docker logs agentgraph-app"
    echo "  - Stop: docker stop agentgraph-app"
    echo "  - Restart: docker restart agentgraph-app"
    echo
}

# Function to setup for local development
setup_local() {
    print_info "Setting up AgentGraph for local development..."
    
    # Setup environment file
    setup_env_file
    
    # Check Python version
    if command_exists python3; then
        python_cmd="python3"
    elif command_exists python; then
        python_cmd="python"
    else
        print_error "Python is not installed. Please install Python 3.11+ first."
        exit 1
    fi
    
    print_success "Python is available"
    
    print_info "Running AgentGraph setup..."
    if $python_cmd main.py --first-run; then
        print_success "Local development setup completed!"
    else
        print_error "Setup failed. Please check the error messages above."
        exit 1
    fi
}

# Main setup logic
echo "Choose your setup method:"
echo "1) Docker (Recommended - Easy and isolated)"
echo "2) Local Development (For development work)"
echo
read -p "Enter your choice (1 or 2): " choice

case $choice in
    1)
        setup_docker
        ;;
    2)
        setup_local
        ;;
    *)
        print_error "Invalid choice. Please run the script again and choose 1 or 2."
        exit 1
        ;;
esac