from typing import Dict, List, Any
from smolagents.tools import Tool

class HealthcareLLMVisualizerTool(Tool):
    name = "healthcare_llm_visualizer"
    description = "Creates interactive visualizations for analyzing LLM applications in Healthcare"
    inputs = {
        'data': {
            'type': 'object',  # Using 'object' for dictionary input
            'description': 'Data for visualization in format: {"items": [{"category": "name", "value": number}]}'
        }
    }
    output_type = "string"  # Changed from 'str' to 'string'

    def prepare_data(self, raw_data: Dict) -> Dict[str, List[Dict[str, Any]]]:
        """Convert raw data into format suitable for visualization"""
        categories = {}
        
        # Process trends
        for item in raw_data.get('trends', []):
            category = item['category']
            categories[category] = categories.get(category, 0) + 1
            
        # Process implementations
        for item in raw_data.get('implementations', []):
            category = item['category']
            categories[category] = categories.get(category, 0) + 1
            
        # Process success cases
        for item in raw_data.get('success_cases', []):
            category = item['category']
            categories[category] = categories.get(category, 0) + 1
        
        # Convert to required format
        return {
            "items": [
                {"category": cat, "value": val}
                for cat, val in categories.items()
            ]
        }

    def forward(self, data: Dict) -> str:
        """Creates a visualization from the provided data"""
        try:
            # Prepare data for visualization
            viz_data = self.prepare_data(data)
            
            # Create React component
            chart_code = """
import React from 'react';
import { BarChart, Bar, XAxis, YAxis, CartesianGrid, Tooltip, Legend, ResponsiveContainer } from 'recharts';

const HealthcareLLMChart = () => {
  const data = DATA_PLACEHOLDER;

  return (
    <div className="w-full max-w-4xl mx-auto p-4">
      <h2 className="text-2xl font-bold mb-4">LLM Applications in Healthcare</h2>
      <div className="h-96">
        <ResponsiveContainer width="100%" height="100%">
          <BarChart data={data.items}>
            <CartesianGrid strokeDasharray="3 3" />
            <XAxis dataKey="category" />
            <YAxis />
            <Tooltip />
            <Legend />
            <Bar dataKey="value" fill="#8884d8" />
          </BarChart>
        </ResponsiveContainer>
      </div>
    </div>
  );
};

export default HealthcareLLMChart;
""".replace('DATA_PLACEHOLDER', str(viz_data))
            
            return chart_code
        except Exception as e:
            return f"Error creating visualization: {str(e)}"