frontend/src/components/features/dashboard/visualizations/RelationFlowAnalysis.tsx

import React, { useMemo } from "react";
import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
import { Badge } from "@/components/ui/badge";
import { useAgentGraph } from "@/context/AgentGraphContext";
import { GitBranch, ArrowRight, Network, TrendingUp } from "lucide-react";

interface RelationStats {
  [relationType: string]: {
    count: number;
    percentage: number;
    sourceTypes: Set<string>;
    targetTypes: Set<string>;
  };
}

interface FlowPattern {
  source: string;
  target: string;
  relation: string;
  count: number;
  percentage: number;
}

interface RelationFlowData {
  totalRelations: number;
  relationStats: RelationStats;
  mostCommonRelation: string;
  topFlowPatterns: FlowPattern[];
  networkComplexity: number;
}

export function RelationFlowAnalysis() {
  const { state } = useAgentGraph();

  const flowData = useMemo<RelationFlowData>(() => {
    const relationStats: RelationStats = {};
    const flowPatterns: { [key: string]: FlowPattern } = {};
    let totalRelations = 0;
    const allEntityTypes = new Set<string>();

    // Process all traces and their knowledge graphs
    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;

            // First pass: collect all entity types
            if (graphData.entities && Array.isArray(graphData.entities)) {
              graphData.entities.forEach((entity: any) => {
                allEntityTypes.add(entity.type || "Unknown");
              });
            }

            if (graphData.relations && Array.isArray(graphData.relations)) {
              // Create entity lookup map
              const entityLookup = new Map<string, string>();
              if (graphData.entities && Array.isArray(graphData.entities)) {
                graphData.entities.forEach((entity: any) => {
                  entityLookup.set(entity.id, entity.type || "Unknown");
                });
              }

              graphData.relations.forEach((relation: any) => {
                const type = relation.type || "Unknown";
                const sourceType =
                  entityLookup.get(relation.source) || "Unknown";
                const targetType =
                  entityLookup.get(relation.target) || "Unknown";

                // Track relation statistics
                if (!relationStats[type]) {
                  relationStats[type] = {
                    count: 0,
                    percentage: 0,
                    sourceTypes: new Set(),
                    targetTypes: new Set(),
                  };
                }

                relationStats[type].count++;
                relationStats[type].sourceTypes.add(sourceType);
                relationStats[type].targetTypes.add(targetType);
                totalRelations++;

                // Track flow patterns
                const flowKey = `${sourceType}_${type}_${targetType}`;
                if (!flowPatterns[flowKey]) {
                  flowPatterns[flowKey] = {
                    source: sourceType,
                    target: targetType,
                    relation: type,
                    count: 0,
                    percentage: 0,
                  };
                }
                flowPatterns[flowKey].count++;
              });
            }
          } catch (error) {
            console.warn("Error parsing graph_data:", error);
          }
        }
      });
    });

    // Calculate percentages
    Object.values(relationStats).forEach((stats) => {
      stats.percentage =
        totalRelations > 0
          ? Math.round((stats.count / totalRelations) * 100)
          : 0;
    });

    Object.values(flowPatterns).forEach((pattern) => {
      pattern.percentage =
        totalRelations > 0
          ? Math.round((pattern.count / totalRelations) * 100)
          : 0;
    });

    // Get top flow patterns
    const topFlowPatterns = Object.values(flowPatterns)
      .sort((a, b) => b.count - a.count)
      .slice(0, 8);

    const mostCommonRelation =
      Object.entries(relationStats).sort(
        ([, a], [, b]) => b.count - a.count
      )[0]?.[0] || "None";

    const networkComplexity =
      allEntityTypes.size > 0
        ? Math.round((totalRelations / allEntityTypes.size) * 10) / 10
        : 0;

    return {
      totalRelations,
      relationStats,
      mostCommonRelation,
      topFlowPatterns,
      networkComplexity,
    };
  }, [state.traces]);

  const getRelationColor = (index: number) => {
    const colors = [
      "bg-emerald-500",
      "bg-cyan-500",
      "bg-violet-500",
      "bg-rose-500",
      "bg-amber-500",
      "bg-indigo-500",
    ];
    return colors[index % colors.length];
  };

  const getImportanceIcon = (percentage: number) => {
    if (percentage > 20)
      return <TrendingUp className="h-4 w-4 text-green-600" />;
    if (percentage > 10)
      return <ArrowRight className="h-4 w-4 text-yellow-600" />;
    return <GitBranch className="h-4 w-4 text-blue-600" />;
  };

  if (flowData.totalRelations === 0) {
    return (
      <Card>
        <CardHeader>
          <CardTitle className="flex items-center gap-2">
            <Network className="h-5 w-5" />
            Relation Flow Analysis
          </CardTitle>
        </CardHeader>
        <CardContent>
          <div className="text-center py-8 text-muted-foreground">
            <Network className="h-12 w-12 mx-auto mb-4 opacity-50" />
            <p>No relation data available</p>
            <p className="text-sm">
              Generate some knowledge graphs to see relationship analysis
            </p>
          </div>
        </CardContent>
      </Card>
    );
  }

  return (
    <div className="space-y-6">
      {/* Summary Stats */}
      <div className="grid grid-cols-1 md:grid-cols-2 lg:grid-cols-3 gap-4">
        <Card>
          <CardContent className="p-4">
            <div className="flex items-center gap-3">
              <div className="p-2 rounded-lg bg-emerald-100">
                <GitBranch className="h-4 w-4 text-emerald-600" />
              </div>
              <div>
                <p className="text-sm text-muted-foreground">Total Relations</p>
                <p className="text-2xl font-bold">{flowData.totalRelations}</p>
              </div>
            </div>
          </CardContent>
        </Card>

        <Card>
          <CardContent className="p-4">
            <div className="flex items-center gap-3">
              <div className="p-2 rounded-lg bg-cyan-100">
                <Network className="h-4 w-4 text-cyan-600" />
              </div>
              <div>
                <p className="text-sm text-muted-foreground">Relation Types</p>
                <p className="text-2xl font-bold">
                  {Object.keys(flowData.relationStats).length}
                </p>
              </div>
            </div>
          </CardContent>
        </Card>

        <Card>
          <CardContent className="p-4">
            <div className="flex items-center gap-3">
              <div className="p-2 rounded-lg bg-violet-100">
                <TrendingUp className="h-4 w-4 text-violet-600" />
              </div>
              <div>
                <p className="text-sm text-muted-foreground">
                  Network Complexity
                </p>
                <p className="text-2xl font-bold">
                  {flowData.networkComplexity}
                </p>
              </div>
            </div>
          </CardContent>
        </Card>
      </div>

      {/* Relation Types Distribution */}
      <Card>
        <CardHeader>
          <CardTitle className="flex items-center gap-2">
            <GitBranch className="h-5 w-5" />
            Relation Types Distribution
          </CardTitle>
        </CardHeader>
        <CardContent>
          <div className="space-y-3">
            {Object.entries(flowData.relationStats)
              .sort(([, a], [, b]) => b.count - a.count)
              .map(([type, stats], index) => (
                <div
                  key={type}
                  className="flex items-center justify-between p-3 bg-muted/30 rounded-lg"
                >
                  <div className="flex items-center gap-3">
                    <div
                      className={`w-4 h-4 rounded-full ${getRelationColor(
                        index
                      )}`}
                    />
                    <div>
                      <span className="font-medium">
                        {type.replace(/_/g, " ")}
                      </span>
                      <div className="text-xs text-muted-foreground mt-1">
                        {stats.sourceTypes.size} → {stats.targetTypes.size}{" "}
                        entity types
                      </div>
                    </div>
                  </div>
                  <div className="flex items-center gap-3">
                    {getImportanceIcon(stats.percentage)}
                    <Badge variant="secondary">{stats.count} relations</Badge>
                    <span className="text-sm text-muted-foreground w-12 text-right">
                      {stats.percentage}%
                    </span>
                  </div>
                </div>
              ))}
          </div>
        </CardContent>
      </Card>

      {/* Top Flow Patterns */}
      <Card>
        <CardHeader>
          <CardTitle className="flex items-center gap-2">
            <ArrowRight className="h-5 w-5" />
            Top Flow Patterns
          </CardTitle>
        </CardHeader>
        <CardContent>
          <div className="space-y-4">
            {flowData.topFlowPatterns.map((pattern, index) => (
              <div
                key={`${pattern.source}_${pattern.relation}_${pattern.target}`}
                className="flex items-center justify-between p-4 border rounded-lg hover:bg-muted/20 transition-colors"
              >
                <div className="flex items-center gap-4">
                  <div
                    className={`w-3 h-3 rounded-full ${getRelationColor(
                      index
                    )}`}
                  />
                  <div className="flex items-center gap-2">
                    <Badge variant="outline" className="text-xs">
                      {pattern.source}
                    </Badge>
                    <ArrowRight className="h-3 w-3 text-muted-foreground" />
                    <span className="text-sm font-medium px-2 py-1 bg-muted/50 rounded">
                      {pattern.relation.replace(/_/g, " ")}
                    </span>
                    <ArrowRight className="h-3 w-3 text-muted-foreground" />
                    <Badge variant="outline" className="text-xs">
                      {pattern.target}
                    </Badge>
                  </div>
                </div>
                <div className="flex items-center gap-3">
                  <Badge variant="secondary">{pattern.count} occurrences</Badge>
                  <span className="text-sm text-muted-foreground w-12 text-right">
                    {pattern.percentage}%
                  </span>
                </div>
              </div>
            ))}
          </div>
        </CardContent>
      </Card>
    </div>
  );
}