frontend/src/components/shared/ModalSystem.tsx

import React from "react";
import {
  Dialog,
  DialogContent,
  DialogHeader,
  DialogTitle,
} from "@/components/ui/dialog";

import { ModalState } from "@/types";
import { TraceDetailsModal } from "@/components/shared/modals/TraceDetailsModal";
import { StageResultsModal } from "@/components/shared/modals/StageResultsModal";
import { KnowledgeGraphModal } from "@/components/shared/modals/KnowledgeGraphModal";
import { TraceSegmentModal } from "@/components/shared/modals/TraceSegmentModal";
import ExampleTraceModal from "./modals/ExampleTraceModal";
import { ObservabilityConnectionDialog } from "@/components/features/observability/ObservabilityConnectionDialog";
import { UploadDialog } from "@/components/features/upload/UploadDialog";
import { AuthModal } from "@/components/auth/AuthModal";

interface ModalSystemProps {
  modalState: ModalState;
  onClose: () => void;
}

export function ModalSystem({ modalState, onClose }: ModalSystemProps) {
  const { isOpen, type, title, data, config } = modalState;

  const getModalSize = () => {
    switch (config?.size) {
      case "sm":
        return "max-w-md";
      case "md":
        return "max-w-lg";
      case "lg":
        return "max-w-2xl";
      case "xl":
        return "max-w-4xl";
      case "full":
        return "max-w-[95vw] max-h-[95vh]";
      default:
        return "max-w-3xl";
    }
  };

  const renderModalContent = () => {
    switch (type) {
      case "trace-details":
        return <TraceDetailsModal data={data} onClose={onClose} />;
      case "stage-results":
        return <StageResultsModal data={data} onClose={onClose} />;
      case "knowledge-graph":
        return <KnowledgeGraphModal data={data} onClose={onClose} />;
      case "trace-segment":
        return <TraceSegmentModal data={data} onClose={onClose} />;
      case "example-traces":
        return <ExampleTraceModal data={data} onClose={onClose} />;
      case "upload-trace":
        return <UploadDialog open={isOpen} onOpenChange={onClose} />;
      case "observability-connection":
        return (
          <ObservabilityConnectionDialog
            open={isOpen}
            onOpenChange={onClose}
            editConnection={data?.editConnection}
          />
        );
      case "auth-login":
        return <AuthModal open={isOpen} onOpenChange={onClose} />;

      default:
        return (
          <div className="p-4 text-center text-muted-foreground">
            Unknown modal type: {type}
          </div>
        );
    }
  };

  // Some modals have their own Dialog wrapper
  const hasOwnDialog =
    type === "upload-trace" ||
    type === "observability-connection" ||
    type === "auth-login";

  if (hasOwnDialog) {
    return <>{renderModalContent()}</>;
  }

  return (
    <Dialog open={isOpen} onOpenChange={onClose}>
      <DialogContent className={`${getModalSize()} p-0`}>
        <DialogHeader className="px-6 py-4 border-b">
          <DialogTitle className="text-lg font-semibold">{title}</DialogTitle>
        </DialogHeader>
        <div className="overflow-auto max-h-[80vh]">{renderModalContent()}</div>
      </DialogContent>
    </Dialog>
  );
}