#!/bin/bash
# MCP-ZERO Installation and Burn Script
# This script installs MCP-ZERO and creates a burn folder outside the project directory

set -e

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

# Constants
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
DEFAULT_INSTALL_DIR="$HOME/mcp-zero"
DEFAULT_BURN_DIR="$HOME/mcp-zero-deploy"

# Print colored message
print_message() {
  local color=$1
  local message=$2
  echo -e "${color}${message}${NC}"
}

# Print usage information
print_usage() {
  echo "Usage: $0 [OPTIONS]"
  echo
  echo "Options:"
  echo "  -h, --help                   Show this help message"
  echo "  --install-dir DIR            MCP-ZERO installation directory (default: $DEFAULT_INSTALL_DIR)"
  echo "  --burn-dir DIR               Burn folder destination (default: $DEFAULT_BURN_DIR)"
  echo "  --offline                    Enable offline mode in the burn folder"
  echo "  --skip-install               Skip MCP-ZERO installation (only burn)"
  echo "  --force                      Overwrite existing directories"
  echo
  echo "Example:"
  echo "  $0 --install-dir ~/my-mcp --burn-dir ~/my-deploy --offline --force"
  exit 1
}

# Parse command line arguments
parse_arguments() {
  INSTALL_DIR="$DEFAULT_INSTALL_DIR"
  BURN_DIR="$DEFAULT_BURN_DIR"
  OFFLINE_MODE=0
  SKIP_INSTALL=0
  FORCE=0
  
  while [[ $# -gt 0 ]]; do
    case $1 in
      -h|--help)
        print_usage
        ;;
      --install-dir)
        INSTALL_DIR="$2"
        shift 2
        ;;
      --burn-dir)
        BURN_DIR="$2"
        shift 2
        ;;
      --offline)
        OFFLINE_MODE=1
        shift
        ;;
      --skip-install)
        SKIP_INSTALL=1
        shift
        ;;
      --force)
        FORCE=1
        shift
        ;;
      *)
        print_message "$RED" "Error: Unknown option $1"
        print_usage
        ;;
    esac
  done
}

# Check for required dependencies
check_dependencies() {
  print_message "$BLUE" "Checking dependencies..."
  
  MISSING_DEPS=0
  
  # Check for essential tools
  for cmd in git python3 pip3 jq; do
    if ! command -v ${cmd} &> /dev/null; then
      print_message "$RED" "Error: Required tool '${cmd}' is not installed."
      MISSING_DEPS=1
    fi
  done
  
  if [ ${MISSING_DEPS} -eq 1 ]; then
    print_message "$RED" "Please install missing dependencies and try again."
    exit 1
  fi
  
  print_message "$GREEN" "All dependencies satisfied."
}

# Install MCP-ZERO
install_mcp_zero() {
  if [ $SKIP_INSTALL -eq 1 ]; then
    print_message "$YELLOW" "Skipping MCP-ZERO installation as requested."
    return 0
  fi
  
  print_message "$BLUE" "Installing MCP-ZERO to $INSTALL_DIR..."
  
  # Check if directory exists and is not empty
  if [ -d "$INSTALL_DIR" ] && [ "$(ls -A "$INSTALL_DIR" 2>/dev/null)" ]; then
    if [ $FORCE -eq 1 ]; then
      print_message "$YELLOW" "Removing existing MCP-ZERO installation..."
      rm -rf "$INSTALL_DIR"
    else
      print_message "$RED" "Error: $INSTALL_DIR already exists and is not empty."
      print_message "$YELLOW" "Use --force to overwrite or --skip-install to skip installation."
      exit 1
    fi
  fi
  
  # Clone the repository
  print_message "$BLUE" "Cloning MCP-ZERO repository..."
  mkdir -p "$INSTALL_DIR"
  
  # Using current directory as source instead of cloning from GitHub
  print_message "$YELLOW" "Using local MCP-ZERO as source..."
  SOURCE_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"
  
  # Copy files from source to install directory
  rsync -a --exclude '.git' --exclude 'node_modules' "$SOURCE_DIR/" "$INSTALL_DIR/"
  
  # Install Python dependencies
  print_message "$BLUE" "Installing Python dependencies..."
  if [ -f "$INSTALL_DIR/requirements.txt" ]; then
    pip3 install -r "$INSTALL_DIR/requirements.txt"
  else
    print_message "$YELLOW" "No requirements.txt found, skipping Python dependencies."
  fi
  
  # Set up environment
  print_message "$BLUE" "Setting up environment..."
  export MCPZERO_ROOT="$INSTALL_DIR"
  
  # Run initialization script if exists
  if [ -f "$INSTALL_DIR/init.sh" ]; then
    print_message "$BLUE" "Running initialization script..."
    chmod +x "$INSTALL_DIR/init.sh"
    (cd "$INSTALL_DIR" && ./init.sh)
  fi
  
  print_message "$GREEN" "MCP-ZERO installation complete!"
}

# Create burn folder
create_burn_folder() {
  print_message "$BLUE" "Creating burn folder at $BURN_DIR..."
  
  # Check if directory exists and is not empty
  if [ -d "$BURN_DIR" ] && [ "$(ls -A "$BURN_DIR" 2>/dev/null)" ]; then
    if [ $FORCE -eq 1 ]; then
      print_message "$YELLOW" "Removing existing burn folder..."
      rm -rf "$BURN_DIR"
    else
      print_message "$RED" "Error: $BURN_DIR already exists and is not empty."
      print_message "$YELLOW" "Use --force to overwrite."
      exit 1
    fi
  fi
  
  # Create burn directory
  mkdir -p "$BURN_DIR"
  
  # Use the burn script from the installation
  if [ -f "$INSTALL_DIR/burn.sh" ]; then
    print_message "$BLUE" "Running burn script from installation..."
    chmod +x "$INSTALL_DIR/burn.sh"
    
    # Build burn arguments
    BURN_ARGS=""
    if [ $OFFLINE_MODE -eq 1 ]; then
      BURN_ARGS="$BURN_ARGS --offline"
    fi
    
    # Add flags for Docker and agents
    BURN_ARGS="$BURN_ARGS --docker --agents"
    
    # Run the burn script
    (cd "$INSTALL_DIR" && ./burn.sh $BURN_ARGS "$BURN_DIR")
  else
    # Fallback to manual burn if burn.sh doesn't exist
    print_message "$YELLOW" "No burn.sh found, performing manual burn process..."
    
    # Create directory structure
    mkdir -p "$BURN_DIR/.mcpzero/playground"
    mkdir -p "$BURN_DIR/.mcpzero/locks"
    mkdir -p "$BURN_DIR/agents/healthcare"
    mkdir -p "$BURN_DIR/logs"
    
    # Copy core components
    rsync -a "$INSTALL_DIR/.mcpzero/playground/" "$BURN_DIR/.mcpzero/playground/"
    
    # Copy healthcare agent
    rsync -a "$INSTALL_DIR/agents/healthcare/" "$BURN_DIR/agents/healthcare/"
    
    # Copy integration scripts
    cp "$SCRIPT_DIR/burn.sh" "$BURN_DIR/"
    chmod +x "$BURN_DIR/burn.sh"
    
    # Copy configuration
    cp "$SCRIPT_DIR/burn_config.json" "$BURN_DIR/"
    
    # Create offline lock if needed
    if [ $OFFLINE_MODE -eq 1 ]; then
      LOCK_FILE="$BURN_DIR/.mcpzero/locks/offline.lock"
      echo "{\"timestamp\": $(date +%s), \"reason\": \"Created during installation\", \"immutable\": true}" > "$LOCK_FILE"
      chmod 444 "$LOCK_FILE"
      print_message "$YELLOW" "Created immutable offline lock file."
    fi
    
    # Run the burn script in the burn folder to finalize setup
    (cd "$BURN_DIR" && ./burn.sh)
  fi
  
  print_message "$GREEN" "Burn folder created successfully!"
}

# Configure healthcare agent
configure_healthcare() {
  print_message "$BLUE" "Configuring healthcare agent in burn folder..."
  
  # Ensure playground_adapter.py is properly linked
  ADAPTER_FILE="$BURN_DIR/agents/healthcare/playground_adapter.py"
  if [ ! -f "$ADAPTER_FILE" ]; then
    if [ -f "$SCRIPT_DIR/agents/healthcare/playground_adapter.py" ]; then
      cp "$SCRIPT_DIR/agents/healthcare/playground_adapter.py" "$ADAPTER_FILE"
      chmod +x "$ADAPTER_FILE"
      print_message "$GREEN" "Copied playground adapter to burn folder."
    else
      print_message "$RED" "Error: Healthcare playground adapter not found."
      exit 1
    fi
  fi
  
  # Link core playground to healthcare agent
  PLAYGROUND_FILE="$BURN_DIR/.mcpzero/playground/zeta_playground.py"
  if [ -f "$PLAYGROUND_FILE" ]; then
    ln -sf "$PLAYGROUND_FILE" "$BURN_DIR/agents/healthcare/zeta_playground.py"
    print_message "$GREEN" "Linked core playground to healthcare agent."
  else
    print_message "$RED" "Error: Core playground not found in burn folder."
    exit 1
  fi
  
  print_message "$GREEN" "Healthcare agent configured successfully!"
}

# Create sample training data for healthcare agent
create_sample_data() {
  print_message "$BLUE" "Creating sample training data for healthcare agent..."
  
  # Create sample data directory
  SAMPLE_DATA_DIR="$BURN_DIR/agents/healthcare/data/sample"
  mkdir -p "$SAMPLE_DATA_DIR"
  
  # Create simple sample training data
  cat > "$SAMPLE_DATA_DIR/medical_queries.jsonl" << 'EOF'
{"query": "What are the symptoms of diabetes?", "response": "Common symptoms of diabetes include increased thirst, frequent urination, extreme hunger, unexplained weight loss, fatigue, blurred vision, and slow-healing sores."}
{"query": "How to lower blood pressure naturally?", "response": "Natural ways to lower blood pressure include reducing sodium intake, regular physical activity, maintaining a healthy weight, limiting alcohol, eating a heart-healthy diet rich in fruits and vegetables, and managing stress."}
{"query": "What causes migraine headaches?", "response": "Migraines may be caused by changes in the brainstem and its interactions with the trigeminal nerve, imbalances in brain chemicals, genetic factors, hormonal changes, certain foods, stress, sensory stimuli, and changes in sleep patterns."}
{"query": "How to treat a common cold?", "response": "Common cold treatments include rest, staying hydrated, using over-the-counter pain relievers, decongestants, gargling with salt water, and using humidifiers. Most colds resolve within 7-10 days."}
{"query": "When should I get a flu shot?", "response": "The best time to get a flu shot is before flu season begins, typically in early fall (September-October). Annual vaccination is recommended as the flu virus changes, and protection decreases over time."}
EOF
  
  print_message "$GREEN" "Sample training data created!"
}

# Train healthcare agent
train_healthcare_agent() {
  print_message "$BLUE" "Training healthcare agent..."
  
  # Ensure training script is executable
  TRAIN_SCRIPT="$BURN_DIR/start_healthcare_agent.sh"
  if [ -f "$TRAIN_SCRIPT" ]; then
    chmod +x "$TRAIN_SCRIPT"
    
    # Create models directory if it doesn't exist
    mkdir -p "$BURN_DIR/agents/healthcare/models/trained"
    
    # Run training with sample data
    print_message "$YELLOW" "Starting training process (this may take a while)..."
    (cd "$BURN_DIR" && ./start_healthcare_agent.sh --mode train --dataset "$BURN_DIR/agents/healthcare/data/sample")
    
    print_message "$GREEN" "Healthcare agent training complete!"
  else
    print_message "$RED" "Error: Healthcare training script not found in burn folder."
    exit 1
  fi
}

# Test healthcare agent
test_healthcare_agent() {
  print_message "$BLUE" "Testing healthcare agent..."
  
  # Run a test query
  TEST_SCRIPT="$BURN_DIR/start_healthcare_agent.sh"
  if [ -f "$TEST_SCRIPT" ]; then
    chmod +x "$TEST_SCRIPT"
    
    print_message "$YELLOW" "Running a test query..."
    (cd "$BURN_DIR" && ./start_healthcare_agent.sh --mode query --query "What are the symptoms of the flu?")
    
    print_message "$GREEN" "Test complete! Check the output above for the response."
  else
    print_message "$RED" "Error: Healthcare test script not found in burn folder."
    exit 1
  fi
}

# Main execution
main() {
  print_message "$BLUE" "MCP-ZERO Installation and Burn Script"
  print_message "$BLUE" "======================================"
  
  # Parse command line arguments
  parse_arguments "$@"
  
  # Check dependencies
  check_dependencies
  
  # Install MCP-ZERO
  install_mcp_zero
  
  # Create burn folder
  create_burn_folder
  
  # Configure healthcare agent
  configure_healthcare
  
  # Create sample data
  create_sample_data
  
  # Train healthcare agent
  train_healthcare_agent
  
  # Test healthcare agent
  test_healthcare_agent
  
  print_message "$GREEN" "Installation and burn process completed successfully!"
  print_message "$YELLOW" "MCP-ZERO installation: $INSTALL_DIR"
  print_message "$YELLOW" "Burn folder: $BURN_DIR"
  print_message "$YELLOW" "You can now use the healthcare agent with:"
  print_message "$YELLOW" "cd $BURN_DIR && ./start_healthcare_agent.sh --mode query --query \"Your medical question\""
}

main "$@"
