#!/bin/bash
# Healthcare Agent Training Script with MCP-ZERO and ZETA integration
# Implements offline-first resilience pattern

set -e

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

# Get the directory of the script
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
MCP_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"

# Define paths
CONFIG_PATH="$MCP_ROOT/agents/templates/specialized/healthcare.json"
DATASET_PATH="$MCP_ROOT/agents/training/datasets/healthcare_samples"
OUTPUT_DIR="$SCRIPT_DIR/models/healthcare"

# Create necessary directories
echo -e "${GREEN}Setting up training environment...${NC}"
mkdir -p "$OUTPUT_DIR/med_large"
mkdir -p "$OUTPUT_DIR/med_small"

# Check for dataset
if [ ! -f "$DATASET_PATH/medical_queries.json" ]; then
    echo -e "${RED}Error: Training dataset not found at $DATASET_PATH/medical_queries.json${NC}"
    exit 1
fi

# Check for training script
if [ ! -f "$MCP_ROOT/agents/training/scripts/train_agent.py" ]; then
    echo -e "${RED}Error: Training script not found at $MCP_ROOT/agents/training/scripts/train_agent.py${NC}"
    exit 1
fi

# Run training with offline-first resilience
echo -e "${GREEN}Starting healthcare agent training with offline-first resilience...${NC}"
echo -e "${YELLOW}Training for primary model (online if available, will fallback to offline)...${NC}"

# Check if we're forcing offline mode
if [ "$1" == "--offline" ]; then
    OFFLINE_ARG="--offline"
    echo -e "${YELLOW}Forcing offline mode for training...${NC}"
else
    OFFLINE_ARG=""
fi

# Run the training script for primary model
python3 "$MCP_ROOT/agents/training/scripts/train_agent.py" \
    --template "$CONFIG_PATH" \
    --dataset "$DATASET_PATH/medical_queries.json" \
    --output "$OUTPUT_DIR/med_large" \
    $OFFLINE_ARG

# Always ensure we have a fallback model by training in offline mode
echo -e "${GREEN}Training fallback model (always in offline mode)...${NC}"
python3 "$MCP_ROOT/agents/training/scripts/train_agent.py" \
    --template "$CONFIG_PATH" \
    --dataset "$DATASET_PATH/medical_queries.json" \
    --output "$OUTPUT_DIR/med_small" \
    --offline

# Copy model info to the main model directories
echo -e "${GREEN}Finalizing models...${NC}"

# Update model info files
if [ -f "$OUTPUT_DIR/med_large/trained_model_full/model_info.json" ]; then
    cp "$OUTPUT_DIR/med_large/trained_model_full/model_info.json" "$OUTPUT_DIR/med_large/model_info.json"
elif [ -f "$OUTPUT_DIR/med_large/trained_model_lite/model_info.json" ]; then
    cp "$OUTPUT_DIR/med_large/trained_model_lite/model_info.json" "$OUTPUT_DIR/med_large/model_info.json"
fi

cp "$OUTPUT_DIR/med_small/trained_model_lite/model_info.json" "$OUTPUT_DIR/med_small/model_info.json"

# Integrate the agent with MCP-ZERO
echo -e "${GREEN}Integrating healthcare agent with MCP-ZERO...${NC}"

# Create a target integration directory
TARGET_DIR="$MCP_ROOT/integrations/healthcare"
mkdir -p "$TARGET_DIR"

# Create a config file for integration
CONFIG_FILE="$TARGET_DIR/integration_config.json"
cat > "$CONFIG_FILE" << EOF
{
  "name": "healthcare-assistant",
  "version": "1.0.0",
  "offline_resilience": true,
  "integration_point": "mcp-zero"
}
EOF

# Run integration script
python3 "$MCP_ROOT/agents/training/scripts/integrate_agent.py" \
    --agent "$CONFIG_PATH" \
    --model-dir "$OUTPUT_DIR" \
    --target "$TARGET_DIR" \
    --config "$CONFIG_FILE" \
    --name "healthcare-assistant"

echo -e "${GREEN}Healthcare Agent trained and integrated successfully!${NC}"
echo -e "${YELLOW}To run the agent:${NC} cd $SCRIPT_DIR && ./start_healthcare_agent.sh"
echo -e "${YELLOW}To run the API server:${NC} cd $SCRIPT_DIR && ./start_healthcare_agent.sh --api"
