#!/bin/bash
# run.sh - Main execution script for Wisal Autism AI Assistant

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

# Function to print colored output
print_status() {
    echo -e "${BLUE}[*]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[+]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[!]${NC} $1"
}

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

# Function to check if conda is available
check_conda() {
    if command -v conda &> /dev/null; then
        return 0
    else
        return 1
    fi
}

# Function to check if a conda environment exists
conda_env_exists() {
    local env_name=$1
    conda info --envs | grep -q "^$env_name "
}

# Function to validate environment variables
validate_env() {
    local missing_vars=()
    
    # Check required environment variables
    [[ -z "$GEMINI_API_KEY" ]] && missing_vars+=("GEMINI_API_KEY")
    [[ -z "$SILICONFLOW_API_KEY" ]] && missing_vars+=("SILICONFLOW_API_KEY")
    [[ -z "$SILICONFLOW_CHAT_URL" ]] && missing_vars+=("SILICONFLOW_CHAT_URL")
    
    if [ ${#missing_vars[@]} -ne 0 ]; then
        print_error "Missing required environment variables:"
        for var in "${missing_vars[@]}"; do
            echo "  - $var"
        done
        echo ""
        echo "Please ensure these variables are set in your .env file."
        return 1
    fi
    
    print_success "All required environment variables are present."
    return 0
}

# Function to check if .env file exists and load it
check_env_file() {
    if [ -f ".env" ]; then
        print_status "Loading environment variables from .env file..."
        export $(grep -v '^#' .env | xargs)
        return 0
    else
        print_warning ".env file not found in current directory."
        echo "Please create a .env file with your API keys."
        echo "See .env.example for reference."
        return 1
    fi
}

# Function to install requirements with pip
install_requirements() {
    print_status "Installing Python requirements..."
    if pip install -r requirements.txt; then
        print_success "Requirements installed successfully."
        return 0
    else
        print_error "Failed to install requirements."
        return 1
    fi
}

# Function to create a new conda environment
create_conda_env() {
    local env_name="wisal-env"
    local python_version="3.10"
    
    print_status "Creating new conda environment: $env_name with Python $python_version..."
    
    if conda create -n "$env_name" python=$python_version -y; then
        print_success "Conda environment '$env_name' created successfully."
        
        # Activate the environment
        eval "$(conda shell.bash hook)"
        conda activate "$env_name"
        
        # Install requirements
        if install_requirements; then
            return 0
        else
            print_error "Failed to install requirements in conda environment."
            return 1
        fi
    else
        print_error "Failed to create conda environment."
        return 1
    fi
}

# Function to use existing conda environment
use_existing_conda_env() {
    local env_name=$1
    
    print_status "Activating existing conda environment: $env_name..."
    
    # Initialize conda for bash shell
    eval "$(conda shell.bash hook)"
    
    if conda activate "$env_name"; then
        print_success "Conda environment '$env_name' activated."
        
        # Install/upgrade requirements
        print_status "Installing/updating Python requirements..."
        if install_requirements; then
            return 0
        else
            print_error "Failed to install requirements."
            return 1
        fi
    else
        print_error "Failed to activate conda environment '$env_name'."
        return 1
    fi
}

# Function to run the application
run_application() {
    print_status "Starting Wisal Autism AI Assistant..."
    
    # Validate environment variables
    if ! validate_env; then
        print_error "Environment validation failed. Cannot start application."
        exit 1
    fi
    
    # Run the main application
    print_success "Launching Wisal application..."
    echo ""
    python main.py
}

# Main script execution
main() {
    echo "==============================================="
    echo "  🚀 Wisal Autism AI Assistant - Launcher"
    echo "==============================================="
    echo ""
    
    # Check if .env file exists
    if ! check_env_file; then
        read -p "Do you want to continue without .env file? (y/N): " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            print_error "Exiting. Please create a .env file with your API keys."
            exit 1
        fi
    fi
    
    echo ""
    echo "Please select an option to run Wisal:"
    echo "  1) Use existing conda environment"
    echo "  2) Create new conda environment"
    echo "  3) Install requirements with current Python environment"
    echo ""
    
    read -p "Enter your choice (1-3): " choice
    
    case $choice in
        1)
            if ! check_conda; then
                print_error "Conda is not installed or not in PATH."
                exit 1
            fi
            
            echo ""
            echo "Available conda environments:"
            conda info --envs | grep -v "#"
            echo ""
            
            read -p "Enter the name of the conda environment to use: " env_name
            
            if [ -z "$env_name" ]; then
                print_error "Environment name cannot be empty."
                exit 1
            fi
            
            if ! conda_env_exists "$env_name"; then
                print_error "Conda environment '$env_name' does not exist."
                exit 1
            fi
            
            if use_existing_conda_env "$env_name"; then
                run_application
            else
                print_error "Failed to set up conda environment."
                exit 1
            fi
            ;;
        2)
            if ! check_conda; then
                print_error "Conda is not installed or not in PATH."
                exit 1
            fi
            
            if create_conda_env; then
                run_application
            else
                print_error "Failed to create conda environment."
                exit 1
            fi
            ;;
        3)
            if install_requirements; then
                run_application
            else
                print_error "Failed to install requirements."
                exit 1
            fi
            ;;
        *)
            print_error "Invalid choice. Please select 1, 2, or 3."
            exit 1
            ;;
    esac
}

# Run main function
main "$@"