frontend/src/components/features/observability/ObservabilityConnectionDialog.tsx

import React, { useState } from "react";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { Label } from "@/components/ui/label";
import {
  Dialog,
  DialogContent,
  DialogDescription,
  DialogHeader,
  DialogTitle,
} from "@/components/ui/dialog";
import {
  Card,
  CardContent,
  CardDescription,
  CardHeader,
  CardTitle,
} from "@/components/ui/card";
import { ArrowLeft, Key, Globe } from "lucide-react";
import { useNotification } from "@/context/NotificationContext";
import { api } from "@/lib/api";

import langfuseIcon from "@/static/langfuse.png";
import langsmithIcon from "@/static/langsmith.png";

interface ObservabilityConnectionDialogProps {
  open: boolean;
  onOpenChange: (open: boolean) => void;
  editConnection?: {
    id: string;
    platform: string;
    status: string;
    connected_at: string;
    last_sync: string | null;
  };
}

type Platform = "langfuse" | "langsmith";

interface LangfuseForm {
  publicKey: string;
  secretKey: string;
  host: string;
}

interface LangSmithForm {
  apiKey: string;
}

export function ObservabilityConnectionDialog({
  open,
  onOpenChange,
  editConnection,
}: ObservabilityConnectionDialogProps) {
  const [selectedPlatform, setSelectedPlatform] = useState<Platform | null>(
    editConnection ? (editConnection.platform as Platform) : null
  );
  const [isConnecting, setIsConnecting] = useState(false);
  const [langfuseForm, setLangfuseForm] = useState<LangfuseForm>({
    publicKey: "",
    secretKey: "",
    host: "https://cloud.langfuse.com",
  });
  const [langsmithForm, setLangsmithForm] = useState<LangSmithForm>({
    apiKey: "",
  });

  const { showNotification } = useNotification();

  const handlePlatformSelect = (platform: Platform) => {
    setSelectedPlatform(platform);
    // Reset forms when switching platforms
    if (platform === "langfuse") {
      setLangfuseForm({
        publicKey: "",
        secretKey: "",
        host: "https://cloud.langfuse.com",
      });
    } else {
      setLangsmithForm({
        apiKey: "",
      });
    }
  };

  const handleBack = () => {
    setSelectedPlatform(null);
    setLangfuseForm({
      publicKey: "",
      secretKey: "",
      host: "https://cloud.langfuse.com",
    });
    setLangsmithForm({
      apiKey: "",
    });
  };

  const handleConnect = async () => {
    if (!selectedPlatform) return;

    // Get current form data based on platform
    let publicKey = "";
    let secretKey = "";
    let host = "";

    if (selectedPlatform === "langfuse") {
      if (!langfuseForm.publicKey.trim() || !langfuseForm.secretKey.trim()) {
        showNotification({
          type: "error",
          title: "Validation Error",
          message: "Please provide both public and secret keys",
        });
        return;
      }
      publicKey = langfuseForm.publicKey.trim();
      secretKey = langfuseForm.secretKey.trim();
      host = langfuseForm.host.trim();
    } else {
      if (!langsmithForm.apiKey.trim()) {
        showNotification({
          type: "error",
          title: "Validation Error",
          message: "Please provide API key",
        });
        return;
      }
      publicKey = langsmithForm.apiKey.trim();
      secretKey = "";
      host = "";
    }

    setIsConnecting(true);

    try {
      const connectionData = {
        platform: selectedPlatform,
        publicKey: publicKey,
        secretKey: secretKey,
        host: host || undefined,
      };

      if (editConnection) {
        // Update existing connection
        await api.observability.updateConnection(
          editConnection.id,
          connectionData
        );
        showNotification({
          type: "success",
          title: "Connection Updated",
          message: `Successfully updated ${
            selectedPlatform === "langfuse" ? "Langfuse" : "LangSmith"
          } connection`,
        });
      } else {
        // Create new connection
        await api.observability.connect(connectionData);
        showNotification({
          type: "success",
          title: "Connection Successful",
          message: `Successfully connected to ${
            selectedPlatform === "langfuse" ? "Langfuse" : "LangSmith"
          }`,
        });
      }

      // Notify ObservabilitySection to refresh connections
      window.dispatchEvent(new CustomEvent("observability-connection-updated"));

      onOpenChange(false);
      handleBack();
    } catch (error) {
      showNotification({
        type: "error",
        title: editConnection ? "Update Failed" : "Connection Failed",
        message:
          error instanceof Error
            ? error.message
            : `Failed to ${editConnection ? "update" : "connect to"} platform`,
      });
    } finally {
      setIsConnecting(false);
    }
  };

  const handleClose = () => {
    if (!isConnecting) {
      onOpenChange(false);
      handleBack();
    }
  };

  const isFormValid = () => {
    if (selectedPlatform === "langfuse") {
      return langfuseForm.publicKey.trim() && langfuseForm.secretKey.trim();
    } else {
      return langsmithForm.apiKey.trim();
    }
  };

  return (
    <Dialog open={open} onOpenChange={handleClose}>
      <DialogContent className="sm:max-w-md max-w-[90vw] w-full">
        <DialogHeader>
          <DialogTitle className="flex items-center gap-2">
            {selectedPlatform && !editConnection && (
              <Button
                variant="ghost"
                size="sm"
                onClick={handleBack}
                disabled={isConnecting}
              >
                <ArrowLeft className="h-4 w-4" />
              </Button>
            )}
            {editConnection
              ? `Edit ${
                  selectedPlatform === "langfuse" ? "Langfuse" : "LangSmith"
                } Connection`
              : "Connect to AI Observability"}
          </DialogTitle>
          <DialogDescription>
            {selectedPlatform
              ? `${editConnection ? "Update" : "Enter"} your ${
                  selectedPlatform === "langfuse" ? "Langfuse" : "LangSmith"
                } credentials ${
                  editConnection ? "to update the connection" : "to connect"
                }`
              : "Choose an AI observability platform to connect"}
          </DialogDescription>
        </DialogHeader>

        {!selectedPlatform ? (
          <div className="grid grid-cols-2 gap-4">
            <Card
              className="cursor-pointer hover:bg-accent transition-colors"
              onClick={() => handlePlatformSelect("langfuse")}
            >
              <CardHeader className="pb-3">
                <div className="w-8 h-8 rounded-md flex items-center justify-center mb-2">
                  <img
                    src={langfuseIcon}
                    alt="Langfuse"
                    className="h-8 w-8 object-contain"
                  />
                </div>
                <CardTitle className="text-lg">Langfuse</CardTitle>
                <CardDescription className="text-sm">
                  Open-source LLM observability
                </CardDescription>
              </CardHeader>
              <CardContent>
                <div className="text-xs text-muted-foreground">
                  Connect to Langfuse Cloud or self-hosted instance
                </div>
              </CardContent>
            </Card>

            <Card
              className="cursor-pointer hover:bg-accent transition-colors"
              onClick={() => handlePlatformSelect("langsmith")}
            >
              <CardHeader className="pb-3">
                <div className="w-8 h-8 rounded-md flex items-center justify-center mb-2">
                  <img
                    src={langsmithIcon}
                    alt="LangSmith"
                    className="h-8 w-8 object-contain"
                  />
                </div>
                <CardTitle className="text-lg">LangSmith</CardTitle>
                <CardDescription className="text-sm">
                  LangChain's observability platform
                </CardDescription>
              </CardHeader>
              <CardContent>
                <div className="text-xs text-muted-foreground">
                  Connect to LangSmith by LangChain
                </div>
              </CardContent>
            </Card>
          </div>
        ) : selectedPlatform === "langfuse" ? (
          <div className="space-y-4">
            <div className="space-y-2">
              <Label htmlFor="publicKey" className="flex items-center gap-2">
                <Key className="h-4 w-4" />
                Public Key
              </Label>
              <Input
                id="publicKey"
                type="text"
                placeholder="pk-lf-..."
                value={langfuseForm.publicKey}
                onChange={(e) =>
                  setLangfuseForm((prev) => ({
                    ...prev,
                    publicKey: e.target.value,
                  }))
                }
                disabled={isConnecting}
              />
            </div>

            <div className="space-y-2">
              <Label htmlFor="secretKey" className="flex items-center gap-2">
                <Key className="h-4 w-4" />
                Secret Key
              </Label>
              <Input
                id="secretKey"
                type="text"
                placeholder="sk-lf-..."
                value={langfuseForm.secretKey}
                onChange={(e) =>
                  setLangfuseForm((prev) => ({
                    ...prev,
                    secretKey: e.target.value,
                  }))
                }
                disabled={isConnecting}
              />
            </div>

            <div className="space-y-2">
              <Label htmlFor="host" className="flex items-center gap-2">
                <Globe className="h-4 w-4" />
                Host (optional)
              </Label>
              <Input
                id="host"
                type="url"
                placeholder="https://cloud.langfuse.com"
                value={langfuseForm.host}
                onChange={(e) =>
                  setLangfuseForm((prev) => ({ ...prev, host: e.target.value }))
                }
                disabled={isConnecting}
              />
              <div className="text-xs text-muted-foreground">
                Leave empty for Langfuse Cloud or enter your self-hosted URL
              </div>
            </div>

            <div className="flex gap-2 pt-4">
              <Button
                variant="outline"
                onClick={handleBack}
                disabled={isConnecting}
                className="flex-1"
              >
                Back
              </Button>
              <Button
                onClick={handleConnect}
                disabled={isConnecting || !isFormValid()}
                className="flex-1"
              >
                {isConnecting
                  ? editConnection
                    ? "Updating..."
                    : "Connecting..."
                  : editConnection
                  ? "Update"
                  : "Connect"}
              </Button>
            </div>
          </div>
        ) : (
          <div className="space-y-4">
            <div className="space-y-2">
              <Label htmlFor="apiKey" className="flex items-center gap-2">
                <Key className="h-4 w-4" />
                API Key
              </Label>
              <Input
                id="apiKey"
                type="text"
                placeholder="ls_..."
                value={langsmithForm.apiKey}
                onChange={(e) =>
                  setLangsmithForm((prev) => ({
                    ...prev,
                    apiKey: e.target.value,
                  }))
                }
                disabled={isConnecting}
              />
            </div>

            <div className="flex gap-2 pt-4">
              <Button
                variant="outline"
                onClick={handleBack}
                disabled={isConnecting}
                className="flex-1"
              >
                Back
              </Button>
              <Button
                onClick={handleConnect}
                disabled={isConnecting || !isFormValid()}
                className="flex-1"
              >
                {isConnecting
                  ? editConnection
                    ? "Updating..."
                    : "Connecting..."
                  : editConnection
                  ? "Update"
                  : "Connect"}
              </Button>
            </div>
          </div>
        )}
      </DialogContent>
    </Dialog>
  );
}