#!/bin/bash
# ai.shell setup script for macOS
# Sets up Python venv environment in ~/.config/ai-shell

set -e

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

# Directories
CONFIG_DIR="$HOME/.config/ai-shell"
VENV_DIR="$CONFIG_DIR/venv"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_DIR="$(dirname "$SCRIPT_DIR")"

echo -e "${BLUE}🚀 Setting up ai.shell environment...${NC}"

# Create config directory
echo -e "${YELLOW}📁 Creating config directory...${NC}"
mkdir -p "$CONFIG_DIR"

# Check Python version
echo -e "${YELLOW}🐍 Checking Python version...${NC}"
if ! command -v python3 &> /dev/null; then
    echo -e "${RED}❌ Python 3 is not installed. Please install Python 3.9 or later.${NC}"
    exit 1
fi

PYTHON_VERSION=$(python3 -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")')
echo -e "${GREEN}✓ Found Python $PYTHON_VERSION${NC}"

# Create virtual environment
if [ ! -d "$VENV_DIR" ]; then
    echo -e "${YELLOW}🔧 Creating virtual environment...${NC}"
    python3 -m venv "$VENV_DIR"
    echo -e "${GREEN}✓ Virtual environment created${NC}"
else
    echo -e "${GREEN}✓ Virtual environment already exists${NC}"
fi

# Activate virtual environment
echo -e "${YELLOW}🔌 Activating virtual environment...${NC}"
source "$VENV_DIR/bin/activate"

# Upgrade pip
echo -e "${YELLOW}📦 Upgrading pip...${NC}"
pip install --upgrade pip

# Install Python dependencies
echo -e "${YELLOW}📦 Installing Python dependencies...${NC}"
pip install -r "$PROJECT_DIR/requirements.txt"

# Copy MCP server to config directory
echo -e "${YELLOW}📄 Copying MCP server...${NC}"
cp "$PROJECT_DIR/mcp_server.py" "$CONFIG_DIR/"

# Create config file if it doesn't exist
if [ ! -f "$CONFIG_DIR/config.toml" ]; then
    echo -e "${YELLOW}⚙️  Creating default config...${NC}"
    cp "$PROJECT_DIR/config/default.toml" "$CONFIG_DIR/config.toml"
fi

# Create bin directory for scripts
mkdir -p "$CONFIG_DIR/bin"

# Create MCP server startup script
cat > "$CONFIG_DIR/bin/mcp-server" << 'EOF'
#!/bin/bash
# MCP Server startup script

CONFIG_DIR="$HOME/.config/ai-shell"
VENV_DIR="$CONFIG_DIR/venv"

# Activate virtual environment
source "$VENV_DIR/bin/activate"

# Start MCP server
cd "$CONFIG_DIR"
python mcp_server.py "$@"
EOF

chmod +x "$CONFIG_DIR/bin/mcp-server"

# Create ai-shell wrapper script
cat > "$CONFIG_DIR/bin/ai-shell" << EOF
#!/bin/bash
# ai-shell wrapper script

CONFIG_DIR="\$HOME/.config/ai-shell"
PROJECT_DIR="$PROJECT_DIR"

# Check if MCP server is running
if ! curl -s http://localhost:8765/health > /dev/null 2>&1; then
    echo "Starting MCP server..."
    "\$CONFIG_DIR/bin/mcp-server" > /dev/null 2>&1 &
    
    # Wait for server to start
    for i in {1..10}; do
        if curl -s http://localhost:8765/health > /dev/null 2>&1; then
            break
        fi
        sleep 0.5
    done
fi

# Run ai-shell CLI
cd "\$PROJECT_DIR"
cargo run --release -- "\$@"
EOF

chmod +x "$CONFIG_DIR/bin/ai-shell"

# Check if Ollama is installed
echo -e "${YELLOW}🦙 Checking Ollama...${NC}"
if ! command -v ollama &> /dev/null; then
    echo -e "${YELLOW}⚠️  Ollama is not installed.${NC}"
    echo -e "   Please install Ollama from: https://ollama.ai"
    echo -e "   Then run: ollama pull qwen2.5-coder:7b"
else
    echo -e "${GREEN}✓ Ollama is installed${NC}"
    
    # Check if model is available
    if ! ollama list | grep -q "qwen2.5-coder:7b"; then
        echo -e "${YELLOW}📥 Pulling qwen2.5-coder:7b model...${NC}"
        echo -e "   This may take a while..."
        ollama pull qwen2.5-coder:7b
    else
        echo -e "${GREEN}✓ Model qwen2.5-coder:7b is available${NC}"
    fi
fi

# Build Rust project
echo -e "${YELLOW}🔨 Building Rust CLI...${NC}"
cd "$PROJECT_DIR"
cargo build --release

echo -e "\n${GREEN}✨ Setup complete!${NC}"
echo -e "\n${BLUE}To use ai.shell:${NC}"
echo -e "1. Add to your PATH: ${YELLOW}export PATH=\"\$HOME/.config/ai-shell/bin:\$PATH\"${NC}"
echo -e "2. Run: ${YELLOW}ai-shell${NC}"
echo -e "\nOr use the full path: ${YELLOW}~/.config/ai-shell/bin/ai-shell${NC}"

# Deactivate virtual environment
deactivate