frontend/src/hooks/useDashboardVisualizationData.ts

import { useState, useEffect, useMemo } from "react";
import { useAgentGraph } from "@/context/AgentGraphContext";
import {
  buildEntityFrequencyMap,
  analyzeTraceContent,
  ContentInsightData,
  TraceContentAnalysis,
} from "@/lib/trace-analyzer";

export interface EntityWordCloudData {
  text: string;
  value: number;
  type: string;
  color: string;
}

export interface RelationTypeData {
  type: string;
  count: number;
  percentage: number;
}

export interface SystemSummaryData {
  word: string;
  weight: number;
  context: string;
}

export interface TraceAnalyticsData {
  readingTime: number;
  complexity: number;
  keywords: string[];
  characterCount: number;
  wordCount: number;
  timeWithoutAgentGraph: number;
  timeSavedPercentage: number;
  comprehensionScore: number;
}

export interface ProcessingMethodData {
  method: string;
  count: number;
  percentage: number;
  color: string;
}

export interface DashboardVisualizationData {
  entityWordCloud: EntityWordCloudData[];
  relationTypes: RelationTypeData[];
  systemSummaries: SystemSummaryData[];
  traceAnalytics: TraceAnalyticsData[];
  processingMethods: ProcessingMethodData[];
  contentInsights: ContentInsightData;
  isLoading: boolean;
  error: string | null;
}

const ENTITY_TYPE_COLORS = {
  person: "#3b82f6",
  system: "#10b981",
  concept: "#8b5cf6",
  action: "#f59e0b",
  object: "#ef4444",
  default: "#6b7280",
};

const PROCESSING_METHOD_COLORS = {
  production: "#059669",
  baseline: "#dc2626",
  hybrid: "#7c3aed",
  sequential: "#ea580c",
  unknown: "#6b7280",
  default: "#6b7280",
};

function getEntityColor(type: string): string {
  return (
    ENTITY_TYPE_COLORS[type as keyof typeof ENTITY_TYPE_COLORS] ||
    ENTITY_TYPE_COLORS.default
  );
}

function getProcessingMethodColor(method: string): string {
  return (
    PROCESSING_METHOD_COLORS[method as keyof typeof PROCESSING_METHOD_COLORS] ||
    PROCESSING_METHOD_COLORS.default
  );
}

export function useDashboardVisualizationData(): DashboardVisualizationData {
  const { state } = useAgentGraph();
  const [isLoading, setIsLoading] = useState(true);
  const [error, setError] = useState<string | null>(null);

  // Process all traces to get comprehensive analysis
  const traceAnalyses = useMemo<TraceContentAnalysis[]>(() => {
    try {
      return state.traces.map((trace) => analyzeTraceContent(trace));
    } catch (err) {
      console.error("Error analyzing traces:", err);
      return [];
    }
  }, [state.traces]);

  const entityWordCloud = useMemo<EntityWordCloudData[]>(() => {
    try {
      const entities = buildEntityFrequencyMap(state.traces);
      return entities.map((entity) => ({
        text: entity.text,
        value: entity.frequency,
        type: entity.type,
        color: getEntityColor(entity.type),
      }));
    } catch (err) {
      console.error("Error processing entity word cloud data:", err);
      return [];
    }
  }, [state.traces]);

  const relationTypes = useMemo<RelationTypeData[]>(() => {
    try {
      const relationFrequency: Record<string, number> = {};
      let totalRelations = 0;

      traceAnalyses.forEach((analysis) => {
        analysis.relations.forEach((relation) => {
          const type = relation.type.replace(/_/g, " ");
          relationFrequency[type] = (relationFrequency[type] || 0) + 1;
          totalRelations += 1;
        });
      });

      if (totalRelations === 0) return [];

      return Object.entries(relationFrequency)
        .map(([type, count]) => ({
          type,
          count,
          percentage: Math.round((count / totalRelations) * 100),
        }))
        .sort((a, b) => b.count - a.count)
        .slice(0, 10);
    } catch (err) {
      console.error("Error processing relation types data:", err);
      return [];
    }
  }, [traceAnalyses]);

  const systemSummaries = useMemo<SystemSummaryData[]>(() => {
    try {
      const wordFrequency: Record<string, number> = {};

      traceAnalyses.forEach((analysis) => {
        analysis.insights.contentThemes.forEach((theme) => {
          wordFrequency[theme.theme] =
            (wordFrequency[theme.theme] || 0) + theme.weight;
        });
      });

      return Object.entries(wordFrequency)
        .map(([word, weight]) => ({
          word,
          weight,
          context: "content_theme",
        }))
        .sort((a, b) => b.weight - a.weight)
        .slice(0, 50);
    } catch (err) {
      console.error("Error processing system summaries data:", err);
      return [];
    }
  }, [traceAnalyses]);

  const traceAnalytics = useMemo<TraceAnalyticsData[]>(() => {
    try {
      return traceAnalyses.map((analysis) => ({
        readingTime: analysis.readingTime,
        complexity: analysis.complexity,
        keywords: analysis.insights.contentThemes
          .slice(0, 10)
          .map((theme) => theme.theme),
        characterCount: analysis.characterCount,
        wordCount: analysis.wordCount,
        timeWithoutAgentGraph: analysis.timeWithoutAgentGraph,
        timeSavedPercentage: analysis.timeSavedPercentage,
        comprehensionScore: analysis.comprehensionScore,
      }));
    } catch (err) {
      console.error("Error processing trace analytics data:", err);
      return [];
    }
  }, [traceAnalyses]);

  const processingMethods = useMemo<ProcessingMethodData[]>(() => {
    try {
      const methodFrequency: Record<string, number> = {};
      let totalMethods = 0;

      state.traces.forEach((trace) => {
        trace.knowledge_graphs?.forEach((kg: any) => {
          const method = kg.processing_metadata?.method_name || "unknown";
          methodFrequency[method] = (methodFrequency[method] || 0) + 1;
          totalMethods += 1;
        });
      });

      if (totalMethods === 0) return [];

      return Object.entries(methodFrequency)
        .map(([method, count]) => ({
          method,
          count,
          percentage: Math.round((count / totalMethods) * 100),
          color: getProcessingMethodColor(method),
        }))
        .sort((a, b) => b.count - a.count);
    } catch (err) {
      console.error("Error processing processing methods data:", err);
      return [];
    }
  }, [state.traces]);

  const contentInsights = useMemo<ContentInsightData>(() => {
    try {
      if (traceAnalyses.length === 0) {
        return {
          commonPatterns: [],
          questionCount: 0,
          commandCount: 0,
          averageResponseLength: 0,
          contentThemes: [],
        };
      }

      const totalQuestions = traceAnalyses.reduce(
        (sum, analysis) => sum + analysis.insights.questionCount,
        0
      );
      const totalCommands = traceAnalyses.reduce(
        (sum, analysis) => sum + analysis.insights.commandCount,
        0
      );
      const totalResponseLength = traceAnalyses.reduce(
        (sum, analysis) => sum + analysis.insights.averageResponseLength,
        0
      );
      const avgResponseLength =
        traceAnalyses.length > 0
          ? Math.round(totalResponseLength / traceAnalyses.length)
          : 0;

      // Aggregate themes from all traces
      const themeFrequency: Record<string, number> = {};
      traceAnalyses.forEach((analysis) => {
        analysis.insights.contentThemes.forEach((theme) => {
          themeFrequency[theme.theme] =
            (themeFrequency[theme.theme] || 0) + theme.weight;
        });
      });

      const contentThemes = Object.entries(themeFrequency)
        .map(([theme, weight]) => ({ theme, weight }))
        .sort((a, b) => b.weight - a.weight)
        .slice(0, 10);

      // Aggregate common patterns
      const patternSet = new Set<string>();
      traceAnalyses.forEach((analysis) => {
        analysis.insights.commonPatterns.forEach((pattern) =>
          patternSet.add(pattern)
        );
      });

      return {
        commonPatterns: Array.from(patternSet),
        questionCount: totalQuestions,
        commandCount: totalCommands,
        averageResponseLength: avgResponseLength,
        contentThemes,
      };
    } catch (err) {
      console.error("Error processing content insights data:", err);
      return {
        commonPatterns: [],
        questionCount: 0,
        commandCount: 0,
        averageResponseLength: 0,
        contentThemes: [],
      };
    }
  }, [traceAnalyses]);

  useEffect(() => {
    const processData = async () => {
      try {
        setIsLoading(true);
        setError(null);

        // Small delay to prevent blocking UI
        await new Promise((resolve) => setTimeout(resolve, 100));

        setIsLoading(false);
      } catch (err) {
        setError(
          err instanceof Error
            ? err.message
            : "Failed to process visualization data"
        );
        setIsLoading(false);
      }
    };

    processData();
  }, [state.traces]);

  return {
    entityWordCloud,
    relationTypes,
    systemSummaries,
    traceAnalytics,
    processingMethods,
    contentInsights,
    isLoading,
    error,
  };
}