frontend/src/components/features/connections/ConnectionsView.tsx

import React, { useState, useEffect } from "react";
import { Card, CardContent } from "@/components/ui/card";
import { Button } from "@/components/ui/button";

import {
  Dialog,
  DialogContent,
  DialogDescription,
  DialogFooter,
  DialogHeader,
  DialogTitle,
} from "@/components/ui/dialog";
import {
  Link,
  Download,
  Eye,
  Upload,
  Hash,
  FileText,
  Database,
  Calendar,
} from "lucide-react";
import { useModal } from "@/context/ModalContext";
import { useToast } from "@/hooks/use-toast";
import { useNavigation } from "@/context/NavigationContext";
import { useAgentGraph } from "@/context/AgentGraphContext";
import { api } from "@/lib/api";
import langfuseIcon from "@/static/langfuse.png";
import langsmithIcon from "@/static/langsmith.png";
import { PreprocessingMethodModal } from "@/components/shared/modals/PreprocessingMethodModal";

interface Connection {
  id: string;
  platform: string;
  status: string;
  connected_at: string;
  last_sync: string | null;
  projects: Array<{
    id: string;
    name: string;
    description?: string;
    created_at?: string;
  }>;
}

interface FetchedTrace {
  id: string;
  name: string;
  platform: string;
  fetched_at: string;
  generated_timestamp?: string;
  imported: boolean;
  data?: any;
  connection_id?: string;
}

export function ConnectionsView() {
  const [connections, setConnections] = useState<Connection[]>([]);
  const [fetchedTraces, setFetchedTraces] = useState<FetchedTrace[]>([]);
  const [selectedTraces, setSelectedTraces] = useState<Set<string>>(new Set());
  const [fetchingConnections, setFetchingConnections] = useState<Set<string>>(
    new Set()
  );
  const [selectedProjects, setSelectedProjects] = useState<Map<string, string>>(
    new Map()
  );
  const [deleteDialogOpen, setDeleteDialogOpen] = useState(false);
  const [connectionToDelete, setConnectionToDelete] =
    useState<Connection | null>(null);
  const [previewDialogOpen, setPreviewDialogOpen] = useState(false);
  const [previewTrace, setPreviewTrace] = useState<FetchedTrace | null>(null);
  const [preprocessingModalOpen, setPreprocessingModalOpen] = useState(false);
  const [tracesToImport, setTracesToImport] = useState<FetchedTrace[]>([]);
  const [isImporting, setIsImporting] = useState(false);
  const { openModal } = useModal();
  const { toast } = useToast();
  const { actions: navigationActions } = useNavigation();
  const { actions: agentGraphActions } = useAgentGraph();

  const loadConnections = async () => {
    try {
      const response = await api.observability.getConnections();
      setConnections(
        response.connections.map((conn: any) => ({
          ...conn,
          last_sync: null, // Add missing field
          projects: conn.projects || [], // Add projects field
        })) || []
      );

      // Load fetched traces for each connection using connection-specific API
      const allFetchedTraces: FetchedTrace[] = [];
      for (const conn of response.connections) {
        try {
          const fetchedResponse = await api.observability.getFetchedTracesByConnection(
            conn.id
          );
          const tracesWithConnectionId = fetchedResponse.traces.map((trace: any) => ({
            ...trace,
            connection_id: conn.id
          }));
          allFetchedTraces.push(...tracesWithConnectionId);
        } catch (error) {
          console.warn(
            `Failed to load fetched traces for connection ${conn.id}:`,
            error
          );
        }
      }
      setFetchedTraces(allFetchedTraces);
    } catch (error) {
      console.warn("Failed to load connections:", error);
      setConnections([]);
    }
  };

  useEffect(() => {
    loadConnections();
  }, []);

  useEffect(() => {
    const handleConnectionUpdate = () => {
      loadConnections();
    };

    window.addEventListener(
      "observability-connection-updated",
      handleConnectionUpdate
    );

    return () => {
      window.removeEventListener(
        "observability-connection-updated",
        handleConnectionUpdate
      );
    };
  }, []);

  // Auto-select first project for connections that don't have a project selected
  useEffect(() => {
    if (connections.length > 0) {
      const newSelected = new Map(selectedProjects);
      let hasChanges = false;
      
      connections.forEach(connection => {
        if (!selectedProjects.has(connection.id) && connection.projects && connection.projects.length > 0) {
          const firstProject = connection.projects[0];
          if (firstProject?.name) {
            newSelected.set(connection.id, firstProject.name);
            hasChanges = true;
          }
        }
      });
      
      if (hasChanges) {
        setSelectedProjects(newSelected);
      }
    }
  }, [connections, selectedProjects]);

  const handleAddConnection = () => {
    openModal("observability-connection", "Connect to AI Observability");
  };

  const handleEditConnection = (connection: Connection) => {
    openModal(
      "observability-connection",
      `Edit ${connection.platform} Connection`,
      {
        editConnection: connection,
      }
    );
  };

  const handleDeleteConnection = async (connection: Connection) => {
    setConnectionToDelete(connection);
    setDeleteDialogOpen(true);
  };

  const confirmDelete = async () => {
    if (!connectionToDelete) return;

    try {
      await api.observability.deleteConnection(connectionToDelete.id);
      loadConnections();
      setDeleteDialogOpen(false);
      setConnectionToDelete(null);
    } catch (error) {
      console.error("Failed to delete connection:", error);
    }
  };

  const cancelDelete = () => {
    setDeleteDialogOpen(false);
    setConnectionToDelete(null);
  };

  const handleFetchTraces = async (connectionId: string, projectName?: string) => {
    setFetchingConnections((prev) => new Set(prev).add(connectionId));
    try {
      // Get the connection to find project info
      const connection = connections.find(conn => conn.id === connectionId);
      const selectedProjectName = selectedProjects.get(connectionId) || projectName;
      let projectId: string | undefined;
      
      if (connection && selectedProjectName) {
        const selectedProject = connection.projects.find(p => p.name === selectedProjectName);
        projectId = selectedProject?.id;
      }

      // Trigger fetch from specific connection with project info
      await api.observability.fetchTracesByConnection(connectionId, 50, projectId, selectedProjectName);

      // Reload all traces to ensure consistency
      loadConnections();
    } catch (error) {
      console.error("Failed to fetch traces:", error);
    } finally {
      setFetchingConnections((prev) => {
        const newSet = new Set(prev);
        newSet.delete(connectionId);
        return newSet;
      });
    }
  };

  const handleDownloadTrace = async (trace: FetchedTrace) => {
    try {
      // Much simpler! Just need the trace ID
      const response = await api.observability.downloadTrace(trace.id);
      const jsonContent = JSON.stringify(response.data, null, 2);

      // Create blob and download
      const blob = new Blob([jsonContent], { type: "application/json" });
      const url = URL.createObjectURL(blob);
      const a = document.createElement("a");
      a.href = url;
      a.download = `${trace.name}_${trace.platform}.json`;
      document.body.appendChild(a);
      a.click();
      document.body.removeChild(a);
      URL.revokeObjectURL(url);

      toast({
        title: "Download Complete",
        description: `${trace.name} has been downloaded successfully.`,
      });
    } catch (error) {
      console.error("Error downloading trace:", error);
      toast({
        title: "Download Failed",
        description: "Failed to download trace. Please try again.",
        variant: "destructive",
      });
    }
  };

  const handleBulkImport = async () => {
    const selectedTracesToImport = fetchedTraces.filter(
      (t) => selectedTraces.has(t.id) && !t.imported
    );

    if (selectedTracesToImport.length === 0) return;

    // Store traces to import and show preprocessing method modal
    setTracesToImport(selectedTracesToImport);
    setPreprocessingModalOpen(true);
  };

  const handlePreprocessingMethodConfirm = async (preprocessingOptions: {
    max_char: number | null;
    topk: number;
    raw: boolean;
    hierarchy: boolean;
    replace: boolean;
  }) => {
    setIsImporting(true);
    setPreprocessingModalOpen(false);

    const successfulImports: string[] = [];
    const failedImports: string[] = [];

    try {
      // Process each trace by uploading the preprocessed data
      for (const trace of tracesToImport) {
        try {
          console.log(
            `Processing trace ${trace.id} from ${trace.platform} with preprocessing options:`, preprocessingOptions
          );

          // Find the connection for this trace
          const connection = connections.find(conn => conn.id === trace.connection_id);
          if (!connection) {
            throw new Error(`No connection found for trace ${trace.id}`);
          }

          // Call backend import API using connection-specific endpoint
          const result = await api.observability.importTracesByConnection(
            connection.id,
            [trace.id],
            preprocessingOptions
          );

          if (result.imported > 0) {
            successfulImports.push(trace.id);
            console.log(`Successfully processed trace ${trace.id}`);
          } else {
            failedImports.push(trace.id);
            console.log(`Failed to process trace ${trace.id}:`, result.errors);
          }
        } catch (error) {
          console.error(`Failed to process trace ${trace.id}:`, error);
          failedImports.push(trace.id);
        }
      }

      // Update UI state for successfully imported traces
      if (successfulImports.length > 0) {
        setFetchedTraces((prev) =>
          prev.map((t) =>
            successfulImports.includes(t.id) ? { ...t, imported: true } : t
          )
        );
      }

      // Clear selection
      setSelectedTraces(new Set());

      // Refresh traces list after successful import
      if (successfulImports.length > 0) {
        try {
          const tracesData = await api.traces.list();
          // Update traces in AgentGraphContext
          agentGraphActions.setTraces(
            Array.isArray(tracesData) ? tracesData : []
          );
        } catch (error) {
          console.error("Failed to refresh traces list:", error);
        }
      }

      // Show appropriate toast based on results
      if (successfulImports.length > 0 && failedImports.length === 0) {
        toast({
          title: "Traces Imported Successfully",
          description: `${successfulImports.length} trace${
            successfulImports.length > 1 ? "s" : ""
          } imported to your traces.`,
          variant: "default",
        });

        // Navigate to main page (traces section) only if all imports succeeded
        navigationActions.setCurrentSection("traces");
      } else if (successfulImports.length > 0 && failedImports.length > 0) {
        toast({
          title: "Partial Import Success",
          description: `${successfulImports.length} traces imported successfully, ${failedImports.length} failed.`,
          variant: "default",
        });
      } else {
        toast({
          title: "Import Failed",
          description:
            "Failed to import all selected traces. Please try again.",
          variant: "destructive",
        });
      }
    } catch (error) {
      console.error("Failed to import traces:", error);
      toast({
        title: "Import Failed",
        description: "Failed to import traces. Please try again.",
        variant: "destructive",
      });
    } finally {
      setIsImporting(false);
      setTracesToImport([]);
    }
  };

  const handlePreprocessingMethodCancel = () => {
    setPreprocessingModalOpen(false);
    setTracesToImport([]);
  };

  const handlePreviewTrace = (trace: FetchedTrace) => {
    setPreviewTrace(trace);
    setPreviewDialogOpen(true);
  };

  const getPlatformIcon = (platform: string) => {
    switch (platform.toLowerCase()) {
      case "langfuse":
        return (
          <img
            src={langfuseIcon}
            alt="Langfuse"
            className="h-6 w-6 object-contain"
          />
        );
      case "langsmith":
        return (
          <img
            src={langsmithIcon}
            alt="LangSmith"
            className="h-6 w-6 object-contain"
          />
        );
      default:
        return "AI";
    }
  };

  if (connections.length === 0) {
    return (
      <div className="flex-1 flex flex-col items-center justify-center p-8">
        <div className="text-center space-y-6 max-w-md">
          <div className="h-16 w-16 rounded-full bg-primary/10 flex items-center justify-center mx-auto">
            <Link className="h-8 w-8 text-primary" />
          </div>
          <div className="space-y-2">
            <h2 className="text-2xl font-bold">Connect AI Observability</h2>
            <p className="text-muted-foreground">
              Connect to your AI observability platforms to fetch and import
              traces
            </p>
          </div>
          <Button onClick={handleAddConnection} size="lg" className="gap-2">
            <Link className="h-4 w-4" />
            Add Connection
          </Button>
        </div>
      </div>
    );
  }

  return (
    <>
      <div className="flex-1 p-6 space-y-6 overflow-y-auto">
        {/* Header */}
        <div className="flex items-center justify-between">
          <div>
            <h1 className="text-2xl font-bold">Connections</h1>
            <p className="text-muted-foreground">
              Manage your AI observability platform connections
            </p>
          </div>
          <Button onClick={handleAddConnection} className="gap-2">
            <Link className="h-4 w-4" />
            Add Connection
          </Button>
        </div>

        {/* Platform Connection Cards */}
        <div className="space-y-6">
          {connections.map((connection) => {
            const isFetching = fetchingConnections.has(connection.id);
            const selectedProject = selectedProjects.get(connection.id);
            const connectionTraces = fetchedTraces.filter(
              (trace) => trace.connection_id === connection.id
            );
            const connectionSelectedTraces = new Set(
              Array.from(selectedTraces).filter((id) =>
                connectionTraces.some((trace) => trace.id === id)
              )
            );

            return (
              <Card key={connection.id} className="overflow-hidden">
                <CardContent className="p-6">
                  {/* Connection Header */}
                  <div className="flex items-center justify-between mb-4">
                    <div className="flex items-center gap-3">
                      <div className="h-10 w-10 rounded-lg bg-primary/10 flex items-center justify-center">
                        {getPlatformIcon(connection.platform)}
                      </div>
                      <div>
                        <h3 className="text-lg font-semibold capitalize">
                          {connection.platform === "langfuse"
                            ? "Langfuse"
                            : "LangSmith"}
                        </h3>
                        <p className="text-sm text-muted-foreground">
                          Connected{" "}
                          {new Date(
                            connection.connected_at
                          ).toLocaleDateString()}
                        </p>
                      </div>
                    </div>
                    <div className="flex items-center gap-2">
                      <Button
                        variant="outline"
                        size="sm"
                        onClick={() => handleEditConnection(connection)}
                      >
                        Edit
                      </Button>
                      <Button
                        variant="outline"
                        size="sm"
                        onClick={() => handleDeleteConnection(connection)}
                      >
                        Delete
                      </Button>
                    </div>
                  </div>

                  {/* Fetch Controls */}
                  <div className="space-y-4 mb-6">
                    <div className="flex items-center gap-3">
                      {/* Project Selection */}
                      {connection.projects &&
                        connection.projects.length > 0 && (
                          <div className="flex-1">
                            <label className="text-sm font-medium text-muted-foreground mb-2 block">
                              Select Project
                            </label>
                            <select
                              value={
                                selectedProject ||
                                connection.projects[0]?.name ||
                                ""
                              }
                              onChange={(e) => {
                                const newSelected = new Map(selectedProjects);
                                newSelected.set(connection.id, e.target.value);
                                setSelectedProjects(newSelected);
                              }}
                              className="w-full px-3 py-2 text-sm border border-gray-300 rounded-md bg-white"
                              disabled={isFetching}
                            >
                              {connection.projects.map((project) => (
                                <option key={project.id} value={project.name}>
                                  {project.name}
                                </option>
                              ))}
                            </select>
                          </div>
                        )}

                      {/* Fetch Button */}
                      <div className="flex-1">
                        <label className="text-sm font-medium text-muted-foreground mb-2 block">
                          Fetch Traces
                        </label>
                        <Button
                          onClick={() =>
                            handleFetchTraces(
                              connection.id,
                              selectedProject
                            )
                          }
                          disabled={isFetching}
                          className="w-full"
                        >
                          {isFetching ? (
                            <>
                              <div className="h-4 w-4 animate-spin rounded-full border-2 border-current border-t-transparent mr-2" />
                              Fetching...
                            </>
                          ) : (
                            <>
                              <Download className="h-4 w-4 mr-2" />
                              Fetch Traces
                            </>
                          )}
                        </Button>
                      </div>
                    </div>
                  </div>

                  {/* Connection Fetched Traces */}
                  {connectionTraces.length > 0 && (
                    <div className="space-y-4">
                      <div className="flex items-center justify-between">
                        <h4 className="text-md font-semibold">
                          Fetched Traces ({connectionTraces.length})
                        </h4>
                        <div className="flex items-center gap-2">
                          {/* Import All Button */}
                          <Button
                            size="sm"
                            variant="outline"
                            onClick={() => {
                              // Import all non-imported traces without changing selection
                              const tracesToImportAll = connectionTraces.filter(
                                (trace) => !trace.imported
                              );
                              
                              if (tracesToImportAll.length === 0) return;
                              
                              // Store traces to import and show preprocessing method modal
                              setTracesToImport(tracesToImportAll);
                              setPreprocessingModalOpen(true);
                            }}
                            disabled={connectionTraces.every(
                              (trace) => trace.imported
                            )}
                            title="Import all traces"
                          >
                            <Upload className="h-3 w-3 mr-1" />
                            Import All (
                            {connectionTraces.filter((t) => !t.imported).length})
                          </Button>

                          {/* Import Selected Button */}
                          <Button
                            size="sm"
                            onClick={handleBulkImport}
                            disabled={
                              connectionSelectedTraces.size === 0 ||
                              !Array.from(connectionSelectedTraces).some(
                                (id) =>
                                  !connectionTraces.find((t) => t.id === id)
                                    ?.imported
                              )
                            }
                            title="Import selected traces"
                          >
                            <Upload className="h-3 w-3 mr-1" />
                            Import Selected ({connectionSelectedTraces.size})
                          </Button>
                        </div>
                      </div>

                      <div className="border rounded-lg overflow-hidden">
                        <div className="overflow-x-auto">
                          <table className="w-full">
                            <thead>
                              <tr className="border-b bg-muted/50">
                                <th className="text-left p-3 font-medium">
                                  <input
                                    type="checkbox"
                                    className="mr-2"
                                    onChange={(e) => {
                                      const newSelected = new Set(
                                        selectedTraces
                                      );
                                      if (e.target.checked) {
                                        connectionTraces.forEach((trace) =>
                                          newSelected.add(trace.id)
                                        );
                                      } else {
                                        connectionTraces.forEach((trace) =>
                                          newSelected.delete(trace.id)
                                        );
                                      }
                                      setSelectedTraces(newSelected);
                                    }}
                                    checked={
                                      connectionTraces.length > 0 &&
                                      connectionTraces.every((trace) =>
                                        selectedTraces.has(trace.id)
                                      )
                                    }
                                  />
                                  Name
                                </th>
                                <th className="text-left p-3 font-medium">
                                  Fetch Time
                                </th>
                                <th className="text-left p-3 font-medium">
                                  Generated At
                                </th>
                                <th className="text-left p-3 font-medium">
                                  Actions
                                </th>
                              </tr>
                            </thead>
                            <tbody>
                              {connectionTraces.map((trace) => (
                                <tr
                                  key={trace.id}
                                  className="border-b hover:bg-muted/20"
                                >
                                  <td className="p-3">
                                    <div className="flex items-center gap-2">
                                      <input
                                        type="checkbox"
                                        checked={selectedTraces.has(trace.id)}
                                        onChange={(e) => {
                                          const newSelected = new Set(
                                            selectedTraces
                                          );
                                          if (e.target.checked) {
                                            newSelected.add(trace.id);
                                          } else {
                                            newSelected.delete(trace.id);
                                          }
                                          setSelectedTraces(newSelected);
                                        }}
                                      />
                                      <span className="font-medium">
                                        {trace.name}
                                      </span>
                                      {trace.imported && (
                                        <span className="text-xs bg-green-100 text-green-800 px-2 py-1 rounded">
                                          Imported
                                        </span>
                                      )}
                                    </div>
                                  </td>
                                  <td className="p-3 text-sm text-muted-foreground">
                                    {new Date(
                                      trace.fetched_at
                                    ).toLocaleString()}
                                  </td>
                                  <td className="p-3 text-sm text-muted-foreground">
                                    {trace.generated_timestamp
                                      ? new Date(
                                          trace.generated_timestamp
                                        ).toLocaleString()
                                      : new Date(
                                          trace.fetched_at
                                        ).toLocaleString()}
                                  </td>
                                  <td className="p-3">
                                    <div className="flex items-center gap-1">
                                      <Button
                                        variant="outline"
                                        size="sm"
                                        onClick={() =>
                                          handlePreviewTrace(trace)
                                        }
                                      >
                                        <Eye className="h-3 w-3" />
                                      </Button>
                                      <Button
                                        size="sm"
                                        onClick={() =>
                                          handleDownloadTrace(trace)
                                        }
                                        variant="outline"
                                        title="Download JSON"
                                      >
                                        <Download className="h-3 w-3" />
                                      </Button>
                                    </div>
                                  </td>
                                </tr>
                              ))}
                            </tbody>
                          </table>
                        </div>
                      </div>
                    </div>
                  )}
                </CardContent>
              </Card>
            );
          })}
        </div>
      </div>

      {/* Delete Confirmation Dialog */}
      <Dialog open={deleteDialogOpen} onOpenChange={setDeleteDialogOpen}>
        <DialogContent>
          <DialogHeader>
            <DialogTitle>Delete Connection</DialogTitle>
            <DialogDescription>
              Are you sure you want to delete the {connectionToDelete?.platform}{" "}
              connection? This action cannot be undone.
            </DialogDescription>
          </DialogHeader>
          <DialogFooter>
            <Button variant="outline" onClick={cancelDelete}>
              Cancel
            </Button>
            <Button variant="destructive" onClick={confirmDelete}>
              Delete
            </Button>
          </DialogFooter>
        </DialogContent>
      </Dialog>

      {/* Preview Trace Dialog */}
      <Dialog open={previewDialogOpen} onOpenChange={setPreviewDialogOpen}>
        <DialogContent className="max-w-4xl max-h-[80vh] overflow-hidden">
          <DialogHeader>
            <DialogTitle>Preview {previewTrace?.name}</DialogTitle>
            <DialogDescription>
              {previewTrace?.platform} trace data
            </DialogDescription>
          </DialogHeader>

          {/* Content Statistics */}
          {previewTrace && (
            <div className="px-4 py-2 bg-muted/50 rounded-lg border">
              <div className="flex flex-wrap items-center gap-4 text-sm text-muted-foreground">
                <span className="flex items-center gap-1 flex-shrink-0">
                  <Hash className="h-4 w-4" />
                  {JSON.stringify(
                    previewTrace.data,
                    null,
                    2
                  ).length.toLocaleString()}{" "}
                  characters
                </span>
                <span className="flex items-center gap-1 flex-shrink-0">
                  <FileText className="h-4 w-4" />
                  {JSON.stringify(previewTrace.data, null, 2)
                    .split("\n")
                    .length.toLocaleString()}{" "}
                  lines
                </span>
                <span className="flex items-center gap-1 flex-shrink-0">
                  <Database className="h-4 w-4" />
                  {(() => {
                    const content = JSON.stringify(previewTrace.data, null, 2);
                    const wordCount = content.trim()
                      ? content.trim().split(/\s+/).length
                      : 0;
                    return wordCount.toLocaleString();
                  })()}{" "}
                  words
                </span>
                <span className="flex items-center gap-1 flex-shrink-0">
                  <Calendar className="h-4 w-4" />
                  <span className="truncate max-w-40">
                    {previewTrace.generated_timestamp
                      ? new Date(
                          previewTrace.generated_timestamp
                        ).toLocaleString()
                      : new Date(previewTrace.fetched_at).toLocaleString()}
                  </span>
                </span>
              </div>
            </div>
          )}

          <div className="flex-1 overflow-hidden">
            <pre className="bg-muted p-4 rounded-md text-sm overflow-auto max-h-[50vh] whitespace-pre-wrap">
              {previewTrace ? JSON.stringify(previewTrace.data, null, 2) : ""}
            </pre>
          </div>
          <DialogFooter>
            <Button
              variant="outline"
              onClick={() => setPreviewDialogOpen(false)}
            >
              Close
            </Button>
          </DialogFooter>
        </DialogContent>
      </Dialog>

      {/* Preprocessing Method Selection Modal */}
      <PreprocessingMethodModal
        open={preprocessingModalOpen}
        onOpenChange={setPreprocessingModalOpen}
        onConfirm={handlePreprocessingMethodConfirm}
        onCancel={handlePreprocessingMethodCancel}
        isLoading={isImporting}
      />
    </>
  );
}