#!/bin/bash
# MCP-ZERO Burn Script Fix
# This script fixes the burn.sh script to ensure proper function definitions

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

echo -e "${BLUE}MCP-ZERO Burn Script Fix${NC}"
echo -e "${BLUE}======================${NC}"

# Path to burn.sh
BURN_SCRIPT="/home/umesh/mcp-zero/burn.sh"

if [ ! -f "$BURN_SCRIPT" ]; then
    echo -e "${RED}Error: burn.sh not found at $BURN_SCRIPT${NC}"
    exit 1
fi

echo -e "${YELLOW}Backing up original burn.sh...${NC}"
cp "$BURN_SCRIPT" "${BURN_SCRIPT}.bak"
echo -e "${GREEN}Backup created at ${BURN_SCRIPT}.bak${NC}"

echo -e "${BLUE}Fixing burn.sh...${NC}"

# Fix 1: Ensure print_message function is defined early and properly
sed -i '20i\
# Print colored message function\
print_message() {\
  local color=$1\
  local message=$2\
  echo -e "${color}${message}${NC}"\
}\
' "$BURN_SCRIPT"

# Fix 2: Fix the source_template function to properly use print_message
sed -i 's/print_message "${YELLOW}" "Warning:/echo -e "${YELLOW}Warning:/g' "$BURN_SCRIPT"
sed -i 's/print_message "${YELLOW}" "Created fallback/echo -e "${YELLOW}Created fallback/g' "$BURN_SCRIPT"

# Fix 3: Add Umeshian Construct integration by adding these lines before the main function
sed -i '/^main() {/i\
# Setup Umeshian Construct integration\
setup_umeshian_construct() {\
  echo -e "${BLUE}Setting up Umeshian Construct integration...${NC}"\
  \
  # Create directories if they dont exist\
  mkdir -p "${TARGET_DIR}/agents/umeshian_construct/contracts/burned"\
  mkdir -p "${TARGET_DIR}/agents/umeshian_construct/scripts"\
  mkdir -p "${TARGET_DIR}/agents/umeshian_construct/skills"\
  mkdir -p "${TARGET_DIR}/agents/umeshian_construct/models"\
  mkdir -p "${TARGET_DIR}/agents/umeshian_construct/data/medical"\
  mkdir -p "${TARGET_DIR}/logs/umeshian"\
  \
  # Copy the medical contract if it exists\
  if [ -f "${SOURCE_DIR}/burn_folder/agents/umeshian_construct/contracts/medical_contract.yaml" ]; then\
    cp "${SOURCE_DIR}/burn_folder/agents/umeshian_construct/contracts/medical_contract.yaml" \\\
       "${TARGET_DIR}/agents/umeshian_construct/contracts/medical_contract.yaml"\
  fi\
  \
  # Copy the base contract if it exists\
  if [ -f "${SOURCE_DIR}/burn_folder/agents/umeshian_construct/contracts/contract.yaml" ]; then\
    cp "${SOURCE_DIR}/burn_folder/agents/umeshian_construct/contracts/contract.yaml" \\\
       "${TARGET_DIR}/agents/umeshian_construct/contracts/contract.yaml"\
  fi\
  \
  # Create symbolic agent startup script\
  cat > "${TARGET_DIR}/start_umeshian_agent.sh" << \x27EOF\x27\
#!/bin/bash\
# Start the Umeshian Construct symbolic agent\
\
set -e\
\
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"\
export MCPZERO_ROOT="${SCRIPT_DIR}"\
\
# Parse arguments\
MODE="query"\
QUERY=""\
DATASET=""\
\
while [[ $# -gt 0 ]]; do\
  case $1 in\
    --mode)\
      MODE="$2"\
      shift 2\
      ;;\
    --query)\
      QUERY="$2"\
      shift 2\
      ;;\
    --dataset)\
      DATASET="$2"\
      shift 2\
      ;;\
    *)\
      echo "Unknown option: $1"\
      exit 1\
      ;;\
  esac\
done\
\
echo "Starting Umeshian Construct symbolic agent in ${MODE} mode..."\
\
# Run the appropriate script based on mode\
case $MODE in\
  query)\
    if [ -z "$QUERY" ]; then\
      echo "Error: No query provided. Use --query \"Your question here\""\
      exit 1\
    fi\
    \
    echo "Processing query: $QUERY"\
    echo "Result: Symbolic agent response would appear here"\
    ;;\
    \
  train)\
    if [ -z "$DATASET" ]; then\
      echo "Error: No dataset provided. Use --dataset path/to/dataset"\
      exit 1\
    fi\
    \
    echo "Training with dataset: $DATASET"\
    ;;\
    \
  status)\
    echo "Checking symbolic agent status..."\
    echo "Status: READY"\
    ;;\
    \
  *)\
    echo "Unknown mode: ${MODE}"\
    echo "Available modes: query, train, status"\
    exit 1\
    ;;\
esac\
EOF\
\
  chmod +x "${TARGET_DIR}/start_umeshian_agent.sh"\
  \
  echo -e "${GREEN}Umeshian Construct integration setup complete.${NC}"\
}\
' "$BURN_SCRIPT"

# Fix 4: Add Umeshian Construct setup to the main function
sed -i '/generate_agent_files/a\\n  # Setup Umeshian Construct if agents are enabled\n  if [ ${INCLUDE_AGENTS} -eq 1 ]; then\n    setup_umeshian_construct\n  fi' "$BURN_SCRIPT"

echo -e "${GREEN}Fixed burn.sh successfully!${NC}"
echo -e "${YELLOW}You can now run burn.sh with:${NC}"
echo -e "${YELLOW}cd /home/umesh/mcp-zero && ./burn.sh --agents --offline${NC}"
