#!/bin/bash
# Helper script to set up CAST.AI MCP server with Claude Desktop
# Usage: ./setup_claude_desktop.sh YOUR_API_KEY

set -e

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

echo "=================================================="
echo "CAST.AI MCP Server - Claude Desktop Setup"
echo "=================================================="
echo ""

# Get API key
if [ -z "$1" ]; then
    echo -e "${YELLOW}No API key provided as argument.${NC}"
    echo ""
    echo "To get your API key:"
    echo "1. Visit https://console.cast.ai"
    echo "2. Go to Settings > API Keys"
    echo "3. Create a new API key (Read-only recommended)"
    echo ""
    read -p "Enter your CASTAI API key: " API_KEY

    if [ -z "$API_KEY" ]; then
        echo -e "${RED}Error: API key is required${NC}"
        exit 1
    fi
else
    API_KEY="$1"
fi

echo ""
echo -e "${BLUE}Step 1: Detecting platform...${NC}"

# Detect platform
if [[ "$OSTYPE" == "darwin"* ]]; then
    PLATFORM="macos"
    CONFIG_DIR="$HOME/Library/Application Support/Claude"
    CONFIG_FILE="$CONFIG_DIR/claude_desktop_config.json"
    LOG_DIR="$HOME/Library/Logs/Claude"
    echo -e "${GREEN}✓ Detected macOS${NC}"
elif [[ "$OSTYPE" == "msys" || "$OSTYPE" == "win32" ]]; then
    PLATFORM="windows"
    CONFIG_DIR="$APPDATA/Claude"
    CONFIG_FILE="$CONFIG_DIR/claude_desktop_config.json"
    echo -e "${GREEN}✓ Detected Windows${NC}"
else
    echo -e "${RED}✗ Unsupported platform: $OSTYPE${NC}"
    exit 1
fi

echo ""
echo -e "${BLUE}Step 2: Checking configuration directory...${NC}"

# Create config directory if it doesn't exist
if [ ! -d "$CONFIG_DIR" ]; then
    echo -e "${YELLOW}Creating config directory: $CONFIG_DIR${NC}"
    mkdir -p "$CONFIG_DIR"
fi
echo -e "${GREEN}✓ Config directory exists${NC}"

echo ""
echo -e "${BLUE}Step 3: Preparing configuration...${NC}"

# Get absolute path to this directory
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
echo "  MCP Server path: $SCRIPT_DIR"

# Create config JSON
CASTAI_CONFIG=$(cat <<EOF
{
  "command": "uv",
  "args": [
    "--directory",
    "$SCRIPT_DIR",
    "run",
    "python",
    "main.py"
  ],
  "env": {
    "CASTAI_API_KEY": "$API_KEY"
  }
}
EOF
)

echo ""
echo -e "${BLUE}Step 4: Updating Claude Desktop configuration...${NC}"

# Backup existing config if it exists
if [ -f "$CONFIG_FILE" ]; then
    BACKUP_FILE="${CONFIG_FILE}.backup.$(date +%Y%m%d_%H%M%S)"
    echo "  Creating backup: $BACKUP_FILE"
    cp "$CONFIG_FILE" "$BACKUP_FILE"

    # Check if config already has mcpServers
    if grep -q '"mcpServers"' "$CONFIG_FILE"; then
        echo "  Found existing MCP servers configuration"

        # Use Python to merge the config
        python3 << PYTHON_SCRIPT
import json
import sys

config_file = "$CONFIG_FILE"

try:
    with open(config_file, 'r') as f:
        config = json.load(f)
except:
    config = {}

if 'mcpServers' not in config:
    config['mcpServers'] = {}

config['mcpServers']['castai'] = json.loads('''$CASTAI_CONFIG''')

with open(config_file, 'w') as f:
    json.dump(config, f, indent=2)

print("  ✓ Added CAST.AI MCP server to existing configuration")
PYTHON_SCRIPT
    else
        # No existing mcpServers, create new config
        cat > "$CONFIG_FILE" << EOF
{
  "mcpServers": {
    "castai": $CASTAI_CONFIG
  }
}
EOF
        echo -e "${GREEN}  ✓ Created new configuration with CAST.AI MCP server${NC}"
    fi
else
    # No config file exists, create new one
    cat > "$CONFIG_FILE" << EOF
{
  "mcpServers": {
    "castai": $CASTAI_CONFIG
  }
}
EOF
    echo -e "${GREEN}  ✓ Created new configuration file${NC}"
fi

echo ""
echo -e "${BLUE}Step 5: Validating configuration...${NC}"

# Validate JSON
if python3 -m json.tool "$CONFIG_FILE" > /dev/null 2>&1; then
    echo -e "${GREEN}✓ Configuration file is valid JSON${NC}"
else
    echo -e "${RED}✗ Configuration file has invalid JSON syntax${NC}"
    echo "  Please check: $CONFIG_FILE"
    exit 1
fi

echo ""
echo -e "${BLUE}Step 6: Testing MCP server...${NC}"

# Test the server
export CASTAI_API_KEY="$API_KEY"
timeout 3s uv run python "$SCRIPT_DIR/main.py" > /tmp/castai_mcp_test.log 2>&1 &
PID=$!
sleep 2

if ps -p $PID > /dev/null 2>&1; then
    kill $PID 2>/dev/null || true
    echo -e "${GREEN}✓ MCP server starts successfully${NC}"
else
    echo -e "${YELLOW}⚠ Server test inconclusive (this may be normal)${NC}"
fi

echo ""
echo "=================================================="
echo -e "${GREEN}Setup Complete!${NC}"
echo "=================================================="
echo ""
echo "Configuration saved to:"
echo "  $CONFIG_FILE"
echo ""
if [ -f "$BACKUP_FILE" ]; then
    echo "Backup of previous config:"
    echo "  $BACKUP_FILE"
    echo ""
fi
echo -e "${YELLOW}Next steps:${NC}"
echo "1. Restart Claude Desktop completely (Cmd+Q on macOS)"
echo "2. Launch Claude Desktop again"
echo "3. Look for the MCP indicator (🔌 icon) in the UI"
echo "4. Try these example queries:"
echo "   - 'Show me all my Kubernetes clusters'"
echo "   - 'What are the details of my production cluster?'"
echo "   - 'What's the optimization score for my cluster?'"
echo ""
if [[ "$PLATFORM" == "macos" ]]; then
    echo -e "${BLUE}To view logs:${NC}"
    echo "  tail -f $LOG_DIR/mcp*.log"
    echo ""
fi
echo -e "${BLUE}To test manually:${NC}"
echo "  cd $SCRIPT_DIR"
echo "  export CASTAI_API_KEY='$API_KEY'"
echo "  uv run python main.py"
echo ""
echo -e "${BLUE}Full documentation:${NC}"
echo "  $SCRIPT_DIR/CLAUDE_DESKTOP_SETUP.md"
echo ""
echo "=================================================="