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

import React, {
  useEffect,
  useState,
  useCallback,
  useMemo,
  useRef,
} from "react";
import { Button } from "@/components/ui/button";
import { ZoomIn, ZoomOut, RotateCcw } from "lucide-react";
import { KnowledgeGraph, OptimizationRecommendation } from "@/types";
import {
  UniversalGraphData,
  UniversalNode,
  UniversalLink,
} from "@/types/graph-visualization";
import { getGraphDataAdapter } from "@/lib/graph-data-adapters";
import { createKnowledgeGraphConfig } from "@/lib/graph-config-factory";
import {
  SchemaCapabilities,
  detectSchemaType,
  getSchemaCapabilities,
} from "@/lib/schema-detection";
import { CytoscapeGraphCore } from "@/lib/cytoscape-graph-core";
import { GraphSelectionCallbacks } from "@/types/graph-visualization";
import { ElementInfoSidebar } from "@/components/features/temporal/ElementInfoSidebar";
import { GraphLegend } from "@/components/shared/GraphLegend";
import { useAgentGraph } from "@/context/AgentGraphContext";
import { api } from "@/lib/api";

interface KnowledgeGraphVisualizerProps {
  knowledgeGraph: KnowledgeGraph;
}

export const KnowledgeGraphVisualizer: React.FC<
  KnowledgeGraphVisualizerProps
> = ({ knowledgeGraph: kg }) => {
  const containerRef = useRef<HTMLDivElement>(null);
  const cytoscapeRef = useRef<CytoscapeGraphCore | null>(null);

  const [graphData, setGraphData] = useState<UniversalGraphData>({
    nodes: [],
    links: [],
  });
  const [loading, setLoading] = useState(true);
  const [error, setError] = useState<string | null>(null);
  const [containerDimensions, setContainerDimensions] = useState({
    width: 800,
    height: 600,
  });

  // Failures extracted from graph data for system-level overview
  const [failures, setFailures] = useState<
    {
      id: string;
      risk_type: string;
      description: string;
    }[]
  >([]);
  const [optimizations, setOptimizations] = useState<
    OptimizationRecommendation[]
  >([]);

  // Selection state
  const [selectedElement, setSelectedElement] = useState<
    UniversalNode | UniversalLink | null
  >(null);
  const [selectedElementType, setSelectedElementType] = useState<
    "node" | "link" | "failure" | null
  >(null);

  // Trace viewer state
  const { state: agState } = useAgentGraph();
  const [traceLines, setTraceLines] = useState<string[]>([]);
  const [highlightRanges, setHighlightRanges] = useState<
    { start: number; end: number }[]
  >([]);

  // Schema capabilities state
  const [schemaCapabilities, setSchemaCapabilities] =
    useState<SchemaCapabilities>({
      supportsLineReferences: false,
      supportsFailures: false,
      supportsImportanceLevels: false,
      supportsInteractionPrompts: false,
    });

  const dataAdapter = useMemo(() => getGraphDataAdapter(), []);

  // Helper function to find affected element by ID with smart matching
  const findAffectedElement = useCallback(
    (affectedId: string, graphData: UniversalGraphData) => {
      // Strategy 1: Try exact ID match first
      let affectedElement =
        graphData.nodes.find((n) => n.id === affectedId) ||
        graphData.links.find((l) => l.id === affectedId);

      // Strategy 2: If not found, try smart matching based on ID pattern
      if (!affectedElement) {
        const match = affectedId.match(
          /^(task|tool|agent|input|output|human)_(\d+)$/i
        );
        if (match) {
          const [, type, index] = match;
          if (type && index) {
            const targetIndex = parseInt(index) - 1; // Convert to 0-based index

            // Find nodes of the matching type
            const nodesOfType = graphData.nodes.filter(
              (n) => n.type && n.type.toLowerCase() === type.toLowerCase()
            );

            // Get the node at the specified index
            if (nodesOfType[targetIndex]) {
              affectedElement = nodesOfType[targetIndex];
              console.log(
                `Found element by type matching: ${type}[${targetIndex}] -> ${affectedElement.id}`
              );
            }
          }
        }
      }

      // Strategy 3: If still not found, try partial name matching
      if (!affectedElement) {
        affectedElement =
          graphData.nodes.find(
            (n) =>
              n.name && n.name.toLowerCase().includes(affectedId.toLowerCase())
          ) ||
          graphData.links.find(
            (l) =>
              l.label &&
              l.label.toLowerCase().includes(affectedId.toLowerCase())
          );

        if (affectedElement) {
          console.log(
            `Found element by partial name matching: ${affectedId} -> ${affectedElement.id}`
          );
        }
      }

      return affectedElement;
    },
    []
  );

  // Fetch graph data from API
  const fetchGraphData = useCallback(async () => {
    try {
      setLoading(true);
      setError(null);

      console.log("Fetching graph data for kg_id:", kg.kg_id);

      // Use the same endpoint as agent_graph_visualizer.html
      const response = await fetch(
        `/api/knowledge-graphs/${encodeURIComponent(kg.kg_id)}`
      );

      console.log("API response status:", response.status);

      if (!response.ok) {
        const errorText = await response.text();
        console.log("API error response:", errorText);
        throw new Error(
          `Failed to fetch graph data: ${response.status} - ${errorText}`
        );
      }

      const data = await response.json();
      console.log("API response data:", data);
      console.log("First entity sample:", data.entities?.[0]);
      console.log("First relation sample:", data.relations?.[0]);

      // Convert the API response using our data adapter
      if (data.entities && data.relations) {
        console.log(
          "Using real API data:",
          data.entities.length,
          "entities,",
          data.relations.length,
          "relations"
        );
        const rawData = await api.knowledgeGraphs.getData(
          kg.kg_id || kg.filename
        );

        // Detect schema type and set capabilities
        const schemaType = detectSchemaType(rawData);
        const capabilities = getSchemaCapabilities(schemaType);
        setSchemaCapabilities(capabilities);

        const adapted = dataAdapter.adapt(rawData);
        const rawFailures = Array.isArray((rawData as any).failures)
          ? (rawData as any).failures
          : [];
        const rawOptimizations = Array.isArray((rawData as any).optimizations)
          ? (rawData as any).optimizations
          : [];

        // Enrich failures with matched element information
        const enrichedFailures = rawFailures.map((failure: any) => {
          if (failure.affected_id) {
            const matchedElement = findAffectedElement(
              failure.affected_id,
              adapted
            );
            if (matchedElement) {
              return {
                ...failure,
                matchedElement,
                displayName:
                  ("name" in matchedElement && matchedElement.name) ||
                  ("label" in matchedElement && matchedElement.label) ||
                  matchedElement.id,
              };
            }
          }
          return failure;
        });

        // Enrich optimizations with matched element information
        const enrichedOptimizations = rawOptimizations.map((opt: any) => {
          const affectedNodes = (opt.affected_ids || []).map((id: string) => {
            const matchedElement = findAffectedElement(id, adapted);
            return {
              id,
              name: matchedElement
                ? ("name" in matchedElement && matchedElement.name) ||
                  ("label" in matchedElement && matchedElement.label) ||
                  id
                : id,
            };
          });
          return { ...opt, affected_nodes: affectedNodes };
        });

        // Debug: Log the data structure to understand the mismatch
        console.log("Raw failures data:", rawFailures);
        console.log("Enriched failures:", enrichedFailures);
        console.log("Raw optimizations data:", rawOptimizations);
        console.log("Enriched optimizations:", enrichedOptimizations);
        console.log("Adapted graph nodes:", adapted.nodes.slice(0, 3));

        setFailures(enrichedFailures);
        setOptimizations(enrichedOptimizations);
        setGraphData(adapted);
      } else {
        throw new Error(
          "API data format unexpected - missing entities or relations"
        );
      }
    } catch (error) {
      console.error("Error fetching graph data:", error);
      setError(
        `Failed to load graph data: ${
          error instanceof Error ? error.message : String(error)
        }`
      );
    } finally {
      setLoading(false);
    }
  }, [kg.kg_id, kg.filename, dataAdapter]);

  // Load data on component mount
  useEffect(() => {
    fetchGraphData();
  }, [fetchGraphData]);

  // Fetch and number trace once when component mounts or selectedTrace changes
  useEffect(() => {
    const fetchTrace = async () => {
      if (!agState.selectedTrace) return;
      try {
        const numbered = await api.traces.getNumberedContent(
          agState.selectedTrace.trace_id
        );
        setTraceLines(numbered.split("\n"));
      } catch (err) {
        console.error("Failed to fetch trace content", err);
      }
    };
    fetchTrace();
  }, [agState.selectedTrace]);

  // Track container dimensions for draggable legend bounds
  useEffect(() => {
    const resizeObserver = new ResizeObserver((entries) => {
      for (const entry of entries) {
        const { width, height } = entry.contentRect;
        setContainerDimensions({ width, height });
      }
    });

    if (containerRef.current) {
      resizeObserver.observe(containerRef.current);
      // Set initial dimensions
      const rect = containerRef.current.getBoundingClientRect();
      setContainerDimensions({ width: rect.width, height: rect.height });
    }

    return () => {
      resizeObserver.disconnect();
    };
  }, []);

  // Compute highlight ranges when element selection changes
  useEffect(() => {
    if (!selectedElement) {
      setHighlightRanges([]);
      return;
    }

    const sel = selectedElement as any;
    const refs = [
      sel.raw_prompt_ref,
      sel.raw_text_ref,
      sel.interaction_prompt_ref,
    ].find((arr) => Array.isArray(arr) && arr.length > 0);
    if (!Array.isArray(refs)) {
      setHighlightRanges([]);
      return;
    }

    const isContinuation = (idx: number): boolean => {
      if (idx < 0 || idx >= traceLines.length) return false;
      const line = traceLines[idx] ?? "";
      // Strip <L#> prefix
      const contentPart = line.replace(/^<L\d+>\s*/, "");
      return contentPart.startsWith("  ");
    };

    const ranges = refs.map((r: any) => {
      const start = r.line_start - 1;
      let end = r.line_end - 1;
      // Extend downward for continuation lines
      while (end + 1 < traceLines.length && isContinuation(end + 1)) {
        end += 1;
      }
      return { start, end };
    });
    setHighlightRanges(ranges);
  }, [selectedElement, traceLines]);

  // Initialize Cytoscape visualization
  useEffect(() => {
    console.log("KnowledgeGraphVisualizer: Cytoscape useEffect triggered", {
      containerRef: !!containerRef.current,
      dataNodes: graphData.nodes.length,
      dataLinks: graphData.links.length,
      loading,
    });

    if (loading || graphData.nodes.length === 0) {
      console.log(
        "KnowledgeGraphVisualizer: Skipping Cytoscape init - still loading or no data"
      );
      return;
    }

    const initializeVisualization = async () => {
      try {
        if (!containerRef.current) {
          console.log(
            "KnowledgeGraphVisualizer: Refs not ready during initialization"
          );
          return;
        }

        // Wait for layout to settle
        await new Promise((resolve) => setTimeout(resolve, 200));

        const container = containerRef.current;
        let width = container.clientWidth || 800;
        let height = container.clientHeight || 600;

        // If container has zero dimensions, try to get parent dimensions
        if (width === 0 || height === 0) {
          const parent = container.parentElement;
          if (parent) {
            width = parent.clientWidth || 800;
            height = parent.clientHeight || 600;
            console.log("KnowledgeGraphVisualizer: Using parent dimensions", {
              parentWidth: width,
              parentHeight: height,
            });
          }
        }

        // Ensure minimum dimensions
        width = Math.max(width, 500);
        height = Math.max(height, 400);

        const config = createKnowledgeGraphConfig({
          width,
          height,
          showToolbar: false,
          showSidebar: false,
          showStats: false,
          enableSearch: false,
          enableZoom: true,
          enablePan: true,
          enableDrag: true,
          enableSelection: true,
        });

        console.log(
          "KnowledgeGraphVisualizer: Creating CytoscapeGraphCore with config",
          config
        );

        // Create selection callbacks
        const selectionCallbacks: GraphSelectionCallbacks = {
          onNodeSelect: (node: UniversalNode) => {
            console.log("Node selected:", node);
            console.log("Node raw_prompt:", node.raw_prompt);
            console.log("Node properties:", node.properties);
            setSelectedElement(node);
            setSelectedElementType("node");
          },
          onLinkSelect: (link: UniversalLink) => {
            console.log("Link selected:", link);
            console.log("Link interaction_prompt:", link.interaction_prompt);
            console.log("Link properties:", link.properties);
            setSelectedElement(link);
            setSelectedElementType("link");
          },
          onClearSelection: () => {
            console.log("Selection cleared");
            setSelectedElement(null);
            setSelectedElementType(null);
          },
        };

        cytoscapeRef.current = new CytoscapeGraphCore(
          containerRef.current,
          config,
          selectionCallbacks
        );

        // Load data
        console.log("KnowledgeGraphVisualizer: Loading data", {
          nodes: graphData.nodes.length,
          links: graphData.links.length,
        });
        cytoscapeRef.current.updateGraph(graphData, true, failures);

        console.log(
          "KnowledgeGraphVisualizer: Cytoscape initialization complete"
        );
      } catch (err) {
        console.error(
          "KnowledgeGraphVisualizer: Error initializing Cytoscape:",
          err
        );
        setError(
          `Failed to initialize visualization: ${
            err instanceof Error ? err.message : String(err)
          }`
        );
      }
    };

    // Function to check if refs are ready and initialize
    const checkAndInitialize = () => {
      if (!containerRef.current) {
        console.log("KnowledgeGraphVisualizer: Missing refs, retrying...");
        return false;
      }
      return true;
    };

    // If refs not ready immediately, set up a retry mechanism
    if (!checkAndInitialize()) {
      let retryCount = 0;
      const maxRetries = 10;

      const retryInterval = setInterval(() => {
        retryCount++;
        console.log(
          `KnowledgeGraphVisualizer: Retry attempt ${retryCount}/${maxRetries}`
        );

        if (checkAndInitialize()) {
          clearInterval(retryInterval);
          initializeVisualization();
        } else if (retryCount >= maxRetries) {
          clearInterval(retryInterval);
          console.error(
            "KnowledgeGraphVisualizer: Max retries reached, giving up"
          );
          setError(
            "Failed to initialize visualization: DOM elements not ready"
          );
        }
      }, 200);

      return () => {
        clearInterval(retryInterval);
        if (cytoscapeRef.current) {
          cytoscapeRef.current.destroy();
          cytoscapeRef.current = null;
        }
      };
    }

    // If refs are ready, initialize immediately
    const timeoutId = setTimeout(initializeVisualization, 100);

    // Cleanup
    return () => {
      clearTimeout(timeoutId);
      if (cytoscapeRef.current) {
        cytoscapeRef.current.destroy();
        cytoscapeRef.current = null;
      }
    };
  }, [graphData, loading]);

  // Update graph when failures change to apply failure zones
  useEffect(() => {
    if (cytoscapeRef.current && graphData.nodes.length > 0 && !loading) {
      console.log("KnowledgeGraphVisualizer: Updating graph with failures", {
        failuresCount: failures.length,
        failureIds: failures
          .map((f) => (f as any).affected_id || (f as any).matchedElement?.id)
          .filter(Boolean),
      });
      cytoscapeRef.current.updateGraph(graphData, false, failures);
    }
  }, [failures, graphData, loading]);

  // Zoom control handlers
  const handleZoomIn = () => {
    if (cytoscapeRef.current) {
      cytoscapeRef.current.zoomIn();
    }
  };

  const handleZoomOut = () => {
    if (cytoscapeRef.current) {
      cytoscapeRef.current.zoomOut();
    }
  };

  const handleResetZoom = () => {
    if (cytoscapeRef.current) {
      cytoscapeRef.current.resetZoom();
    }
  };

  const handleEntitySelect = (entityId: string) => {
    if (cytoscapeRef.current) {
      cytoscapeRef.current.selectNodeById(entityId);
    }
  };

  if (loading) {
    return (
      <div className="flex flex-col h-screen bg-background">
        <div className="flex items-center justify-center flex-1">
          <div className="text-center">
            <div className="animate-spin rounded-full h-12 w-12 border-b-2 border-primary mx-auto mb-4"></div>
            <p className="text-lg font-medium">Loading Knowledge Graph...</p>
            <p className="text-muted-foreground">Preparing visualization</p>
          </div>
        </div>
      </div>
    );
  }

  if (error) {
    return (
      <div className="flex flex-col h-screen bg-background">
        <div className="flex items-center justify-center flex-1">
          <div className="text-center">
            <div className="text-red-500 mb-4 text-4xl">⚠️</div>
            <p className="text-lg font-medium text-red-600">
              Failed to Load Graph
            </p>
            <p className="text-muted-foreground mt-2">{error}</p>
            <button
              onClick={fetchGraphData}
              className="mt-4 px-4 py-2 bg-blue-600 text-white rounded hover:bg-blue-700"
            >
              Retry
            </button>
          </div>
        </div>
      </div>
    );
  }

  // Debug logging
  console.log("KnowledgeGraphVisualizer render:", {
    loading,
    error,
    graphDataNodes: graphData.nodes.length,
    graphDataLinks: graphData.links.length,
  });

  return (
    <div className="flex flex-col h-screen bg-background">
      {/* Visualization with sidebars */}
      <div className="flex-1 flex min-h-0 overflow-x-auto">
        {/* Graph container */}
        <div
          className="flex-1 relative p-4 min-w-0"
          style={{ minWidth: "400px" }}
        >
          <div
            ref={containerRef}
            className="w-full h-full bg-white border rounded-lg relative overflow-hidden shadow-sm"
            style={{ minHeight: "500px" }}
          >
            {/* Draggable Graph Legend */}
            <GraphLegend containerBounds={containerDimensions} />

            {/* Zoom Controls - Top Right */}
            <div className="absolute top-4 right-4 flex flex-col gap-2 z-50 pointer-events-auto">
              <Button
                variant="outline"
                size="sm"
                onClick={handleZoomIn}
                disabled={loading}
              >
                <ZoomIn className="h-4 w-4" />
              </Button>
              <Button
                variant="outline"
                size="sm"
                onClick={handleZoomOut}
                disabled={loading}
              >
                <ZoomOut className="h-4 w-4" />
              </Button>
              <Button
                variant="outline"
                size="sm"
                onClick={handleResetZoom}
                disabled={loading}
              >
                <RotateCcw className="h-4 w-4" />
              </Button>
            </div>
          </div>
        </div>

        {/* Right sidebar with all tabs - always show */}
        <ElementInfoSidebar
          selectedElement={selectedElement}
          selectedElementType={selectedElementType}
          currentData={
            cytoscapeRef.current?.getCurrentData() || { nodes: [], links: [] }
          }
          failures={failures}
          optimizations={optimizations}
          onFailureSelect={(failure) => {
            console.log("Failure selected:", failure);

            // Find the affected element using our smart matching
            let matchedElement = null;
            if ((failure as any).affected_id) {
              matchedElement = findAffectedElement(
                (failure as any).affected_id,
                graphData
              );
              if (matchedElement) {
                console.log("Matched element for failure:", matchedElement);
                // Set the matched element as selected instead of the failure
                setSelectedElement(matchedElement);
                setSelectedElementType("node");
              } else {
                // Fallback to failure if no element found
                setSelectedElement(failure);
                setSelectedElementType("failure");
              }
            } else {
              // No affected_id, select the failure
              setSelectedElement(failure);
              setSelectedElementType("failure");
            }

            // Auto-select the affected element in the graph if it exists
            if ((failure as any).affected_id && cytoscapeRef.current) {
              const affectedId = (failure as any).affected_id;
              console.log("Looking for affected element with ID:", affectedId);

              const cy = cytoscapeRef.current.getCytoscape();
              if (cy) {
                // Strategy 1: Try exact ID match first
                let affectedElement = cy.getElementById(affectedId);

                // Strategy 2: If not found, try smart matching based on ID pattern
                if (affectedElement.length === 0) {
                  const match = affectedId.match(
                    /^(task|tool|agent|input|output|human)_(\d+)$/i
                  );
                  if (match) {
                    const [, type, index] = match;
                    const targetIndex = parseInt(index) - 1; // Convert to 0-based index

                    // Find elements of the matching type
                    const elementsOfType = cy.nodes().filter((node: any) => {
                      const nodeType = node.data("type");
                      return (
                        nodeType &&
                        nodeType.toLowerCase() === type.toLowerCase()
                      );
                    });

                    // Get the element at the specified index
                    if (
                      elementsOfType.length > targetIndex &&
                      targetIndex >= 0
                    ) {
                      const targetElement = elementsOfType[targetIndex];
                      if (targetElement) {
                        affectedElement = cy.getElementById(targetElement.id());
                        console.log(
                          `Found element by type matching: ${type}[${targetIndex}] -> ${targetElement.id()}`
                        );
                      }
                    }
                  }
                }

                // Strategy 3: If still not found, try partial name matching
                if (affectedElement.length === 0) {
                  const matchingElements = cy.nodes().filter((node: any) => {
                    const nodeName = node.data("name") || node.data("label");
                    return (
                      nodeName &&
                      nodeName.toLowerCase().includes(affectedId.toLowerCase())
                    );
                  });

                  if (matchingElements.length > 0) {
                    const firstMatch = matchingElements.first();
                    affectedElement = cy.getElementById(firstMatch.id());
                    console.log(
                      `Found element by partial name matching: ${affectedId} -> ${firstMatch.id()}`
                    );
                  }
                }

                if (affectedElement.length > 0) {
                  console.log("Found affected element in graph, selecting it");

                  // Clear current selection
                  cy.elements().unselect();

                  // Select the affected element
                  affectedElement.select();

                  // Center the view on the affected element
                  cy.center(affectedElement);

                  // Optional: Zoom to the element for better visibility
                  cy.fit(affectedElement, 100); // 100px padding
                } else {
                  console.log(
                    "Affected element not found in graph:",
                    affectedId
                  );
                  console.log(
                    "Available graph element IDs:",
                    cy
                      .elements()
                      .map((el: any) => el.id())
                      .join(", ")
                  );
                }
              }
            }
          }}
          knowledgeGraph={kg}
          numberedLines={traceLines}
          highlightRanges={highlightRanges}
          showTraceTab={schemaCapabilities.supportsLineReferences}
          onShowInTrace={(ranges) => {
            console.log("Show in trace requested with ranges:", ranges);
            // Set the highlight ranges with a small delay to ensure tab transition is complete
            setTimeout(() => {
              setHighlightRanges(ranges);
            }, 100);
          }}
          onEntitySelect={handleEntitySelect}
        />
      </div>
    </div>
  );
};