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

import React, { useState, useCallback } from "react";
import { Button } from "@/components/ui/button";
import { Badge } from "@/components/ui/badge";
import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
import {
  FileText,
  Database,
  Download,
  GitBranch,
  HardDrive,
  Hash,
  Activity,
  Trash2,
  AlertCircle,
  Plus,
} from "lucide-react";
import { Trace, KnowledgeGraph } from "@/types";
import { formatFileSize, formatRelativeTime } from "@/lib/utils";
import { getTraceStatusColor } from "@/lib/status-utils";
import { api } from "@/lib/api";
import { useModal } from "@/context/ModalContext";
import { useNotification } from "@/context/NotificationContext";
import { useNavigation } from "@/context/NavigationContext";
import { useAgentGraph } from "@/context/AgentGraphContext";
import { KnowledgeGraphTree } from "@/components/features/traces/KnowledgeGraphTree";
import { SplitterSelectionModal, SplitterType } from "./SplitterSelectionModal";
import { useTaskPolling } from "@/hooks/useTaskPolling";
import { ContextDocumentsSection } from "@/components/features/context/ContextDocumentsSection";

interface TraceDetailsModalProps {
  data: {
    trace: Trace;
    knowledgeGraphs?: KnowledgeGraph[];
  };
  onClose: () => void;
}

export function TraceDetailsModal({
  data,
  onClose: _onClose,
}: TraceDetailsModalProps) {
  const { trace, knowledgeGraphs = [] } = data;
  const { openModal } = useModal();
  const { showNotification } = useNotification();
  const { actions: navigationActions } = useNavigation();

  // 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]
  );

  const { actions } = useAgentGraph();
  const [isSplitterModalOpen, setIsSplitterModalOpen] = useState(false);
  const [isGenerating, setIsGenerating] = useState(false);
  const [generationProgress, setGenerationProgress] = useState(0);

  const formatTraceType = (type?: string) => {
    if (!type) return "Unknown";
    return type.replace(/[-_]/g, " ").replace(/\b\w/g, (l) => l.toUpperCase());
  };

  const handleViewKnowledgeGraph = (kgId: string) => {
    const kg = knowledgeGraphs.find((k) => k.kg_id === kgId || k.id === kgId);
    if (!kg) return;

    openModal(
      "knowledge-graph",
      `Knowledge Graph - ${
        kg.filename || `Agent Graph ${kg.kg_id?.slice(0, 8) || "unknown"}`
      }`,
      {
        knowledgeGraph: kg,
        windowGraphs: kg.window_knowledge_graphs || [],
      },
      {
        size: "xl",
        closable: true,
      }
    );
  };

  const handleViewSegment = async (
    traceId: string,
    startChar: number,
    endChar: number,
    windowIndex: number
  ) => {
    try {
      const segmentData = await api.traces.extractSegment(
        traceId,
        startChar,
        endChar
      );

      openModal(
        "trace-segment",
        `Trace Segment - Window ${windowIndex}`,
        {
          trace: trace,
          segment: {
            content: segmentData.content,
            startChar,
            endChar,
            windowIndex,
          },
        },
        {
          size: "xl",
          closable: true,
        }
      );
    } catch (error) {
      console.error("Error loading trace segment:", error);
      showSystemNotification({
        type: "error",
        title: "Failed to Load Segment",
        message: "Could not load the trace segment. Please try again.",
      });
    }
  };

  const handleDeleteKnowledgeGraph = async (kgId: string, name: string) => {
    const confirmMessage = `Are you sure you want to delete the knowledge graph "${name}"?

This will permanently remove:
• The knowledge graph and all entities/relationships
• All prompt reconstruction data
• All perturbation test results
• All causal analysis results

This action cannot be undone.`;

    if (!confirm(confirmMessage)) {
      return;
    }

    try {
      const response = await fetch(`/api/knowledge-graphs/${kgId}`, {
        method: "DELETE",
        headers: {
          "Content-Type": "application/json",
        },
      });

      if (!response.ok) {
        throw new Error(`API error: ${response.status} ${response.statusText}`);
      }

      const data = await response.json();
      console.log("Knowledge graph deleted:", data);

      // TODO: Add success notification and refresh knowledge graphs list
      alert("Knowledge graph deleted successfully");
    } catch (error) {
      console.error("Error deleting knowledge graph:", error);
      alert(
        `Error deleting knowledge graph: ${
          error instanceof Error ? error.message : String(error)
        }`
      );
    }
  };

  const handleReplayKg = (
    kgId: string,
    traceId: string,
    processingRunId?: string
  ) => {
    // Handle replay functionality
    console.log("Replay KG:", { kgId, traceId, processingRunId });
    // Implement actual replay logic
  };

  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 the trace data to get the updated knowledge graphs
      actions.setTraces([]);
      actions.setKnowledgeGraphs([]);
      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,
    modelName?: string,
    chunkingConfig?: { min_chunk_size?: number; max_chunk_size?: number }
  ) => {
    const finalMethodName = methodName || "production";
    console.log("TraceDetailsModal: Using method name:", finalMethodName);
    console.log("TraceDetailsModal: 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,
        finalMethodName,
        modelName || "gpt-5-mini",
        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
        const updatedTraces = await api.traces.list();
        actions.setTraces(updatedTraces);
        const updatedKGs = await api.knowledgeGraphs.list();
        actions.setKnowledgeGraphs(updatedKGs);
        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);
      showSystemNotification({
        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
  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;
    }
  };

  const handleViewContent = async () => {
    try {
      // Get numbered content for ContentReference compatibility
      const numberedContent = await api.traces.getNumberedContent(
        trace.trace_id
      );
      openModal(
        "trace-content",
        `Trace Content - ${trace.filename}`,
        {
          trace: trace,
          content: { content: numberedContent, type: "numbered" },
        },
        {
          size: "xl",
          closable: true,
        }
      );
    } catch (error) {
      console.error("Failed to load numbered content:", error);
      // Fallback to existing content if numbered content fails
      openModal(
        "trace-content",
        `Trace Content - ${trace.filename}`,
        {
          trace: trace,
          content: { content: trace.content || "", type: "raw" },
        },
        {
          size: "xl",
          closable: true,
        }
      );
    }
  };

  const handleDownloadTrace = async () => {
    try {
      // Get the full trace content instead of just metadata
      const traceContent = await api.traces.getContent(trace.trace_id);

      // Create blob with the actual trace content
      const blob = new Blob([traceContent], {
        type: "application/json",
      });
      const url = URL.createObjectURL(blob);

      const exportFileDefaultName = `${trace.filename || "trace"}.json`;

      const linkElement = document.createElement("a");
      linkElement.setAttribute("href", url);
      linkElement.setAttribute("download", exportFileDefaultName);
      linkElement.click();

      // Clean up the object URL
      URL.revokeObjectURL(url);

      console.log("Trace downloaded successfully:", trace.trace_id);
    } catch (error) {
      console.error("Failed to download trace:", error);
      alert("Failed to download trace. Please try again.");
    }
  };

  const handleDeleteTrace = async () => {
    const kgCount = knowledgeGraphs.length;
    const displayFilename =
      trace.filename.length > 50
        ? `${trace.filename.substring(0, 47)}...`
        : trace.filename;
    const confirmMessage = `Are you sure you want to delete the trace "${displayFilename}"?

This will permanently remove:
• The trace and all its content
• All associated knowledge graphs (${kgCount} found)
• All entities and relationships
• All prompt reconstruction data
• All perturbation test results
• All causal analysis results

This action cannot be undone.`;

    if (!confirm(confirmMessage)) {
      return;
    }

    try {
      await api.traces.delete(trace.trace_id);

      // Close modal and refresh page
      window.location.reload();
    } catch (error) {
      console.error("Error deleting trace:", error);
      alert(
        `Error deleting trace: ${
          error instanceof Error ? error.message : String(error)
        }`
      );
    }
  };

  // Calculate summary statistics - use consistent counting
  const finalKGsArray = knowledgeGraphs.filter(
    (kg) =>
      kg.is_final === true ||
      (kg.window_index === null && kg.window_total !== null)
  );
  const finalKGs = finalKGsArray.length; // Count of final KGs
  const totalKGs = finalKGs; // Use final KG count for consistency
  const completedPipelines = knowledgeGraphs.filter(
    (kg) => kg.is_analyzed
  ).length;

  return (
    <div className="space-y-6 p-6 max-h-[90vh] overflow-y-auto">
      {/* Trace Header */}
      <div className="flex items-center justify-between">
        <div className="flex items-center gap-3">
          <div className="p-2 rounded-lg bg-primary/10">
            <FileText className="h-6 w-6 text-primary" />
          </div>
          <div>
            <h2 className="text-xl font-semibold">{trace.filename}</h2>
            <p className="text-sm text-muted-foreground">
              Uploaded{" "}
              {formatRelativeTime(
                trace.upload_timestamp || trace.timestamp || ""
              )}
            </p>
          </div>
        </div>
        <Badge variant="outline" className={getTraceStatusColor(trace.status)}>
          {trace.status}
        </Badge>
      </div>

      {/* Quick Stats */}
      <div className="grid grid-cols-1 sm:grid-cols-2 lg:grid-cols-4 gap-4">
        <Card className="p-4 quick-stats-card">
          <div className="flex items-center gap-2">
            <HardDrive className="h-4 w-4 text-muted-foreground" />
            <div>
              <p className="text-2xl font-bold">
                {formatFileSize(trace.character_count || trace.size || 0)}
              </p>
              <p className="text-xs text-muted-foreground">File Size</p>
            </div>
          </div>
        </Card>

        <Card className="p-4 quick-stats-card">
          <div className="flex items-center gap-2">
            <GitBranch className="h-4 w-4 text-muted-foreground" />
            <div>
              <p className="text-2xl font-bold">{trace.turn_count || 0}</p>
              <p className="text-xs text-muted-foreground">Turns</p>
            </div>
          </div>
        </Card>

        <Card className="p-4 quick-stats-card">
          <div className="flex items-center gap-2">
            <Activity className="h-4 w-4 text-muted-foreground" />
            <div>
              <p className="text-2xl font-bold">{finalKGs}</p>
              <p className="text-xs text-muted-foreground">
                Final Agent Graphs
              </p>
            </div>
          </div>
        </Card>

        <Card className="p-4 quick-stats-card">
          <div className="flex items-center gap-2">
            <Database className="h-4 w-4 text-muted-foreground" />
            <div>
              <p className="text-2xl font-bold">{completedPipelines}</p>
              <p className="text-xs text-muted-foreground">Completed</p>
            </div>
          </div>
        </Card>
      </div>

      {/* Trace Information */}
      <Card>
        <CardHeader>
          <CardTitle className="flex items-center gap-2">
            <Database className="h-5 w-5" />
            Trace Information
          </CardTitle>
        </CardHeader>
        <CardContent className="space-y-6">
          {/* Core Information Grid */}
          <div className="grid grid-cols-1 md:grid-cols-2 gap-4">
            {/* Characters */}
            <Card className="p-4 bg-muted/30">
              <div className="flex items-center gap-3">
                <div className="p-2 rounded-lg bg-primary/10">
                  <Hash className="h-4 w-4 text-primary" />
                </div>
                <div className="flex-1">
                  <p className="text-xs text-muted-foreground uppercase tracking-wide font-medium">
                    Content Size
                  </p>
                  <p className="text-lg font-semibold">
                    {trace.character_count?.toLocaleString() || "N/A"} chars
                  </p>
                  {trace.turn_count && (
                    <p className="text-sm text-muted-foreground">
                      {trace.turn_count} turns
                    </p>
                  )}
                </div>
              </div>
            </Card>

            {/* Type & Upload Info */}
            <Card className="p-4 bg-muted/30">
              <div className="flex items-center gap-3">
                <div className="p-2 rounded-lg bg-primary/10">
                  <FileText className="h-4 w-4 text-primary" />
                </div>
                <div className="flex-1">
                  <p className="text-xs text-muted-foreground uppercase tracking-wide font-medium">
                    Type & Date
                  </p>
                  <p className="text-lg font-semibold">
                    {formatTraceType(trace.trace_type)}
                  </p>
                  <p className="text-sm text-muted-foreground">
                    {new Date(
                      trace.upload_timestamp || trace.timestamp || ""
                    ).toLocaleDateString()}
                  </p>
                </div>
              </div>
            </Card>
          </div>

          {/* Description */}
          {trace.description && (
            <div className="space-y-2">
              <p className="text-sm font-medium flex items-center gap-2">
                <FileText className="h-4 w-4 text-muted-foreground" />
                Description
              </p>
              <div className="p-4 rounded-lg bg-muted/30 border-l-4 border-primary">
                <p className="text-sm text-muted-foreground leading-relaxed">
                  {trace.description}
                </p>
              </div>
            </div>
          )}

          {/* Error Message */}
          {trace.error_message && (
            <div className="space-y-2">
              <p className="text-sm font-medium text-destructive flex items-center gap-2">
                <AlertCircle className="h-4 w-4" />
                Error Message
              </p>
              <div className="p-4 rounded-lg bg-destructive/10 border border-destructive/20">
                <p className="text-sm text-destructive font-mono">
                  {trace.error_message}
                </p>
              </div>
            </div>
          )}

          {/* Tags */}
          {trace.tags && trace.tags.length > 0 && (
            <div className="space-y-3">
              <p className="text-sm font-medium flex items-center gap-2">
                <Hash className="h-4 w-4 text-muted-foreground" />
                Tags
              </p>
              <div className="flex flex-wrap gap-2">
                {trace.tags.map((tag) => (
                  <Badge
                    key={tag}
                    variant="secondary"
                    className="text-xs px-3 py-1"
                  >
                    {tag}
                  </Badge>
                ))}
              </div>
            </div>
          )}
        </CardContent>
      </Card>

      {/* Context Documents Section */}
      <ContextDocumentsSection traceId={trace.trace_id} />

      {/* Knowledge Graphs using KnowledgeGraphTree */}
      {knowledgeGraphs.length > 0 ? (
        <Card>
          <CardHeader>
            <CardTitle className="flex items-center justify-between">
              <div className="flex items-center gap-2">
                <Activity className="h-5 w-5" />
                Agent Graphs ({totalKGs})
              </div>
              <div className="flex items-center gap-2">
                <Button
                  onClick={handleGenerateKG}
                  disabled={isGenerating || isPolling}
                  size="sm"
                  className="bg-primary hover:bg-primary/90"
                >
                  {isGenerating || isPolling ? (
                    <>
                      <div className="w-3 h-3 border-2 border-white border-t-transparent rounded-full animate-spin mr-2" />
                      {generationProgress > 0
                        ? `${Math.round(generationProgress)}%`
                        : "Generating..."}
                    </>
                  ) : (
                    <>
                      <Plus className="h-3 w-3 mr-2" />
                      Generate
                    </>
                  )}
                </Button>
                <div className="flex items-center gap-2 text-sm text-muted-foreground">
                  <span>{finalKGs} Final</span>
                  <span>•</span>
                  <span>{completedPipelines} Completed</span>
                </div>
              </div>
            </CardTitle>
          </CardHeader>
          <CardContent className="pt-0">
            <KnowledgeGraphTree
              knowledgeGraphs={knowledgeGraphs}
              _onViewKg={handleViewKnowledgeGraph}
              onDeleteKg={handleDeleteKnowledgeGraph}
              onReplayKg={handleReplayKg}
              onViewSegment={handleViewSegment}
              currentTraceId={trace.trace_id}
            />
          </CardContent>
        </Card>
      ) : (
        <Card>
          <CardContent className="p-6 text-center empty-state">
            <div className="text-muted-foreground">
              <Activity className="h-8 w-8 mx-auto mb-3 empty-state-icon" />
              <p className="font-medium mb-2">Ready to analyze this trace</p>
              <p className="text-sm text-muted-foreground mb-1">
                Transform your trace data into an interactive agent graph.
              </p>
              <p className="text-xs text-muted-foreground mb-4">
                Choose from different chunking strategies for optimal analysis.
              </p>
            </div>
            <Button
              onClick={handleGenerateKG}
              className="mt-2 bg-primary hover:bg-primary/90"
              size="sm"
            >
              <Plus className="h-4 w-4 mr-2" />
              Generate Agent Graph
            </Button>
          </CardContent>
        </Card>
      )}

      {/* Metadata */}
      {trace.metadata && Object.keys(trace.metadata).length > 0 && (
        <Card>
          <CardHeader>
            <CardTitle>Metadata</CardTitle>
          </CardHeader>
          <CardContent>
            <div className="space-y-2">
              {Object.entries(trace.metadata).map(([key, value]) => (
                <div key={key} className="flex justify-between">
                  <span className="text-sm font-medium">{key}:</span>
                  <span className="text-sm text-muted-foreground">
                    {typeof value === "object"
                      ? JSON.stringify(value, null, 2)
                      : String(value)}
                  </span>
                </div>
              ))}
            </div>
          </CardContent>
        </Card>
      )}

      {/* Actions */}
      <div className="flex gap-2 pt-4 border-t action-button-group">
        <Button onClick={handleViewContent}>
          <FileText className="h-4 w-4 mr-2" />
          View Content
        </Button>
        <Button variant="outline" onClick={handleGenerateKG}>
          <Activity className="h-4 w-4 mr-2" />
          Generate Knowledge Graph
        </Button>
        <Button variant="outline" onClick={handleDownloadTrace}>
          <Download className="h-4 w-4 mr-2" />
          Download Trace
        </Button>
        <Button variant="destructive" onClick={handleDeleteTrace}>
          <Trash2 className="h-4 w-4 mr-2" />
          Delete Trace
        </Button>
      </div>

      <SplitterSelectionModal
        open={isSplitterModalOpen}
        onOpenChange={setIsSplitterModalOpen}
        onConfirm={handleSplitterConfirm}
        isLoading={isGenerating}
      />
    </div>
  );
}