frontend/src/components/features/traces/TraceKnowledgeGraphView.tsx

import React, { useState, useEffect, useCallback } from "react";
import { TooltipProvider } from "@/components/ui/tooltip";
import { Trace, KnowledgeGraph } from "@/types";
import { useAgentGraph } from "@/context/AgentGraphContext";
import { useKGDisplayMode } from "@/context/KGDisplayModeContext";
import { filterKnowledgeGraphsByMode } from "@/lib/kg-selection";
import { AgentGraphsSection } from "@/components/features/traces/AgentGraphsSection";
import { TraceOverviewSection } from "@/components/features/traces/TraceOverviewSection";
import { ContextDocumentsModal } from "@/components/shared/modals/ContextDocumentsModal";
import { TraceContentModal } from "@/components/shared/modals/TraceContentModal";
import { api } from "@/lib/api";
import { useModal } from "@/context/ModalContext";
import { useNotification } from "@/context/NotificationContext";
import { useNavigation } from "@/context/NavigationContext";
import {
  SplitterSelectionModal,
  SplitterType,
  // ChunkingConfig,
} from "@/components/shared/modals/SplitterSelectionModal";
import { useTaskPolling } from "@/hooks/useTaskPolling";
import { MetadataCardConfig } from "@/components/shared/CompactMetadataCard";
import {
  MetadataCardSelector,
  AVAILABLE_CARDS,
  DEFAULT_CARDS,
} from "@/components/shared/modals/MetadataCardSelector";

interface TraceKnowledgeGraphViewProps {
  trace: Trace;
  knowledgeGraphs: KnowledgeGraph[];
}

export function TraceKnowledgeGraphView({
  trace,
  knowledgeGraphs: _knowledgeGraphs,
}: TraceKnowledgeGraphViewProps) {
  const { actions } = useAgentGraph();
  const {
    setActiveView, // Used in TraceContentModal onEditTrace
  } = actions;
  const { mode: kgDisplayMode } = useKGDisplayMode();
  const [kgData, setKgData] = useState<KnowledgeGraph[]>([]);
  const [loading, setLoading] = useState(true);
  const [error, setError] = useState<string | null>(null);
  useModal(); // Keep for potential future use
  const { showNotification } = useNotification();
  const { actions: navigationActions } = useNavigation();

  // Add state for trace content
  const [traceContent, setTraceContent] = useState<string>("");
  const [traceContentLoading, setTraceContentLoading] = useState(false);
  const [traceContentError, setTraceContentError] = useState<string | null>(
    null
  );

  // State for metadata regeneration
  const [metadataOutdated, setMetadataOutdated] = useState(false);
  const [updatedMetadata, setUpdatedMetadata] = useState<any>(null);
  const [updatedTraceData, setUpdatedTraceData] = useState<any>(null);

  // State for enhanced statistics
  const [enhancedStats, setEnhancedStats] = useState<any>(null);
  const [enhancedStatsLoading, setEnhancedStatsLoading] = useState(false);

  // Get current metadata (use updated if available, otherwise original)
  const currentMetadata = updatedMetadata || trace.metadata;

  // Helper to show both toast and persistent notifications
  const showSystemNotification = useCallback(
    (notification: {
      type: "success" | "error" | "warning" | "info";
      title: string;
      message: string;
    }) => {
      // Show toast notification
      showNotification(notification);

      // Add to persistent notifications for all types
      navigationActions.addNotification({
        type: notification.type,
        title: notification.title,
        message: notification.message,
      });
    },
    [showNotification, navigationActions]
  );

  // Function to check if metadata is outdated
  const checkMetadataFreshness = useCallback(
    (trace: Trace, content: string, metadata: any, updatedTraceData?: any) => {
      if (!metadata?.schema_analytics || !content) {
        return false; // No metadata to check
      }

      // Check 1: JSON validity - if current content is invalid JSON, metadata is definitely outdated
      try {
        JSON.parse(content);
      } catch (e) {
        return true; // Invalid JSON means metadata is outdated
      }

      // Check 2: Content length difference (>5% threshold)
      const currentLength = content.length;
      // Use updated trace data if available, otherwise use original trace
      const metadataLength =
        updatedTraceData?.character_count || trace.character_count || 0;
      if (metadataLength > 0) {
        const lengthDifference =
          Math.abs(currentLength - metadataLength) / metadataLength;
        if (lengthDifference > 0.05) {
          return true;
        }
      }

      // Check 3: Character count mismatch (exact)
      if (currentLength !== metadataLength) {
        return true;
      }

      return false;
    },
    []
  );

  // Check metadata freshness when trace content is loaded
  useEffect(() => {
    if (traceContent && trace) {
      const outdated = checkMetadataFreshness(
        trace,
        traceContent,
        currentMetadata,
        updatedTraceData
      );
      setMetadataOutdated(outdated);
    }
  }, [
    traceContent,
    trace,
    currentMetadata,
    updatedTraceData,
    checkMetadataFreshness,
  ]);

  // Auto-refresh metadata when outdated
  useEffect(() => {
    const autoRefreshMetadata = async () => {
      if (!metadataOutdated || !trace.trace_id) return;

      try {
        // Call the universal parser endpoint to regenerate metadata
        const response = await fetch(
          `/api/traces/${trace.trace_id}/regenerate-metadata`,
          {
            method: "POST",
            headers: {
              "Content-Type": "application/json",
            },
          }
        );

        if (!response.ok) {
          throw new Error("Failed to regenerate metadata");
        }

        // Fetch updated trace data
        const traceResponse = await fetch(`/api/traces/${trace.trace_id}`);
        if (traceResponse.ok) {
          const response = await traceResponse.json();
          const updatedTrace = response.trace; // Extract trace from response wrapper
          // Store the updated metadata and trace data locally
          setUpdatedMetadata(updatedTrace.metadata);
          setUpdatedTraceData(updatedTrace);

          // Reset metadata outdated flag
          setMetadataOutdated(false);
        }
      } catch (error) {
        console.error("Auto-refresh metadata failed:", error);
      }
    };

    // Auto-refresh metadata every 90 seconds if outdated (reduced frequency for HF Spaces)
    const interval = setInterval(() => {
      autoRefreshMetadata();
    }, 90000); // 90 seconds (increased from 30s to avoid 429 errors)

    return () => clearInterval(interval);
  }, [metadataOutdated, trace.trace_id]);

  const [isSplitterModalOpen, setIsSplitterModalOpen] = useState(false);
  const [isContextDocumentsModalOpen, setIsContextDocumentsModalOpen] =
    useState(false);
  const [isTraceContentModalOpen, setIsTraceContentModalOpen] = useState(false);
  const [isGenerating, setIsGenerating] = useState(false);
  const [generationProgress, setGenerationProgress] = useState(0);

  // Load trace content (clean content for overview display)
  const fetchTraceContent = useCallback(async () => {
    if (!trace.trace_id) return;

    setTraceContentLoading(true);
    setTraceContentError(null);

    try {
      // Use clean content for overview display, not numbered content
      const content = await api.traces.getContent(trace.trace_id);
      setTraceContent(content);
    } catch (err) {
      setTraceContentError(
        err instanceof Error ? err.message : "Failed to load trace content"
      );
    } finally {
      setTraceContentLoading(false);
    }
  }, [trace.trace_id]);

  const fetchKnowledgeGraphs = useCallback(async () => {
    if (!trace.trace_id) return;

    setLoading(true);
    setError(null);

    try {
      // Also check if trace already has knowledge graphs
      console.log("Trace object:", trace);
      console.log("Trace knowledge_graphs:", trace.knowledge_graphs);
      const response = await fetch(
        `/api/traces/${trace.trace_id}/knowledge-graphs`
      );

      if (!response.ok) {
        throw new Error(
          `Failed to fetch knowledge graphs: ${response.statusText}`
        );
      }

      const data = await response.json();

      console.log("Raw API response:", data);

      if (data?.knowledge_graphs?.length) {
        console.log(
          "Total knowledge graphs received:",
          data.knowledge_graphs.length
        );

        // Log each KG to see their structure
        data.knowledge_graphs.forEach((kg: any, index: number) => {
          console.log(`KG ${index}:`, {
            id: kg.id,
            kg_id: kg.kg_id,
            is_final: kg.is_final,
            window_index: kg.window_index,
            window_total: kg.window_total,
            filename: kg.filename,
          });
        });

        // Filter for final graphs using the exact logic from stage_processor.js
        const finalKGs = data.knowledge_graphs.filter((kg: any) => {
          const hasValidId = Boolean(kg.kg_id || kg.id);
          const isFinal =
            kg.is_final === true ||
            (kg.window_index === null && kg.window_total !== null);

          if (!hasValidId) {
            console.warn(
              "Found knowledge graph with missing ID - skipping:",
              kg
            );
            return false;
          }

          if (!isFinal) {
            console.log(
              "Skipping non-final KG:",
              kg.kg_id,
              "window_index:",
              kg.window_index,
              "is_final:",
              kg.is_final
            );
          } else {
            console.log(
              "Including final KG:",
              kg.kg_id,
              "is_final:",
              kg.is_final,
              "window_index:",
              kg.window_index,
              "window_total:",
              kg.window_total
            );
          }

          return isFinal;
        });

        console.log("Filtered final KGs:", finalKGs.length, finalKGs);

        // Apply display mode filtering
        const filteredByMode = filterKnowledgeGraphsByMode(
          finalKGs,
          kgDisplayMode
        );
        console.log(
          `Applied ${kgDisplayMode} mode filtering:`,
          filteredByMode.length,
          "KGs"
        );

        setKgData(filteredByMode);
      } else {
        console.log(
          "No knowledge graphs in API response, checking trace object"
        );
        // Fallback to trace.knowledge_graphs if API returns empty
        if (trace.knowledge_graphs && trace.knowledge_graphs.length > 0) {
          console.log("Using trace.knowledge_graphs:", trace.knowledge_graphs);
          setKgData(trace.knowledge_graphs);
        } else {
          setKgData([]);
        }
      }
    } catch (err) {
      console.log("API call failed, trying trace.knowledge_graphs fallback");
      setError(
        err instanceof Error ? err.message : "Failed to fetch knowledge graphs"
      );

      // Fallback to trace data if API fails
      if (trace.knowledge_graphs && trace.knowledge_graphs.length > 0) {
        console.log(
          "Using trace.knowledge_graphs as fallback:",
          trace.knowledge_graphs
        );
        setKgData(trace.knowledge_graphs);
        setError(null); // Clear error since we found data
      } else {
        setKgData([]);
      }
    } finally {
      setLoading(false);
    }
  }, [trace, kgDisplayMode]);

  // Fetch enhanced statistics
  const fetchEnhancedStatistics = useCallback(async () => {
    if (!trace.trace_id) return;

    setEnhancedStatsLoading(true);
    try {
      const response = await api.traces.getEnhancedStatistics(trace.trace_id);
      setEnhancedStats(response.enhanced_statistics);
    } catch (err) {
      console.error("Failed to fetch enhanced statistics:", err);
      setEnhancedStats(null);
    } finally {
      setEnhancedStatsLoading(false);
    }
  }, [trace.trace_id]);

  useEffect(() => {
    fetchKnowledgeGraphs();
    fetchTraceContent();
    fetchEnhancedStatistics();
  }, [fetchKnowledgeGraphs, fetchTraceContent, fetchEnhancedStatistics]);

  // Auto-refresh knowledge graphs every 10 seconds
  useEffect(() => {
    const interval = setInterval(() => {
      if (!loading && !isGenerating) {
        fetchKnowledgeGraphs();
      }
    }, 10000); // 10 seconds

    return () => clearInterval(interval);
  }, [fetchKnowledgeGraphs, loading, isGenerating]);

  const handleGenerateKG = () => {
    setIsSplitterModalOpen(true);
  };

  // Task polling setup for knowledge graph generation
  const { pollTaskStatus, isPolling } = useTaskPolling({
    onSuccess: (taskId) => {
      console.log("Knowledge graph generation completed:", taskId);
      showSystemNotification({
        type: "success",
        title: "Agent Graph Generation Complete",
        message: "Your agent graph has been generated successfully!",
      });

      // Refresh knowledge graphs list
      fetchKnowledgeGraphs();
      setIsGenerating(false);
      setGenerationProgress(0);
    },
    onError: (error, taskId) => {
      console.error("Knowledge graph generation failed:", error, taskId);

      // Check if this is a timeout vs a real failure
      const isTimeout =
        error.includes("timeout") || error.includes("timed out");

      showSystemNotification({
        type: isTimeout ? "warning" : "error",
        title: isTimeout
          ? "Generation Taking Longer Than Expected"
          : "Agent Graph Generation Failed",
        message: isTimeout
          ? `${error} You can refresh the page to check if it completed.`
          : error,
      });

      // Only reset state if it's not a timeout (user might want to keep waiting)
      if (!isTimeout) {
        setIsGenerating(false);
        setGenerationProgress(0);
      }
    },
    onProgress: (progress, message) => {
      setGenerationProgress(progress);
      if (message) {
        console.log("Generation progress:", progress, message);
      }
    },
    maxAttempts: 180, // 15 minutes with exponential backoff
    interval: 5000, // Start with 5 seconds
    enableExponentialBackoff: true,
  });

  const handleSplitterConfirm = async (
    splitterType: SplitterType,
    methodName?: string,
    model?: string,
    chunkingConfig?: { min_chunk_size?: number; max_chunk_size?: number }
  ) => {
    const finalMethodName = methodName || "production";
    const finalModel = model || "gpt-5-mini";
    console.log("TraceKnowledgeGraphView: Using method name:", finalMethodName);
    console.log("TraceKnowledgeGraphView: Using model:", finalModel);
    console.log(
      "TraceKnowledgeGraphView: Using chunking config:",
      chunkingConfig
    );

    setIsGenerating(true);
    setIsSplitterModalOpen(false); // Safety fallback - modal should already be closed
    setGenerationProgress(0);

    try {
      const response = await api.traces.generateKnowledgeGraph(
        trace.trace_id,
        splitterType,
        true, // force_regenerate = true to allow generating new graphs even if existing ones exist
        finalMethodName,
        finalModel,
        chunkingConfig
      );

      showSystemNotification({
        type: "info",
        title: "Agent Graph Generation Started",
        message: `Using ${getSplitterDisplayName(
          splitterType
        )} splitter with ${finalMethodName} method. This may take several minutes.`,
      });

      // Start polling for task status if we got a task_id
      if (response.task_id) {
        pollTaskStatus(response.task_id);
      } else {
        // If no task_id, just refresh immediately and mark as complete
        await fetchKnowledgeGraphs();
        setIsGenerating(false);
        showSystemNotification({
          type: "success",
          title: "Agent Graph Generation Complete",
          message: "Your agent graph has been generated successfully!",
        });
      }
    } catch (error) {
      console.error("Failed to generate knowledge graph:", error);
      showNotification({
        type: "error",
        title: "Agent Graph Generation Failed",
        message: "Failed to start agent graph generation. Please try again.",
      });
      setIsGenerating(false);
      setGenerationProgress(0);
    }
  };

  // Helper function to get display name for splitter type (matching stage_processor.js)
  const getSplitterDisplayName = (splitterType: SplitterType) => {
    switch (splitterType) {
      case "agent_semantic":
        return "Agent Semantic";
      case "json":
        return "JSON";
      case "prompt_interaction":
        return "Prompt Interaction";
      default:
        return splitterType;
    }
  };

  // Metadata cards state management
  const [selectedMetadataCards, setSelectedMetadataCards] = useState<string[]>(
    () => {
      const saved = localStorage.getItem("metadata-cards-selection");
      return saved ? JSON.parse(saved) : DEFAULT_CARDS;
    }
  );
  const [isMetadataCardSelectorOpen, setIsMetadataCardSelectorOpen] =
    useState(false);

  // Save metadata card selection to localStorage
  const handleSaveMetadataCards = (cards: string[]) => {
    setSelectedMetadataCards(cards);
    localStorage.setItem("metadata-cards-selection", JSON.stringify(cards));
  };

  // Generate metadata card configs
  const generateMetadataCards = (): MetadataCardConfig[] => {
    const formatDate = (timestamp?: string) => {
      if (!timestamp) return "N/A";
      return new Date(timestamp).toLocaleDateString();
    };

    const formatCurrency = (amount: number | string): string => {
      const num = typeof amount === "string" ? parseFloat(amount) : amount;
      if (isNaN(num) || num === 0) return "$0.00";
      if (num < 0.01) return `$${num.toFixed(6)}`;
      return `$${num.toFixed(4)}`;
    };

    const getCardValue = (cardId: string): string | number => {
      // Get schema analytics from trace metadata
      const schemaAnalytics = currentMetadata?.schema_analytics;

      // Use enhanced stats if available and not loading
      const stats =
        !enhancedStatsLoading && enhancedStats ? enhancedStats : null;

      switch (cardId) {
        case "size":
          return trace.character_count || 0;
        case "type":
          return trace.trace_type || "Unknown";
        case "uploaded":
          return formatDate(trace.upload_timestamp);
        case "graphs":
          return kgData.length;
        case "modified":
          return formatDate(trace.upload_timestamp);

        case "source":
          return "Manual Upload";
        case "method":
          return "Production";

        // Schema analytics cards
        case "depth":
          return (
            stats?.components?.max_depth ||
            schemaAnalytics?.numerical_overview?.component_stats?.max_depth ||
            "N/A"
          );
        case "execution_time": {
          const totalTimeMs =
            stats?.performance?.total_execution_time_ms ||
            schemaAnalytics?.numerical_overview?.timing_analytics
              ?.total_execution_time_ms;
          if (totalTimeMs && totalTimeMs > 0) {
            if (totalTimeMs >= 1000) {
              return `${(totalTimeMs / 1000).toFixed(1)}s`;
            } else {
              return `${totalTimeMs}ms`;
            }
          }
          return "N/A";
        }
        case "total_tokens":
          return (
            stats?.tokens?.total_tokens ||
            schemaAnalytics?.numerical_overview?.token_analytics
              ?.total_tokens ||
            "N/A"
          );
        case "prompt_calls":
          return (
            stats?.prompt_calls?.total_calls ||
            schemaAnalytics?.prompt_analytics?.prompt_calls_detected ||
            "N/A"
          );
        case "components":
          return (
            stats?.components?.total_components ||
            schemaAnalytics?.numerical_overview?.component_stats
              ?.total_components ||
            "N/A"
          );
        case "success_rate": {
          const successRate =
            stats?.components?.success_rate ||
            schemaAnalytics?.numerical_overview?.component_stats?.success_rate;
          if (successRate !== undefined && successRate !== null) {
            return `${successRate.toFixed(1)}%`;
          }
          return "N/A";
        }

        // Cost cards
        case "total_cost":
          return enhancedStatsLoading
            ? "Loading..."
            : stats?.cost?.total_cost_usd
            ? formatCurrency(stats.cost.total_cost_usd)
            : "N/A";
        case "cost_per_call":
          return enhancedStatsLoading
            ? "Loading..."
            : stats?.cost?.avg_cost_per_call_usd
            ? formatCurrency(stats.cost.avg_cost_per_call_usd)
            : "N/A";
        case "input_cost":
          return enhancedStatsLoading
            ? "Loading..."
            : stats?.cost?.input_cost_usd
            ? formatCurrency(stats.cost.input_cost_usd)
            : "N/A";
        case "output_cost":
          return enhancedStatsLoading
            ? "Loading..."
            : stats?.cost?.output_cost_usd
            ? formatCurrency(stats.cost.output_cost_usd)
            : "N/A";

        // Token analytics cards
        case "avg_input_tokens":
          return enhancedStatsLoading
            ? "Loading..."
            : stats?.prompt_calls?.avg_prompt_tokens_per_call
            ? Math.round(stats.prompt_calls.avg_prompt_tokens_per_call)
            : "N/A";
        case "avg_output_tokens":
          return enhancedStatsLoading
            ? "Loading..."
            : stats?.prompt_calls?.avg_completion_tokens_per_call
            ? Math.round(stats.prompt_calls.avg_completion_tokens_per_call)
            : "N/A";

        default:
          return "N/A";
      }
    };

    return selectedMetadataCards
      .map((cardId) => {
        const cardConfig = AVAILABLE_CARDS.find((card) => card.id === cardId);
        if (!cardConfig) return null;

        return {
          ...cardConfig,
          value: getCardValue(cardId),
        };
      })
      .filter((card): card is MetadataCardConfig => card !== null);
  };

  return (
    <TooltipProvider delayDuration={300} skipDelayDuration={100}>
      <div className="flex-1 flex flex-col p-4 space-y-4 min-h-0">
        {/* Main Content */}
        <div className="flex-1 space-y-4">
          {/* Main Side-by-Side Section - Agent Graphs (Left) and Trace Overview (Right) */}
          <div className="flex-1 flex bg-background border border-border/30 rounded-lg min-h-0 gap-4 p-4">
            {/* Left Column - Agent Graphs */}
            <div className="flex-1 flex flex-col min-h-0">
              <AgentGraphsSection
                trace={trace}
                kgData={kgData}
                loading={loading}
                error={error}
                isGenerating={isGenerating}
                isPolling={isPolling}
                generationProgress={generationProgress}
                onGenerateKG={handleGenerateKG}
                onRefreshKnowledgeGraphs={fetchKnowledgeGraphs}
                onOpenContextDocuments={() =>
                  setIsContextDocumentsModalOpen(true)
                }
              />
            </div>

            {/* Right Column - Trace Overview */}
            <div className="flex-1 flex flex-col min-h-0">
              <TraceOverviewSection
                trace={trace}
                metadataOutdated={metadataOutdated}
                currentMetadata={currentMetadata}
                enhancedStats={enhancedStats}
                enhancedStatsLoading={enhancedStatsLoading}
                onOpenTraceContent={() => setIsTraceContentModalOpen(true)}
                metadataCards={generateMetadataCards()}
                onOpenMetadataCardSelector={() =>
                  setIsMetadataCardSelectorOpen(true)
                }
              />
            </div>
          </div>
        </div>

        {/* Splitter Selection Modal */}
        <SplitterSelectionModal
          open={isSplitterModalOpen}
          onOpenChange={setIsSplitterModalOpen}
          onConfirm={handleSplitterConfirm}
          isLoading={isGenerating}
        />

        {/* Metadata Card Selector Modal */}
        <MetadataCardSelector
          open={isMetadataCardSelectorOpen}
          onOpenChange={setIsMetadataCardSelectorOpen}
          selectedCards={selectedMetadataCards}
          onSave={handleSaveMetadataCards}
          trace={trace}
          graphCount={kgData.length}
        />

        {/* Context Documents Modal */}
        <ContextDocumentsModal
          isOpen={isContextDocumentsModalOpen}
          onClose={() => setIsContextDocumentsModalOpen(false)}
          traceId={trace.trace_id}
        />

        {/* Trace Content Modal */}
        <TraceContentModal
          isOpen={isTraceContentModalOpen}
          onClose={() => setIsTraceContentModalOpen(false)}
          trace={trace}
          traceContent={traceContent}
          traceContentLoading={traceContentLoading}
          traceContentError={traceContentError}
          onEditTrace={() => setActiveView("trace-editor")}
        />
      </div>
    </TooltipProvider>
  );
}