#!/bin/bash
# Healthcare Agent Training with ZETA playground integration
# Uses MCP-ZERO core infrastructure for offline-first resilience

set -e

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

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

# Initialize variables
FORCE_OFFLINE=""
MODEL_SIZE="standard"
EPOCHS=5
VERBOSE=0

# Parse command line arguments
while [[ $# -gt 0 ]]; do
  key="$1"
  case $key in
    --offline)
      FORCE_OFFLINE="--offline"
      shift
      ;;
    --model-size)
      MODEL_SIZE="$2"
      shift 2
      ;;
    --epochs)
      EPOCHS="$2"
      shift 2
      ;;
    --verbose)
      VERBOSE=1
      shift
      ;;
    *)
      echo -e "${RED}Unknown option: $key${NC}"
      echo "Usage: $0 [--offline] [--model-size small|standard|large] [--epochs N] [--verbose]"
      exit 1
      ;;
  esac
done

echo -e "${GREEN}=== Healthcare Agent Training with ZETA Playground ===${NC}"
echo -e "${BLUE}Using MCP-ZERO core infrastructure${NC}"

# Check for MCP-ZERO CLI
MCP_CLI="$MCP_CORE/mcp-cli.sh"
if [[ ! -f "$MCP_CLI" ]]; then
    echo -e "${RED}Error: MCP-ZERO CLI not found at $MCP_CLI${NC}"
    echo -e "${YELLOW}Running in standalone mode without MCP-ZERO core integration${NC}"
    MCP_CLI=""
fi

# Set up training environment through MCP-ZERO if available
if [[ -n "$MCP_CLI" ]]; then
    echo -e "${GREEN}Setting up training environment through MCP-ZERO...${NC}"
    $MCP_CLI setup --agent healthcare --force
    
    # Check offline status
    if [[ -n "$FORCE_OFFLINE" || "$($MCP_CLI status --offline)" == *"enabled"* ]]; then
        echo -e "${YELLOW}Running in offline mode${NC}"
        OFFLINE_STATUS="--offline"
    else
        echo -e "${GREEN}Running in online mode${NC}"
        OFFLINE_STATUS=""
    fi
else
    # Standalone setup
    echo -e "${GREEN}Setting up training environment...${NC}"
    mkdir -p "$SCRIPT_DIR/models/healthcare"
    mkdir -p "$SCRIPT_DIR/logs"
    
    if [[ -n "$FORCE_OFFLINE" ]]; then
        OFFLINE_STATUS="--offline"
        echo -e "${YELLOW}Running in offline mode${NC}"
    else
        OFFLINE_STATUS=""
        echo -e "${GREEN}Running in online mode${NC}"
    fi
fi

# Launch ZETA playground training
echo -e "${GREEN}Starting ZETA playground for healthcare agent training...${NC}"

# Create training configuration
TRAIN_CONFIG="$SCRIPT_DIR/training_config.json"
cat > "$TRAIN_CONFIG" << EOF
{
  "agent": "healthcare",
  "model_size": "$MODEL_SIZE",
  "epochs": $EPOCHS,
  "dataset": "healthcare_samples/medical_queries.json",
  "offline_first_resilience": true,
  "single_connection_attempt": true,
  "output_dir": "models/healthcare",
  "zeta_components": ["llm", "tensor", "face"]
}
EOF

# Launch ZETA playground with training configuration
echo -e "${GREEN}Launching ZETA playground training...${NC}"

ZETA_PLAYGROUND="$MCP_ROOT/zeta/playground.py"
if [[ -f "$ZETA_PLAYGROUND" ]]; then
    python3 "$ZETA_PLAYGROUND" --mode train --config "$TRAIN_CONFIG" $OFFLINE_STATUS
else
    echo -e "${YELLOW}ZETA playground not found at $ZETA_PLAYGROUND${NC}"
    echo -e "${YELLOW}Using alternate training method...${NC}"
    
    # Fall back to using our healthcare agent's built-in training
    echo -e "${GREEN}Training healthcare agent using built-in methods...${NC}"
    python3 "$SCRIPT_DIR/zeta_playground.py" --train --model "$MODEL_SIZE" --epochs "$EPOCHS" $OFFLINE_STATUS
fi

echo -e "${GREEN}Training complete!${NC}"
echo -e "${YELLOW}To start the healthcare agent:${NC} cd $SCRIPT_DIR && ./start_healthcare_agent.sh"
echo -e "${YELLOW}To start the healthcare API:${NC} cd $SCRIPT_DIR && ./start_healthcare_agent.sh --api"
