frontend/src/components/shared/EntityRelationTreeChart.tsx

import React, { useMemo } from "react";
import { useAgentGraph } from "@/context/AgentGraphContext";

interface EntityRelationTreeChartProps {
  className?: string;
}

interface ChartData {
  name: string;
  count: number;
  color: string;
  type: "entity" | "relation";
}

export function EntityRelationTreeChart({
  className,
}: EntityRelationTreeChartProps) {
  const { state } = useAgentGraph();

  // Process data to get entity and relation type distributions
  const chartData = useMemo(() => {
    const entityTypes = new Map<string, number>();
    const relationTypes = new Map<string, number>();

    state.traces.forEach((trace) => {
      trace.knowledge_graphs?.forEach((kg: any) => {
        if (kg.graph_data) {
          try {
            const graphData =
              typeof kg.graph_data === "string"
                ? JSON.parse(kg.graph_data)
                : kg.graph_data;

            // Count entity types
            if (graphData.entities && Array.isArray(graphData.entities)) {
              graphData.entities.forEach((entity: any) => {
                const type = entity.type || "Unknown";
                entityTypes.set(type, (entityTypes.get(type) || 0) + 1);
              });
            }

            // Count relation types
            if (graphData.relations && Array.isArray(graphData.relations)) {
              graphData.relations.forEach((relation: any) => {
                const type = relation.type || "Unknown";
                relationTypes.set(type, (relationTypes.get(type) || 0) + 1);
              });
            }
          } catch (error) {
            console.warn("Error parsing graph_data:", error);
          }
        }
      });
    });

    // Entity type colors
    const entityColors = {
      Agent: "#3b82f6",
      Task: "#10b981",
      Tool: "#f59e0b",
      Input: "#8b5cf6",
      Output: "#ef4444",
      Human: "#06b6d4",
      Unknown: "#6b7280",
    };

    // Relation type colors (using different shades)
    const relationColors = {
      CONSUMED_BY: "#93c5fd",
      PERFORMS: "#86efac",
      ASSIGNED_TO: "#fbbf24",
      USES: "#c4b5fd",
      REQUIRED_BY: "#fca5a5",
      SUBTASK_OF: "#67e8f9",
      NEXT: "#a78bfa",
      PRODUCES: "#34d399",
      DELIVERS_TO: "#60a5fa",
      INTERVENES: "#fb7185",
      Unknown: "#9ca3af",
    };

    const data: ChartData[] = [];

    // Add entity type data
    entityTypes.forEach((count, type) => {
      data.push({
        name: type,
        count,
        color:
          entityColors[type as keyof typeof entityColors] ||
          entityColors.Unknown,
        type: "entity",
      });
    });

    // Add relation type data
    relationTypes.forEach((count, type) => {
      data.push({
        name: type.replace(/_/g, " "), // Make relation names more readable
        count,
        color:
          relationColors[type as keyof typeof relationColors] ||
          relationColors.Unknown,
        type: "relation",
      });
    });

    // Sort by count descending and take top 8 to fit in the space
    return data.sort((a, b) => b.count - a.count).slice(0, 8);
  }, [state.traces]);

  if (chartData.length === 0) {
    return (
      <div
        className={`h-20 w-full flex items-center justify-center ${className}`}
      >
        <div className="text-xs text-muted-foreground/50">No graph data</div>
      </div>
    );
  }

  // Calculate max count for scaling
  const maxCount = Math.max(...chartData.map((d) => d.count));

  return (
    <div
      className={`h-20 w-full flex items-end justify-center gap-1 px-2 ${className}`}
    >
      {chartData.map((item, index) => {
        const height = Math.max(8, (item.count / maxCount) * 60); // Min height 8px, max 60px

        return (
          <div
            key={`${item.type}-${item.name}-${index}`}
            className="flex flex-col items-center justify-end flex-1 min-w-0"
            title={`${item.name}: ${item.count}`}
          >
            {/* Bar */}
            <div
              className="w-full rounded-t-sm transition-all duration-300 hover:opacity-80"
              style={{
                height: `${height}px`,
                backgroundColor: item.color,
                boxShadow: `0 0 4px ${item.color}40`,
              }}
            />

            {/* Label */}
            <div className="text-[8px] text-muted-foreground/70 mt-1 truncate w-full text-center font-medium">
              {item.name.length > 6 ? `${item.name.slice(0, 5)}...` : item.name}
            </div>
          </div>
        );
      })}
    </div>
  );
}