frontend/src/components/shared/modals/KnowledgeGraphModal.tsx

import React, { useState, useEffect } from "react";
import { Button } from "@/components/ui/button";
import { Badge } from "@/components/ui/badge";
import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
import {
  Database,
  Eye,
  Download,
  GitBranch,
  Play,
  Calendar,
  Hash,
  FileText,
  Settings,
  Network,
  Loader,
} from "lucide-react";
import {
  KnowledgeGraph,
  WindowKnowledgeGraph,
  Entity,
  Relation,
} from "@/types";
import { formatRelativeTime } from "@/lib/utils";
import { SimpleGraphVisualization } from "../../features/workspace/SimpleGraphVisualization";
import { api } from "@/lib/api";
import { useAgentGraph } from "@/context/AgentGraphContext";

interface KnowledgeGraphModalProps {
  data: {
    knowledgeGraph: KnowledgeGraph;
    windowGraphs?: WindowKnowledgeGraph[];
  };
  onClose: () => void;
}

interface GraphData {
  entities: Entity[];
  relations: Relation[];
  metadata?: Record<string, any>;
}

export function KnowledgeGraphModal({
  data,
  onClose,
}: KnowledgeGraphModalProps) {
  const { knowledgeGraph, windowGraphs = [] } = data;
  const [graphData, setGraphData] = useState<GraphData | null>(null);
  const [isLoading, setIsLoading] = useState(false);
  const [error, setError] = useState<string | null>(null);
  const { actions } = useAgentGraph();
  // Removed visualization type toggle - only using simple visualization in modal

  // Fetch the actual graph data when component mounts
  useEffect(() => {
    const fetchGraphData = async () => {
      if (!knowledgeGraph.id) return;

      setIsLoading(true);
      setError(null);

      try {
        const data = await api.knowledgeGraphs.getData(knowledgeGraph.id);
        setGraphData(data);
      } catch (err) {
        console.error("Failed to fetch knowledge graph data:", err);
        setError(
          err instanceof Error ? err.message : "Failed to load graph data"
        );
      } finally {
        setIsLoading(false);
      }
    };

    fetchGraphData();
  }, [knowledgeGraph.id]);

  const getStatusColor = (status: KnowledgeGraph["status"]) => {
    switch (status) {
      case "created":
        return "bg-green-500/10 text-green-700 border-green-200";
      case "enriched":
        return "bg-blue-500/10 text-blue-700 border-blue-200";
      case "perturbed":
        return "bg-purple-500/10 text-purple-700 border-purple-200";
      case "analyzed":
        return "bg-orange-500/10 text-orange-700 border-orange-200";
      case "causal":
        return "bg-red-500/10 text-red-700 border-red-200";
      default:
        return "bg-yellow-500/10 text-yellow-700 border-yellow-200";
    }
  };

  const getPipelineStageStatus = (status?: string) => {
    switch (status) {
      case "completed":
        return "bg-green-500/10 text-green-700";
      case "running":
        return "bg-blue-500/10 text-blue-700";
      case "error":
        return "bg-red-500/10 text-red-700";
      default:
        return "bg-gray-500/10 text-gray-700";
    }
  };

  const formatDisplayName = (filename?: string, id?: string) => {
    if (!filename) return `Knowledge Graph ${id?.slice(0, 8)}`;

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

    return filename;
  };

  return (
    <Tabs defaultValue="details" className="w-full">
      <TabsList className="grid w-full grid-cols-2">
        <TabsTrigger value="details">Details</TabsTrigger>
        <TabsTrigger value="visualization">
          <Network className="h-4 w-4 mr-2" />
          Visualization
        </TabsTrigger>
      </TabsList>

      <TabsContent value="details" className="space-y-6 p-6">
        <div className="space-y-6">
          {/* Header */}
          <div className="flex items-center justify-between">
            <div className="flex items-center gap-3">
              <Database className="h-6 w-6" />
              <div>
                <h2 className="text-xl font-semibold">
                  {formatDisplayName(
                    knowledgeGraph.filename,
                    knowledgeGraph.id
                  )}
                </h2>
                <p className="text-sm text-muted-foreground">
                  Created {formatRelativeTime(knowledgeGraph.created_at)}
                </p>
              </div>
            </div>
            <div className="flex items-center gap-2">
              <Badge
                variant="outline"
                className={getStatusColor(knowledgeGraph.status)}
              >
                {knowledgeGraph.status}
              </Badge>
              {knowledgeGraph.is_final && (
                <Badge variant="secondary">Final</Badge>
              )}
            </div>
          </div>

          {/* Knowledge Graph Information */}
          <Card>
            <CardHeader>
              <CardTitle className="flex items-center gap-2">
                <Database className="h-5 w-5" />
                Graph Information
              </CardTitle>
            </CardHeader>
            <CardContent>
              <div className="grid grid-cols-2 md:grid-cols-4 gap-4">
                <div className="flex items-center gap-2">
                  <Hash className="h-4 w-4 text-muted-foreground" />
                  <div>
                    <p className="text-sm font-medium">Entities</p>
                    <p className="text-sm text-muted-foreground">
                      {knowledgeGraph.entity_count?.toLocaleString() || "N/A"}
                    </p>
                  </div>
                </div>

                <div className="flex items-center gap-2">
                  <GitBranch className="h-4 w-4 text-muted-foreground" />
                  <div>
                    <p className="text-sm font-medium">Relations</p>
                    <p className="text-sm text-muted-foreground">
                      {knowledgeGraph.relation_count?.toLocaleString() || "N/A"}
                    </p>
                  </div>
                </div>

                <div className="flex items-center gap-2">
                  <Calendar className="h-4 w-4 text-muted-foreground" />
                  <div>
                    <p className="text-sm font-medium">Created</p>
                    <p className="text-sm text-muted-foreground">
                      {new Date(knowledgeGraph.created_at).toLocaleDateString()}
                    </p>
                  </div>
                </div>

                <div className="flex items-center gap-2">
                  <FileText className="h-4 w-4 text-muted-foreground" />
                  <div>
                    <p className="text-sm font-medium">Status</p>
                    <p className="text-sm text-muted-foreground">
                      {knowledgeGraph.status}
                    </p>
                  </div>
                </div>

                {knowledgeGraph.window_total && (
                  <div className="flex items-center gap-2">
                    <Settings className="h-4 w-4 text-muted-foreground" />
                    <div>
                      <p className="text-sm font-medium">Windows</p>
                      <p className="text-sm text-muted-foreground">
                        {knowledgeGraph.window_total} total
                      </p>
                    </div>
                  </div>
                )}

                {knowledgeGraph.processing_run_id && (
                  <div className="flex items-center gap-2">
                    <Hash className="h-4 w-4 text-muted-foreground" />
                    <div>
                      <p className="text-sm font-medium">Run ID</p>
                      <p className="text-sm text-muted-foreground font-mono">
                        {knowledgeGraph.processing_run_id.slice(0, 8)}
                      </p>
                    </div>
                  </div>
                )}
              </div>
            </CardContent>
          </Card>

          {/* Pipeline Status */}
          <Card>
            <CardHeader>
              <CardTitle className="flex items-center gap-2">
                <GitBranch className="h-5 w-5" />
                Pipeline Status
              </CardTitle>
            </CardHeader>
            <CardContent>
              <div className="space-y-3">
                <div className="flex items-center justify-between p-3 border rounded-lg">
                  <div className="flex items-center gap-2">
                    <Settings className="h-4 w-4" />
                    <span className="text-sm font-medium">
                      Prompt Reconstruction
                    </span>
                  </div>
                  <Badge
                    className={getPipelineStageStatus(
                      knowledgeGraph.is_enriched ? "completed" : "pending"
                    )}
                  >
                    {knowledgeGraph.is_enriched ? "Completed" : "Pending"}
                  </Badge>
                </div>

                <div className="flex items-center justify-between p-3 border rounded-lg">
                  <div className="flex items-center gap-2">
                    <Settings className="h-4 w-4" />
                    <span className="text-sm font-medium">
                      Perturbation Testing
                    </span>
                  </div>
                  <Badge
                    className={getPipelineStageStatus(
                      knowledgeGraph.is_perturbed ? "completed" : "pending"
                    )}
                  >
                    {knowledgeGraph.is_perturbed ? "Completed" : "Pending"}
                  </Badge>
                </div>

                <div className="flex items-center justify-between p-3 border rounded-lg">
                  <div className="flex items-center gap-2">
                    <Settings className="h-4 w-4" />
                    <span className="text-sm font-medium">Causal Analysis</span>
                  </div>
                  <Badge
                    className={getPipelineStageStatus(
                      knowledgeGraph.is_analyzed ? "completed" : "pending"
                    )}
                  >
                    {knowledgeGraph.is_analyzed ? "Completed" : "Pending"}
                  </Badge>
                </div>
              </div>
            </CardContent>
          </Card>

          {/* Window Agent Graphs */}
          {windowGraphs.length > 0 && (
            <Card>
              <CardHeader>
                <CardTitle className="flex items-center gap-2">
                  <Database className="h-5 w-5" />
                  Window Agent Graphs ({windowGraphs.length})
                </CardTitle>
              </CardHeader>
              <CardContent>
                <div className="space-y-3">
                  {windowGraphs.map((windowKg, index) => {
                    const entityCount = windowKg.entity_count || 0;
                    const relationCount = windowKg.relation_count || 0;

                    return (
                      <div
                        key={windowKg.kg_id}
                        className="flex items-center justify-between p-3 border rounded-lg"
                      >
                        <div className="flex items-center gap-3">
                          <div className="text-center min-w-[2rem]">
                            <span className="text-sm font-medium">
                              #{index + 1}
                            </span>
                          </div>
                          <div>
                            <p className="text-sm font-medium">
                              Characters{" "}
                              {windowKg.window_start_char?.toLocaleString()} -{" "}
                              {windowKg.window_end_char?.toLocaleString()}
                            </p>
                            <div className="flex items-center gap-3 text-xs text-muted-foreground">
                              <span>{entityCount} entities</span>
                              <span>{relationCount} relations</span>
                              <span>Window {index + 1}</span>
                            </div>
                          </div>
                        </div>

                        <div className="flex items-center gap-2">
                          <Button
                            variant="outline"
                            size="sm"
                            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",
                              };
                              onClose();
                              actions.setSelectedKnowledgeGraph(
                                kgForVisualizer
                              );
                              actions.setActiveView("kg-visualizer");
                            }}
                          >
                            <Eye className="h-3 w-3 mr-1" />
                            View
                          </Button>
                        </div>
                      </div>
                    );
                  })}
                </div>
              </CardContent>
            </Card>
          )}

          {/* Actions */}
          <div className="flex gap-2 pt-4 border-t">
            <Button variant="outline">
              <Download className="h-4 w-4 mr-2" />
              Download
            </Button>
            {knowledgeGraph.processing_run_id && windowGraphs.length > 1 && (
              <Button variant="outline">
                <Play className="h-4 w-4 mr-2" />
                Replay Creation
              </Button>
            )}
          </div>
        </div>
      </TabsContent>

      <TabsContent value="visualization" className="p-6">
        {isLoading ? (
          <div className="flex items-center justify-center py-12">
            <div className="text-center">
              <Loader className="h-8 w-8 animate-spin mx-auto mb-4" />
              <p className="text-sm text-muted-foreground">
                Loading graph data...
              </p>
            </div>
          </div>
        ) : error ? (
          <div className="text-center py-12">
            <div className="text-red-500 mb-4">Error: {error}</div>
            <Button variant="outline" onClick={() => window.location.reload()}>
              Retry
            </Button>
          </div>
        ) : (
          <div className="space-y-4">
            {/* Simple Visualization */}
            <div className="text-center">
              <h3 className="text-lg font-semibold">
                Agent Graph Visualization
              </h3>
              <p className="text-sm text-muted-foreground">
                For advanced interactive visualization, use the full-page
                visualizer
              </p>
            </div>

            <SimpleGraphVisualization
              knowledgeGraph={knowledgeGraph}
              width={550}
              height={400}
              graphData={graphData}
            />
          </div>
        )}
      </TabsContent>
    </Tabs>
  );
}