#!/bin/bash
# Healthcare Agent Runner with MCP-ZERO and ZETA playground 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="$SCRIPT_DIR/../templates/specialized/healthcare.json"
MODEL_DIR="$SCRIPT_DIR/models"
DATA_DIR="$SCRIPT_DIR/data"
LOGS_DIR="$SCRIPT_DIR/logs"

# Create necessary directories
echo -e "${GREEN}Setting up Healthcare Agent directories...${NC}"
mkdir -p "$MODEL_DIR/healthcare/med_large"
mkdir -p "$MODEL_DIR/healthcare/med_small"
mkdir -p "$DATA_DIR"
mkdir -p "$LOGS_DIR"

# Create medical reference database
echo -e "${GREEN}Creating medical reference databases...${NC}"
sqlite3 "$DATA_DIR/medical_reference.db" << EOF
CREATE TABLE IF NOT EXISTS medical_references (
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  keyword TEXT NOT NULL,
  information TEXT NOT NULL,
  source TEXT,
  last_updated TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);

-- Add some sample medical references
INSERT OR IGNORE INTO medical_references (keyword, information, source) VALUES 
  ('diabetes', 'Type 2 diabetes is characterized by insulin resistance and relative insulin deficiency. Treatment includes lifestyle modifications and medications.', 'National Diabetes Association'),
  ('hypertension', 'High blood pressure often has no symptoms but can lead to heart disease and stroke. Lifestyle changes and medications can help manage it.', 'American Heart Association'),
  ('covid-19', 'COVID-19 is caused by the SARS-CoV-2 virus and primarily affects the respiratory system. Symptoms include fever, cough, and difficulty breathing.', 'World Health Organization'),
  ('influenza', 'Seasonal flu is a contagious respiratory illness that can cause mild to severe illness. Annual vaccination is the best way to prevent infection.', 'CDC'),
  ('asthma', 'Asthma is a chronic condition affecting the airways, characterized by recurrent episodes of wheezing, breathlessness, and coughing.', 'National Asthma Education Program');
EOF

# Create a fallback database for offline mode
echo -e "${GREEN}Creating fallback database for offline mode...${NC}"
cp "$DATA_DIR/medical_reference.db" "$DATA_DIR/medical_cache.db"

# Create essential medical data JSON for extreme fallback
cat > "$DATA_DIR/essential_medical_data.json" << EOF
{
  "emergency_protocols": {
    "cardiac_arrest": "Begin CPR immediately. Push hard and fast in the center of the chest at a rate of 100-120 compressions per minute.",
    "stroke": "Remember FAST: Face drooping, Arm weakness, Speech difficulty, Time to call emergency services.",
    "severe_bleeding": "Apply direct pressure to the wound with a clean cloth or bandage. If bleeding doesn't stop, apply more pressure.",
    "allergic_reaction": "For severe reactions with difficulty breathing, use epinephrine auto-injector if available and seek immediate medical help.",
    "choking": "Perform abdominal thrusts (Heimlich maneuver) until the object is expelled or the person becomes unconscious."
  },
  "common_conditions": {
    "fever": "Rest and drink plenty of fluids. Consider acetaminophen or ibuprofen for temperature over 101°F (38.3°C).",
    "headache": "Rest in a quiet, dark room. Try over-the-counter pain relievers. Stay hydrated.",
    "cough": "Stay hydrated and use honey (if over 1 year old) to soothe throat. Avoid smoke and other irritants.",
    "nausea": "Eat small, bland meals. Avoid fatty or spicy foods. Stay hydrated with small sips of clear fluids.",
    "diarrhea": "Stay hydrated with water, clear broths, and electrolyte solutions. Avoid dairy, caffeine, and fatty foods."
  },
  "medication_dosing": {
    "acetaminophen": "Adults: 325-650mg every 4-6 hours as needed, not exceeding 3000mg per day.",
    "ibuprofen": "Adults: 200-400mg every 4-6 hours as needed, not exceeding 1200mg per day unless directed by doctor.",
    "aspirin": "Adults: 325-650mg every 4-6 hours as needed, not exceeding 4000mg per day."
  },
  "disclaimers": [
    "This information is for educational purposes only and not a substitute for professional medical advice.",
    "Always consult with a qualified healthcare provider for medical conditions.",
    "If experiencing a medical emergency, call emergency services immediately."
  ]
}
EOF

# Create simulation model files
echo -e "${GREEN}Creating simulation model files...${NC}"
cat > "$MODEL_DIR/healthcare/med_large/model_info.json" << EOF
{
  "name": "healthcare-large",
  "version": "1.0.0",
  "type": "transformer",
  "size": "large",
  "parameters": "7B",
  "specialization": "medical",
  "requires_cuda": true,
  "quantization": "4-bit",
  "offline_capable": false
}
EOF

cat > "$MODEL_DIR/healthcare/med_small/model_info.json" << EOF
{
  "name": "healthcare-small",
  "version": "1.0.0",
  "type": "transformer",
  "size": "small",
  "parameters": "1.3B",
  "specialization": "medical",
  "requires_cuda": false,
  "quantization": "8-bit",
  "offline_capable": true
}
EOF

# Create ZETA playground integration file
echo -e "${GREEN}Creating ZETA playground integration...${NC}"
cat > "$SCRIPT_DIR/zeta_playground.py" << EOF
#!/usr/bin/env python3
# ZETA Playground Integration for Healthcare Agent

import os
import sys
import json
import time
from pathlib import Path

# Add the src directory to the Python path
src_dir = Path(__file__).parent / "src"
sys.path.append(str(src_dir))

from healthcare_agent import HealthcareAgent, logger

class ZetaPlayground:
    """
    ZETA Playground integration for Healthcare Agent
    Implements offline-first resilience pattern
    """
    
    def __init__(self, config_path=None):
        """Initialize the ZETA playground with Healthcare Agent"""
        if config_path is None:
            base_dir = Path(__file__).parent
            config_path = str(base_dir.parent / "templates" / "specialized" / "healthcare.json")
        
        logger.info(f"Initializing ZETA playground with config: {config_path}")
        self.agent = HealthcareAgent(config_path)
        self.history = []
    
    def process_query(self, query, patient_id=None, context=None):
        """Process a medical query through the agent with offline-first resilience"""
        response = self.agent.process_medical_query(query, patient_id, context)
        
        # Add to history
        self.history.append({
            "timestamp": time.time(),
            "query": query,
            "patient_id": patient_id,
            "response": response
        })
        
        return response
    
    def get_history(self):
        """Get interaction history"""
        return self.history
    
    def get_agent_status(self):
        """Get agent status including offline mode"""
        return {
            "agent_name": self.agent.config['agent']['name'],
            "offline_mode": self.agent.offline_manager.is_in_offline_mode(),
            "model_path": self.agent.zeta_components['llm']['model_path'] if 'llm' in self.agent.zeta_components else None,
            "zeta_components": list(self.agent.zeta_components.keys()),
            "history_count": len(self.history)
        }

# Web API for ZETA playground integration
def create_api():
    """Create a simple Flask API for ZETA playground integration"""
    try:
        from flask import Flask, request, jsonify
        
        app = Flask("ZETA Healthcare Playground")
        playground = ZetaPlayground()
        
        @app.route('/api/health', methods=['GET'])
        def health_check():
            return jsonify({"status": "ok", "agent": playground.get_agent_status()})
        
        @app.route('/api/query', methods=['POST'])
        def process_query():
            data = request.json
            query = data.get('query')
            patient_id = data.get('patient_id')
            context = data.get('context')
            
            if not query:
                return jsonify({"error": "Query is required"}), 400
            
            response = playground.process_query(query, patient_id, context)
            return jsonify(response)
        
        @app.route('/api/history', methods=['GET'])
        def get_history():
            return jsonify(playground.get_history())
        
        return app
    except ImportError:
        logger.error("Flask not installed. Cannot create web API.")
        return None

# Command-line interface
def main():
    import argparse
    
    parser = argparse.ArgumentParser(description="ZETA Healthcare Playground")
    parser.add_argument('--api', action='store_true', help='Start the web API')
    parser.add_argument('--port', type=int, default=8765, help='Port for the web API')
    parser.add_argument('--query', help='Medical query to process')
    parser.add_argument('--patient', help='Patient ID')
    
    args = parser.parse_args()
    
    if args.api:
        app = create_api()
        if app:
            app.run(host='0.0.0.0', port=args.port)
        else:
            print("Error: Could not create web API. Flask may not be installed.")
            return 1
    elif args.query:
        playground = ZetaPlayground()
        response = playground.process_query(args.query, args.patient)
        print(json.dumps(response, indent=2))
    else:
        # Interactive mode
        playground = ZetaPlayground()
        status = playground.get_agent_status()
        
        print(f"ZETA Healthcare Playground")
        print(f"Agent: {status['agent_name']} (Offline mode: {'enabled' if status['offline_mode'] else 'disabled'})")
        print(f"Model: {status['model_path']}")
        print(f"ZETA components: {', '.join(status['zeta_components'])}")
        print("Enter 'exit' to quit")
        
        while True:
            query = input("\nMedical query> ")
            if query.lower() == 'exit':
                break
            
            print("Processing query...")
            response = playground.process_query(query)
            print(json.dumps(response, indent=2))
    
    return 0

if __name__ == "__main__":
    sys.exit(main())
EOF

# Make files executable
chmod +x "$SCRIPT_DIR/src/healthcare_agent.py" 
chmod +x "$SCRIPT_DIR/zeta_playground.py"

# Create a convenience wrapper script
cat > "$SCRIPT_DIR/start_healthcare_agent.sh" << EOF
#!/bin/bash
# Start the Healthcare Agent with ZETA playground integration

cd "\$(dirname "\${BASH_SOURCE[0]}")"
export PYTHONPATH="\$PYTHONPATH:./src"

# If Flask is not installed, prompt to install it
if ! python3 -c "import flask" 2>/dev/null; then
    read -p "Flask is not installed. Install it for web API support? (y/n): " answer
    if [[ "\$answer" == "y" ]]; then
        pip install flask
    fi
fi

# Start the agent
if [[ "\$1" == "--api" ]]; then
    echo "Starting Healthcare Agent API on port 8765..."
    python3 ./zeta_playground.py --api --port 8765
else
    echo "Starting Healthcare Agent in interactive mode..."
    python3 ./zeta_playground.py
fi
EOF

chmod +x "$SCRIPT_DIR/start_healthcare_agent.sh"

echo -e "${GREEN}Healthcare Agent with MCP-ZERO and ZETA playground integration is ready!${NC}"
echo -e "${YELLOW}To start in interactive mode:${NC} cd $SCRIPT_DIR && ./start_healthcare_agent.sh"
echo -e "${YELLOW}To start the API server:${NC} cd $SCRIPT_DIR && ./start_healthcare_agent.sh --api"
