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

import React, { useState, useEffect, useRef, useCallback } from "react";
import { Button } from "@/components/ui/button";
import { Card, CardContent } from "@/components/ui/card";
import { Badge } from "@/components/ui/badge";
import {
  ArrowLeft,
  ZoomIn,
  ZoomOut,
  RotateCcw,
  AlertCircle,
} from "lucide-react";
// import { useAgentGraph } from "@/context/AgentGraphContext"; // Not used in this component
import {
  TemporalGraphData,
  TemporalControlState,
  TemporalNode,
  TemporalLink,
} from "@/types/temporal";
import {
  CytoscapeGraphCore,
  ElementSelectionCallbacks,
} from "@/lib/cytoscape-graph-core";
import { TemporalDataAdapter } from "@/lib/graph-data-adapters";
import { createTemporalGraphConfig } from "@/lib/graph-config-factory";
import { TemporalControls } from "../temporal/TemporalControls";
import { ElementInfoSidebar } from "../temporal/ElementInfoSidebar";

interface TemporalGraphVisualizerProps {
  temporalData: TemporalGraphData;
  onBack?: () => void;
}

export const TemporalGraphVisualizer: React.FC<
  TemporalGraphVisualizerProps
> = ({ temporalData, onBack }) => {
  // const { actions } = useAgentGraph();
  // SVG ref removed - using container for Cytoscape
  const containerRef = useRef<HTMLDivElement>(null);
  const cytoscapeRef = useRef<CytoscapeGraphCore | null>(null);
  const animationIntervalRef = useRef<NodeJS.Timeout | null>(null);
  const dataAdapter = useRef(new TemporalDataAdapter());
  const isInitializedRef = useRef(false);
  const isDestroyingRef = useRef(false);

  // State for temporal controls
  const [controlState, setControlState] = useState<TemporalControlState>({
    isPlaying: false,
    currentWindowIndex: 0,
    animationSpeed: 1.0,
    showingFullVersion: false,
    totalWindows: temporalData.windows.length,
  });

  const [loading, setLoading] = useState(true);
  const [error, setError] = useState<string | null>(null);
  const [graphStats, setGraphStats] = useState<{
    nodes: number;
    links: number;
    label: string;
  }>({ nodes: 0, links: 0, label: "Loading..." });

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

  // Load specific window data
  const loadWindow = useCallback(
    (windowIndex: number) => {
      console.log(
        `TemporalGraphVisualizer: loadWindow called with index ${windowIndex}`
      );

      if (
        !cytoscapeRef.current ||
        windowIndex < 0 ||
        windowIndex >= temporalData.windows.length
      ) {
        console.warn(
          `TemporalGraphVisualizer: Invalid window index ${windowIndex} or missing Cytoscape core`
        );
        return;
      }

      try {
        const window = temporalData.windows[windowIndex];
        console.log(
          `TemporalGraphVisualizer: Loading window ${windowIndex}:`,
          window
        );

        if (!window || (!window.entities && !window.relations)) {
          console.warn(
            `TemporalGraphVisualizer: Window ${windowIndex} has no entities or relations`
          );
          setError(`Window ${windowIndex + 1} contains no graph data`);
          return;
        }

        const adaptedData = dataAdapter.current.adapt({
          nodes: window.entities,
          links: window.relations,
          metadata: window.metadata,
        });
        cytoscapeRef.current.updateGraph(adaptedData, true);

        setGraphStats({
          nodes: window?.entities?.length || 0,
          links: window?.relations?.length || 0,
          label: `Window ${windowIndex + 1}`,
        });

        setControlState((prev) => ({
          ...prev,
          currentWindowIndex: windowIndex,
          showingFullVersion: false,
        }));

        console.log(
          `TemporalGraphVisualizer: Successfully loaded window ${windowIndex}`
        );
      } catch (err) {
        console.error(
          `TemporalGraphVisualizer: Error loading window ${windowIndex}:`,
          err
        );
        setError(
          `Failed to load window ${windowIndex + 1}: ${
            err instanceof Error ? err.message : String(err)
          }`
        );
      }
    },
    [temporalData]
  );

  // Initialize Cytoscape visualization
  useEffect(() => {
    console.log("TemporalGraphVisualizer: useEffect triggered", {
      containerRef: !!containerRef.current,
      temporalData: temporalData,
    });

    // Reset initialization state when data changes
    isInitializedRef.current = false;
    isDestroyingRef.current = false;

    // Add a small delay to ensure DOM is fully rendered
    const initializeVisualization = async () => {
      try {
        if (
          !containerRef.current ||
          isInitializedRef.current ||
          isDestroyingRef.current
        ) {
          console.log(
            "TemporalGraphVisualizer: Refs not ready or already initialized",
            {
              containerRef: !!containerRef.current,
              isInitialized: isInitializedRef.current,
              isDestroying: isDestroyingRef.current,
            }
          );
          return;
        }

        isInitializedRef.current = true;

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

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

        // For full-screen temporal visualization, use viewport dimensions if container is small
        if (width < 800 || height < 600) {
          const viewportWidth = window.innerWidth;
          const viewportHeight = window.innerHeight;

          // Account for header, controls, padding, and sidebar
          const headerHeight = 80; // Approximate header height
          const controlsHeight = 60; // Approximate controls height
          const padding = 32; // Total padding
          const sidebarWidth = 384; // Approximate sidebar width (w-96 = 384px)

          width = Math.max(width, viewportWidth - sidebarWidth - 50); // Account for sidebar and margin
          height = Math.max(
            height,
            viewportHeight - headerHeight - controlsHeight - padding
          );

          console.log(
            "TemporalGraphVisualizer: Using viewport-based dimensions",
            {
              viewportWidth,
              viewportHeight,
              sidebarWidth,
              calculatedWidth: width,
              calculatedHeight: height,
            }
          );
        }

        console.log("TemporalGraphVisualizer: Container dimensions", {
          width,
          height,
          clientWidth: container.clientWidth,
          clientHeight: container.clientHeight,
          offsetWidth: container.offsetWidth,
          offsetHeight: container.offsetHeight,
          scrollWidth: container.scrollWidth,
          scrollHeight: container.scrollHeight,
        });

        // If container has zero dimensions, try to get parent dimensions
        let finalWidth = width;
        let finalHeight = height;

        if (width === 0 || height === 0) {
          const parent = container.parentElement;
          if (parent) {
            finalWidth = parent.clientWidth || 800;
            finalHeight = parent.clientHeight || 600;
            console.log("TemporalGraphVisualizer: Using parent dimensions", {
              parentWidth: finalWidth,
              parentHeight: finalHeight,
            });
          }
        }

        if (finalWidth === 0 || finalHeight === 0) {
          console.warn(
            "TemporalGraphVisualizer: Container has zero dimensions, using fallbacks"
          );
          finalWidth = Math.max(finalWidth, 800);
          finalHeight = Math.max(finalHeight, 500); // Ensure minimum height
        }

        // Ensure minimum dimensions for usability
        finalWidth = Math.max(finalWidth, 400);
        finalHeight = Math.max(finalHeight, 400);

        const config = createTemporalGraphConfig({
          width: finalWidth,
          height: finalHeight,
        });

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

        // Create selection callbacks
        const selectionCallbacks: ElementSelectionCallbacks = {
          onNodeSelect: (node: TemporalNode) => {
            setSelectedElement(node);
            setSelectedElementType("node");
          },
          onLinkSelect: (link: TemporalLink) => {
            setSelectedElement(link);
            setSelectedElementType("link");
          },
          onClearSelection: () => {
            setSelectedElement(null);
            setSelectedElementType(null);
          },
        };

        // Check if container is already being used by Cytoscape
        if (containerRef.current.hasAttribute("data-cytoscape-active")) {
          console.warn(
            "Container already has active Cytoscape instance, skipping initialization"
          );
          return;
        }

        // Mark container as active
        containerRef.current.setAttribute("data-cytoscape-active", "true");

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

        // Load initial window
        if (temporalData.windows.length > 0) {
          console.log(
            "TemporalGraphVisualizer: Loading initial window",
            temporalData.windows[0]
          );
          loadWindow(0);
        } else {
          console.warn(
            "TemporalGraphVisualizer: No windows available in temporal data"
          );
          setError("No temporal windows available for visualization");
        }

        console.log(
          "TemporalGraphVisualizer: Initialization complete, setting loading to false"
        );
        setLoading(false);
      } catch (err) {
        console.error(
          "TemporalGraphVisualizer: Error initializing Cytoscape visualization:",
          err
        );
        setError(
          `Failed to initialize visualization: ${
            err instanceof Error ? err.message : String(err)
          }`
        );
        setLoading(false);
      }
    };

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

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

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

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

      return () => {
        clearInterval(retryInterval);
        isDestroyingRef.current = true;
        if (cytoscapeRef.current) {
          cytoscapeRef.current.destroy();
          cytoscapeRef.current = null;
        }
        if (animationIntervalRef.current) {
          clearInterval(animationIntervalRef.current);
          animationIntervalRef.current = null;
        }
        isInitializedRef.current = false;
      };
    }

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

    // Cleanup
    return () => {
      clearTimeout(timeoutId);
      isDestroyingRef.current = true;
      if (cytoscapeRef.current) {
        cytoscapeRef.current.destroy();
        cytoscapeRef.current = null;
      }
      if (animationIntervalRef.current) {
        clearInterval(animationIntervalRef.current);
        animationIntervalRef.current = null;
      }
      isInitializedRef.current = false;
    };
  }, [temporalData, loadWindow]);

  // Handle window resize for full-screen mode
  useEffect(() => {
    const handleResize = () => {
      if (cytoscapeRef.current && containerRef.current) {
        const container = containerRef.current;
        let width = container.clientWidth;
        let height = container.clientHeight;

        // Use viewport dimensions for better full-screen experience
        if (width < 800 || height < 600) {
          const viewportWidth = window.innerWidth;
          const viewportHeight = window.innerHeight;

          const headerHeight = 80;
          const controlsHeight = 60;
          const padding = 32;
          const sidebarWidth = 384; // Account for sidebar

          width = Math.max(width, viewportWidth - sidebarWidth - 50);
          height = Math.max(
            height,
            viewportHeight - headerHeight - controlsHeight - padding
          );
        }

        console.log("TemporalGraphVisualizer: Resizing to", { width, height });
        cytoscapeRef.current.resize(width, height);
      }
    };

    window.addEventListener("resize", handleResize);
    return () => window.removeEventListener("resize", handleResize);
  }, []);

  // Cleanup effect for component unmount
  useEffect(() => {
    return () => {
      // Final cleanup when component unmounts
      isDestroyingRef.current = true;
      if (cytoscapeRef.current) {
        cytoscapeRef.current.destroy();
        cytoscapeRef.current = null;
      }
      if (animationIntervalRef.current) {
        clearInterval(animationIntervalRef.current);
        animationIntervalRef.current = null;
      }
      isInitializedRef.current = false;
    };
  }, []);

  // Load full version
  const loadFullVersion = useCallback(() => {
    if (!cytoscapeRef.current || !temporalData.full_kg) {
      return;
    }

    try {
      console.log("Loading full version KG:", temporalData.full_kg);

      const adaptedData = dataAdapter.current.adapt({
        nodes: temporalData.full_kg.entities,
        links: temporalData.full_kg.relations,
        metadata: temporalData.full_kg.metadata,
      });
      cytoscapeRef.current.updateGraph(adaptedData, true);

      setGraphStats({
        nodes: temporalData.full_kg.entities?.length || 0,
        links: temporalData.full_kg.relations?.length || 0,
        label: "Full Version",
      });

      setControlState((prev) => ({
        ...prev,
        showingFullVersion: true,
      }));
    } catch (err) {
      console.error("Error loading full version:", err);
      setError("Failed to load full version");
    }
  }, [temporalData]);

  // Animation control functions
  const startAnimation = useCallback(() => {
    if (controlState.isPlaying || temporalData.windows.length < 2) {
      return;
    }

    setControlState((prev) => ({
      ...prev,
      isPlaying: true,
      currentWindowIndex: 0,
    }));
    loadWindow(0);

    animationIntervalRef.current = setInterval(() => {
      setControlState((prev) => {
        const nextIndex = prev.currentWindowIndex + 1;
        if (nextIndex >= temporalData.windows.length) {
          // Animation complete
          if (animationIntervalRef.current) {
            clearInterval(animationIntervalRef.current);
          }
          return { ...prev, isPlaying: false };
        } else {
          loadWindow(nextIndex);
          return { ...prev, currentWindowIndex: nextIndex };
        }
      });
    }, 2000 / controlState.animationSpeed);
  }, [
    controlState.isPlaying,
    controlState.animationSpeed,
    temporalData.windows.length,
    loadWindow,
  ]);

  const stopAnimation = useCallback(() => {
    if (animationIntervalRef.current) {
      clearInterval(animationIntervalRef.current);
      animationIntervalRef.current = null;
    }
    setControlState((prev) => ({ ...prev, isPlaying: false }));
  }, []);

  const handlePlay = () => startAnimation();
  const handlePause = () => stopAnimation();

  const handlePrevious = () => {
    if (controlState.currentWindowIndex > 0) {
      loadWindow(controlState.currentWindowIndex - 1);
    }
  };

  const handleNext = () => {
    if (controlState.currentWindowIndex < temporalData.windows.length - 1) {
      loadWindow(controlState.currentWindowIndex + 1);
    }
  };

  const handleWindowChange = (windowIndex: number) => {
    loadWindow(windowIndex);
  };

  const handleSpeedChange = (speed: number) => {
    setControlState((prev) => ({ ...prev, animationSpeed: speed }));
  };

  const handleToggleFullVersion = () => {
    if (controlState.showingFullVersion) {
      loadWindow(controlState.currentWindowIndex);
    } else {
      loadFullVersion();
    }
  };

  const handleReset = () => {
    stopAnimation();
    if (cytoscapeRef.current) {
      cytoscapeRef.current.resetZoom();
    }
    loadWindow(0);
  };

  const handleZoomIn = () => {
    if (cytoscapeRef.current) {
      cytoscapeRef.current.zoomIn();
    }
  };

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

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

  if (error) {
    return (
      <div className="flex-1 flex items-center justify-center">
        <Card className="w-full max-w-md">
          <CardContent className="pt-6">
            <div className="text-center">
              <AlertCircle className="h-12 w-12 text-destructive mx-auto mb-4" />
              <h3 className="text-lg font-semibold mb-2">
                Visualization Error
              </h3>
              <p className="text-muted-foreground mb-4">{error}</p>
              <Button onClick={() => window.location.reload()}>
                Try Again
              </Button>
            </div>
          </CardContent>
        </Card>
      </div>
    );
  }

  return (
    <div className="w-full h-full flex flex-col">
      {/* Header */}
      <div className="flex items-center justify-between px-6 py-4 border-b bg-background/50">
        <div className="flex items-center gap-4">
          <Button variant="ghost" size="sm" onClick={onBack}>
            <ArrowLeft className="h-4 w-4 mr-2" />
            Back
          </Button>
          <div>
            <h1 className="text-xl font-semibold">
              {temporalData.trace_title}
            </h1>
            <p className="text-sm text-muted-foreground">
              Temporal Knowledge Graph Visualization
            </p>
          </div>
        </div>
        <div className="flex items-center gap-2">
          <Badge variant="outline">
            {graphStats.nodes} nodes, {graphStats.links} relations
          </Badge>
          <Badge variant="secondary">{graphStats.label}</Badge>
        </div>
      </div>

      {/* Controls */}
      <div className="px-6 py-3 border-b bg-background/30">
        <TemporalControls
          state={controlState}
          onPlay={handlePlay}
          onPause={handlePause}
          onPrevious={handlePrevious}
          onNext={handleNext}
          onWindowChange={handleWindowChange}
          onSpeedChange={handleSpeedChange}
          onToggleFullVersion={handleToggleFullVersion}
          onReset={handleReset}
          hasFullVersion={temporalData.has_full_version}
          disabled={loading}
        />
      </div>

      {/* Visualization */}
      <div className="flex-1 flex min-h-0">
        <div className="flex-1 relative p-4">
          <div
            ref={containerRef}
            key={`temporal-graph-${temporalData.trace_id || "default"}`}
            className="w-full h-full bg-background border rounded-lg relative overflow-hidden shadow-sm"
          >
            {/* Cytoscape will render here - no SVG needed */}

            {/* Loading Overlay */}
            {loading && (
              <div className="absolute inset-0 bg-background/80 backdrop-blur-sm flex items-center justify-center">
                <div className="text-center">
                  <div className="w-8 h-8 border-4 border-primary border-t-transparent rounded-full animate-spin mx-auto mb-4" />
                  <p className="text-muted-foreground">
                    Loading temporal visualization...
                  </p>
                </div>
              </div>
            )}

            {/* Zoom Controls */}
            <div className="absolute top-4 right-4 flex flex-col gap-2">
              <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>

        {/* Element Info Sidebar */}
        <ElementInfoSidebar
          selectedElement={selectedElement}
          selectedElementType={selectedElementType}
          currentData={
            cytoscapeRef.current?.getCurrentData() || { nodes: [], links: [] }
          }
        />
      </div>
    </div>
  );
};