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

import React, { useState } from "react";
import { Button } from "@/components/ui/button";
import { Badge } from "@/components/ui/badge";
import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";

import { Collapsible, CollapsibleContent } from "@/components/ui/collapsible";
import {
  ChevronDown,
  ChevronUp,
  Eye,
  Play,
  FileText,
  BarChart3,
  Trash2,
  Activity,
  Settings,
  Scissors,
} from "lucide-react";
import { KnowledgeGraph } from "@/types";

import { useAgentGraph } from "@/context/AgentGraphContext";

interface KnowledgeGraphTreeProps {
  knowledgeGraphs: KnowledgeGraph[];
  selectedKgId?: string;
  onSelectKg?: (kgId: string) => void;
  _onViewKg?: (kgId: string) => void;
  onDeleteKg?: (kgId: string, name: string) => void;
  onReplayKg?: (
    kgId: string,
    traceId: string,
    processingRunId?: string
  ) => void;
  onViewSegment?: (
    traceId: string,
    startChar: number,
    endChar: number,
    windowIndex: number
  ) => void;
  currentTraceId?: string;
}

export const KnowledgeGraphTree: React.FC<KnowledgeGraphTreeProps> = ({
  knowledgeGraphs,
  selectedKgId: _selectedKgId,
  onSelectKg: _onSelectKg,
  _onViewKg,
  onDeleteKg,
  onReplayKg,
  onViewSegment,
  currentTraceId,
}) => {
  const [expandedKgs, setExpandedKgs] = useState<Set<string>>(new Set());
  const { actions } = useAgentGraph();

  // Filter to show final/main knowledge graphs with more permissive logic
  const finalKGs = knowledgeGraphs.filter((kg) => {
    const hasValidId = Boolean(kg.kg_id || kg.id);

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

    // Use the exact logic from stage_processor.js to determine if a KG is final
    const isFinal =
      kg.is_final === true ||
      (kg.window_index === null && kg.window_total !== null);

    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;
  });

  // Sort by creation timestamp, newest first (matching stage_processor.js)
  const sortedFinalKGs = finalKGs.sort((a, b) => {
    const dateA = a.created_at ? new Date(a.created_at) : new Date(0);
    const dateB = b.created_at ? new Date(b.created_at) : new Date(0);
    return dateB.getTime() - dateA.getTime();
  });

  const toggleExpanded = (kgId: string) => {
    const newExpanded = new Set(expandedKgs);
    if (newExpanded.has(kgId)) {
      newExpanded.delete(kgId);
    } else {
      newExpanded.add(kgId);
    }
    setExpandedKgs(newExpanded);
  };

  const getDisplayName = (kg: KnowledgeGraph) => {
    // Use system_name if available, otherwise fall back to filename
    if (kg.system_name) {
      return kg.system_name;
    }

    // Fallback to original filename logic
    const baseId = kg.kg_id || "unknown";
    let displayName = kg.filename || `Knowledge Graph #${baseId}`;

    // Strip the _knowledge_graph_{timestamp}_{uuid}.json part if present
    if (displayName.includes("_knowledge_graph_")) {
      displayName = displayName.split("_knowledge_graph_")[0] || displayName;
    }

    return displayName;
  };

  const getSystemSummary = (kg: KnowledgeGraph) => {
    return kg.system_summary || "";
  };

  if (sortedFinalKGs.length === 0) {
    return (
      <Card>
        <CardContent className="p-6 text-center">
          <div className="text-muted-foreground">
            <Activity className="h-8 w-8 mx-auto mb-2" />
            <p>No final knowledge graphs found for this trace.</p>
            <p className="text-sm mt-1">
              Generate a knowledge graph first to continue with pipeline
              processing.
            </p>
          </div>
        </CardContent>
      </Card>
    );
  }

  return (
    <div className="space-y-4 flex-1">
      {sortedFinalKGs.map((kg) => {
        const kgId = kg.kg_id || "unknown";
        const isSelected = false; // Selection removed
        const isExpanded = expandedKgs.has(kgId);
        const windowCount = kg.window_knowledge_graphs
          ? kg.window_knowledge_graphs.length
          : kg.window_total || 0;
        const formattedDate = kg.created_at
          ? new Date(kg.created_at).toLocaleString()
          : "Unknown";

        return (
          <Card
            key={kgId}
            className={`transition-all ${
              isSelected ? "ring-2 ring-primary" : ""
            }`}
          >
            <CardHeader className="pb-3">
              <div className="space-y-3">
                <div className="flex items-center justify-between">
                  <div className="flex items-center gap-2 flex-1 min-w-0">
                    <Activity className="h-5 w-5 text-primary" />
                    <CardTitle className="text-base truncate">
                      {getDisplayName(kg)}
                    </CardTitle>
                  </div>

                  {/* Delete button in top right */}
                  <Button
                    variant="ghost"
                    size="sm"
                    onClick={() => onDeleteKg?.(kgId, getDisplayName(kg))}
                    className="h-8 w-8 p-0 text-muted-foreground hover:text-destructive"
                    title="Delete knowledge graph"
                  >
                    <Trash2 className="h-4 w-4" />
                  </Button>
                </div>

                {/* System summary if available */}
                {getSystemSummary(kg) && (
                  <div className="text-sm text-muted-foreground bg-muted/30 rounded-md p-2 border-l-2 border-primary/20">
                    {getSystemSummary(kg)}
                  </div>
                )}

                <div className="flex flex-wrap gap-2 text-sm text-muted-foreground items-center">
                  <span>{formattedDate}</span>

                  {/* Processing metadata */}
                  {(kg as any).processing_metadata && (
                    <>
                      {(kg as any).processing_metadata.method_name &&
                        (kg as any).processing_metadata.method_name !==
                          "unknown" && (
                          <Badge
                            variant="outline"
                            className="text-xs h-5 flex items-center gap-1"
                          >
                            <Settings className="h-3 w-3" />
                            Method:{" "}
                            {(kg as any).processing_metadata.method_name}
                          </Badge>
                        )}
                      {(kg as any).processing_metadata.splitter_type &&
                        (kg as any).processing_metadata.splitter_type !==
                          "unknown" && (
                          <Badge
                            variant="outline"
                            className="text-xs h-5 flex items-center gap-1"
                          >
                            <Scissors className="h-3 w-3" />
                            Splitter:{" "}
                            {(kg as any).processing_metadata.splitter_type}
                          </Badge>
                        )}
                    </>
                  )}

                  {/* Entity and Relation Statistics - next to splitter info */}
                  {(kg.entity_count !== undefined ||
                    kg.relation_count !== undefined) && (
                    <div className="text-xs text-muted-foreground">
                      Entities: {kg.entity_count || 0} • Relations:{" "}
                      {kg.relation_count || 0}
                    </div>
                  )}
                </div>

                <div className="flex items-center gap-3">
                  {/* Primary actions */}
                  <div className="flex items-center gap-2">
                    <Button
                      variant="outline"
                      size="sm"
                      onClick={() => {
                        // Navigate directly to the knowledge graph visualizer
                        actions.setSelectedKnowledgeGraph(kg);
                        actions.setActiveView("kg-visualizer");
                      }}
                    >
                      <Eye className="h-4 w-4 mr-1" />
                      View
                    </Button>

                    <Button
                      variant={
                        kg.is_enriched || kg.is_perturbed || kg.is_analyzed
                          ? "default"
                          : "outline"
                      }
                      size="sm"
                      onClick={() => {
                        // Navigate to dedicated advanced processing page
                        actions.setSelectedKnowledgeGraph(kg);
                        actions.setActiveView("advanced-processing");
                      }}
                      title={`${
                        kg.is_enriched || kg.is_perturbed || kg.is_analyzed
                          ? "View and manage"
                          : "Run"
                      } advanced processing on this knowledge graph`}
                    >
                      <Activity className="h-4 w-4 mr-1" />
                      Process
                    </Button>
                  </div>

                  {/* Secondary actions */}
                  {windowCount > 0 && (
                    <div className="flex items-center gap-2 border-l pl-3">
                      <Button
                        variant="outline"
                        size="sm"
                        onClick={() => toggleExpanded(kgId)}
                        title={`${
                          isExpanded ? "Hide" : "Show"
                        } ${windowCount} window ${
                          windowCount === 1 ? "graph" : "graphs"
                        }`}
                      >
                        <BarChart3 className="h-4 w-4 mr-1" />
                        {isExpanded ? (
                          <ChevronUp className="h-4 w-4 mr-1" />
                        ) : (
                          <ChevronDown className="h-4 w-4 mr-1" />
                        )}
                        {windowCount} {windowCount === 1 ? "Window" : "Windows"}
                      </Button>

                      {windowCount > 0 && onReplayKg && currentTraceId && (
                        <Button
                          variant="outline"
                          size="sm"
                          onClick={() =>
                            onReplayKg(
                              kgId,
                              currentTraceId,
                              kg.processing_run_id
                            )
                          }
                          title="Replay creation process in temporal visualizer"
                        >
                          <Play className="h-4 w-4 mr-1" />
                          Replay
                        </Button>
                      )}
                    </div>
                  )}
                </div>
              </div>
            </CardHeader>

            {/* Window Knowledge Graphs */}
            {windowCount > 0 && (
              <Collapsible
                open={isExpanded}
                onOpenChange={() => toggleExpanded(kgId)}
              >
                <CollapsibleContent>
                  <CardContent className="pt-0">
                    <div className="border-t pt-4">
                      <div className="mb-6">
                        <h4 className="font-semibold text-lg flex items-center gap-2 mb-3">
                          <BarChart3 className="h-5 w-5 text-primary" />
                          Window Agent Graphs
                          <Badge
                            variant="secondary"
                            className="ml-1 font-normal"
                          >
                            {windowCount}
                          </Badge>
                        </h4>
                        <div className="bg-muted/30 rounded-lg p-4 border border-border/30">
                          <p className="text-sm text-muted-foreground flex items-start gap-2 leading-relaxed">
                            <Activity className="h-4 w-4 mt-0.5 text-primary/60 flex-shrink-0" />
                            <span>
                              These window agent graphs are automatically merged
                              to create the final agent graph shown above. Each
                              window represents a portion of the original trace
                              content, allowing you to examine specific sections
                              and their relationships in detail.
                            </span>
                          </p>
                        </div>
                      </div>

                      {/* Window Agent Graphs Table */}
                      <div className="space-y-3">
                        <div className="flex text-xs font-semibold text-muted-foreground/80 uppercase tracking-wider border-b border-border/30 pb-3">
                          <div className="w-32">Window</div>
                          <div className="w-44">Character Range</div>
                          <div className="w-44">Statistics</div>
                          <div className="w-52">Actions</div>
                        </div>

                        {kg.window_knowledge_graphs ? (
                          kg.window_knowledge_graphs.map((windowKg) => {
                            const entityCount = windowKg.entity_count || 0;
                            const relationCount = windowKg.relation_count || 0;

                            return (
                              <div
                                key={windowKg.kg_id}
                                className="flex py-4 border-b border-border/30 last:border-b-0 hover:bg-muted/30 transition-colors rounded-sm -mx-1 px-1"
                              >
                                <div className="w-32 flex items-center">
                                  <Badge
                                    variant="secondary"
                                    className="font-medium bg-primary/10 text-primary border-primary/20"
                                  >
                                    Window {(windowKg.window_index || 0) + 1}
                                  </Badge>
                                </div>

                                <div className="w-44 flex items-center text-sm font-mono text-muted-foreground">
                                  {windowKg.window_start_char?.toLocaleString() ||
                                    "N/A"}{" "}
                                  -{" "}
                                  {windowKg.window_end_char?.toLocaleString() ||
                                    "N/A"}
                                </div>

                                <div className="w-44 flex items-center py-1">
                                  <div className="text-xs text-muted-foreground">
                                    Entities: {entityCount} • Relations:{" "}
                                    {relationCount}
                                  </div>
                                </div>

                                <div className="flex-1 flex items-center gap-1 min-w-0 justify-end">
                                  <Button
                                    variant="outline"
                                    size="sm"
                                    className="flex-shrink-0 hover:bg-blue-50 hover:border-blue-200 transition-colors text-xs px-2"
                                    onClick={() => {
                                      // Navigate to the window knowledge graph visualizer
                                      const kgForVisualizer: KnowledgeGraph = {
                                        ...windowKg,
                                        created_at: new Date().toISOString(),
                                        filename:
                                          windowKg.filename ||
                                          `Window ${windowKg.window_index}`,
                                        status:
                                          (windowKg.status as KnowledgeGraph["status"]) ||
                                          "created",
                                      };
                                      actions.setSelectedKnowledgeGraph(
                                        kgForVisualizer
                                      );
                                      actions.setActiveView("kg-visualizer");
                                    }}
                                    title="View this window agent graph"
                                  >
                                    <Eye className="h-3 w-3 mr-1" />
                                    View
                                  </Button>

                                  <Button
                                    variant={
                                      (windowKg as any).is_enriched ||
                                      (windowKg as any).is_perturbed ||
                                      (windowKg as any).is_analyzed
                                        ? "default"
                                        : "outline"
                                    }
                                    size="sm"
                                    className={`flex-shrink-0 transition-colors text-xs px-2 ${
                                      (windowKg as any).is_enriched ||
                                      (windowKg as any).is_perturbed ||
                                      (windowKg as any).is_analyzed
                                        ? "bg-primary hover:bg-primary/90"
                                        : "hover:bg-purple-50 hover:border-purple-200"
                                    }`}
                                    onClick={() => {
                                      // Run advanced processing on this window knowledge graph
                                      const kgForProcessing: KnowledgeGraph = {
                                        ...windowKg,
                                        created_at: new Date().toISOString(),
                                        filename:
                                          windowKg.filename ||
                                          `Window ${windowKg.window_index}`,
                                        status:
                                          (windowKg.status as KnowledgeGraph["status"]) ||
                                          "created",
                                      };
                                      actions.setSelectedKnowledgeGraph(
                                        kgForProcessing
                                      );
                                      actions.setActiveView(
                                        "advanced-processing"
                                      );
                                    }}
                                    title={`${
                                      (windowKg as any).is_enriched ||
                                      (windowKg as any).is_perturbed ||
                                      (windowKg as any).is_analyzed
                                        ? "View and manage"
                                        : "Run"
                                    } advanced processing on this window agent graph`}
                                  >
                                    <Activity className="h-3 w-3 mr-1" />
                                    Process
                                  </Button>

                                  {onViewSegment && currentTraceId && (
                                    <Button
                                      variant="outline"
                                      size="sm"
                                      className="flex-shrink-0 hover:bg-green-50 hover:border-green-200 transition-colors text-xs px-2"
                                      onClick={() =>
                                        onViewSegment(
                                          currentTraceId,
                                          windowKg.window_start_char || 0,
                                          windowKg.window_end_char || 0,
                                          (windowKg.window_index || 0) + 1
                                        )
                                      }
                                      title="View the trace segment for this window"
                                    >
                                      <FileText className="h-3 w-3 mr-1" />
                                      Seg
                                    </Button>
                                  )}
                                </div>
                              </div>
                            );
                          })
                        ) : (
                          <div className="col-span-12 text-center py-4 text-muted-foreground">
                            <Activity className="h-4 w-4 mx-auto mb-1" />
                            <span className="text-sm">
                              No window agent graphs are available for this
                              trace. This may occur if the trace was processed
                              as a single unit without windowing.
                            </span>
                          </div>
                        )}
                      </div>
                    </div>
                  </CardContent>
                </CollapsibleContent>
              </Collapsible>
            )}
          </Card>
        );
      })}
    </div>
  );
};