viewer/src/contexts/UrdfContext.tsx

import React, {
  createContext,
  useState,
  useCallback,
  ReactNode,
  useRef,
  useEffect,
} from "react";
import { toast } from "sonner";
import { UrdfProcessor, readUrdfFileContent } from "@/lib/UrdfDragAndDrop";
import { UrdfFileModel } from "@/lib/types";

// Define the result interface for URDF detection
interface UrdfDetectionResult {
  hasUrdf: boolean;
  modelName?: string;
}

// Define the context type
export type UrdfContextType = {
  urdfProcessor: UrdfProcessor | null;
  registerUrdfProcessor: (processor: UrdfProcessor) => void;
  onUrdfDetected: (
    callback: (result: UrdfDetectionResult) => void
  ) => () => void;
  processUrdfFiles: (
    files: Record<string, File>,
    availableModels: string[]
  ) => Promise<void>;
  urdfBlobUrls: Record<string, string>;
  alternativeUrdfModels: string[];
  isSelectionModalOpen: boolean;
  setIsSelectionModalOpen: (isOpen: boolean) => void;
  urdfModelOptions: UrdfFileModel[];
  selectUrdfModel: (model: UrdfFileModel) => void;

  isDefaultModel: boolean;
  setIsDefaultModel: (isDefault: boolean) => void;
  resetToDefaultModel: () => void;
  urdfContent: string | null;
};

// Create the context
export const UrdfContext = createContext<UrdfContextType | undefined>(
  undefined
);

// Props for the provider component
interface UrdfProviderProps {
  children: ReactNode;
}

export const UrdfProvider: React.FC<UrdfProviderProps> = ({ children }) => {
  // State for URDF processor
  const [urdfProcessor, setUrdfProcessor] = useState<UrdfProcessor | null>(
    null
  );

  // State for blob URLs (replacing window.urdfBlobUrls)
  const [urdfBlobUrls, setUrdfBlobUrls] = useState<Record<string, string>>({});

  // State for alternative models (replacing window.alternativeUrdfModels)
  const [alternativeUrdfModels, setAlternativeUrdfModels] = useState<string[]>(
    []
  );

  // State for the URDF selection modal
  const [isSelectionModalOpen, setIsSelectionModalOpen] = useState(false);
  const [urdfModelOptions, setUrdfModelOptions] = useState<UrdfFileModel[]>([]);

  // New state for centralized robot data management
  const [isDefaultModel, setIsDefaultModel] = useState(true);
  const [urdfContent, setUrdfContent] = useState<string | null>(null);

  // Fetch the default URDF content when the component mounts
  useEffect(() => {
    // Only fetch if we don't have content and we're using the default model
    if (isDefaultModel && !urdfContent) {
      const fetchDefaultUrdf = async () => {
        try {
          // Path to the default T12 URDF file
          const defaultUrdfPath = "/urdf/T12/urdf/T12.URDF";

          // Fetch the URDF content
          const response = await fetch(defaultUrdfPath);

          if (!response.ok) {
            throw new Error(
              `Failed to fetch default URDF: ${response.statusText}`
            );
          }

          const defaultUrdfContent = await response.text();
          console.log(
            `📄 Default URDF content loaded, length: ${defaultUrdfContent.length} characters`
          );

          // Set the URDF content in state
          setUrdfContent(defaultUrdfContent);
        } catch (error) {
          console.error("❌ Error loading default URDF content:", error);
        }
      };

      fetchDefaultUrdf();
    }
  }, [isDefaultModel, urdfContent]);

  // Reference for callbacks
  const urdfCallbacksRef = useRef<((result: UrdfDetectionResult) => void)[]>(
    []
  );

  // Reset to default model
  const resetToDefaultModel = useCallback(() => {
    setIsDefaultModel(true);
    setUrdfContent(null);

    toast.info("Switched to default model", {
      description: "The default T12 robot model is now displayed.",
    });
  }, []);

  // Register a callback for URDF detection
  const onUrdfDetected = useCallback(
    (callback: (result: UrdfDetectionResult) => void) => {
      urdfCallbacksRef.current.push(callback);

      return () => {
        urdfCallbacksRef.current = urdfCallbacksRef.current.filter(
          (cb) => cb !== callback
        );
      };
    },
    []
  );

  // Register a URDF processor
  const registerUrdfProcessor = useCallback((processor: UrdfProcessor) => {
    setUrdfProcessor(processor);
  }, []);

  // Internal function to notify callbacks and update central state
  const notifyUrdfCallbacks = useCallback(
    (result: UrdfDetectionResult) => {
      console.log("📣 Notifying URDF callbacks with result:", result);

      // Update our internal state based on the result
      if (result.hasUrdf) {
        // Always ensure we set isDefaultModel to false when we have a URDF
        setIsDefaultModel(false);

        if (result.modelName) {
          setUrdfContent(result.modelName);
        }

        // Set description if available
        if (result.modelName) {
          setUrdfContent(
            "A detailed 3D model of a robotic system with articulated joints and components."
          );
        }
      } else {
        // If no URDF, reset to default
        resetToDefaultModel();
      }

      // Call all registered callbacks
      urdfCallbacksRef.current.forEach((callback) => callback(result));
    },
    [resetToDefaultModel]
  );

  // Helper function to process the selected URDF model
  const processSelectedUrdf = useCallback(
    async (model: UrdfFileModel) => {
      if (!urdfProcessor) return;

      // Find the file in our files record
      const files = Object.values(urdfBlobUrls)
        .filter((url) => url === model.blobUrl)
        .map((url) => {
          const path = Object.keys(urdfBlobUrls).find(
            (key) => urdfBlobUrls[key] === url
          );
          return path ? { path, url } : null;
        })
        .filter((item) => item !== null);

      if (files.length === 0) {
        console.error("❌ Could not find file for selected URDF model");
        return;
      }

      // Show a toast notification that we're parsing the URDF
      const parsingToast = toast.loading("Analyzing URDF model...", {
        description: "Extracting robot information",
        duration: 10000, // Long duration since we'll dismiss it manually
      });

      try {
        // Get the file from our record
        const filePath = files[0]?.path;
        if (!filePath || !urdfBlobUrls[filePath]) {
          throw new Error("File not found in records");
        }

        // Get the actual File object
        const response = await fetch(model.blobUrl);
        const blob = await response.blob();
        const file = new File(
          [blob],
          filePath.split("/").pop() || "model.urdf",
          {
            type: "application/xml",
          }
        );

        // Read the URDF content
        const urdfContent = await readUrdfFileContent(file);

        console.log(
          `📏 URDF content read, length: ${urdfContent.length} characters`
        );

        // Store the URDF content in state
        setUrdfContent(urdfContent);

        // Dismiss the toast
        toast.dismiss(parsingToast);

        // Always set isDefaultModel to false when processing a custom URDF
        setIsDefaultModel(false);
      } catch (error) {
        // Error case
        console.error("❌ Error processing selected URDF:", error);
        toast.dismiss(parsingToast);
        toast.error("Error analyzing URDF", {
          description: `Error: ${
            error instanceof Error ? error.message : String(error)
          }`,
          duration: 3000,
        });

        // Keep showing the custom model even if parsing failed
        // No need to reset to default unless user explicitly chooses to
      }
    },
    [urdfBlobUrls, urdfProcessor]
  );

  // Function to handle selecting a URDF model from the modal
  const selectUrdfModel = useCallback(
    (model: UrdfFileModel) => {
      if (!urdfProcessor) {
        console.error("❌ No URDF processor available");
        return;
      }

      console.log(`🤖 Selected model: ${model.name || model.path}`);

      // Close the modal
      setIsSelectionModalOpen(false);

      // Extract model name
      const modelName =
        model.name ||
        model.path
          .split("/")
          .pop()
          ?.replace(/\.urdf$/i, "") ||
        "Unknown";

      // Load the selected URDF model
      urdfProcessor.loadUrdf(model.blobUrl);

      // Update our state immediately even before parsing
      setIsDefaultModel(false);

      // Show a toast notification that we're loading the model
      toast.info(`Loading model: ${modelName}`, {
        description: "Preparing 3D visualization",
        duration: 2000,
      });

      // Notify callbacks about the selection before parsing
      notifyUrdfCallbacks({
        hasUrdf: true,
        modelName,
      });

      // Try to parse the model - this will update the UI when complete
      processSelectedUrdf(model);
    },
    [urdfProcessor, notifyUrdfCallbacks, processSelectedUrdf]
  );

  // Process URDF files - moved from DragAndDropContext
  const processUrdfFiles = useCallback(
    async (files: Record<string, File>, availableModels: string[]) => {
      // Clear previous blob URLs to prevent memory leaks
      Object.values(urdfBlobUrls).forEach(URL.revokeObjectURL);
      setUrdfBlobUrls({});
      setAlternativeUrdfModels([]);
      setUrdfModelOptions([]);

      try {
        // Check if we have any URDF files
        if (availableModels.length > 0 && urdfProcessor) {
          console.log(
            `🤖 Found ${availableModels.length} URDF models:`,
            availableModels
          );

          // Create blob URLs for all models
          const newUrdfBlobUrls: Record<string, string> = {};
          availableModels.forEach((path) => {
            if (files[path]) {
              newUrdfBlobUrls[path] = URL.createObjectURL(files[path]);
            }
          });
          setUrdfBlobUrls(newUrdfBlobUrls);

          // Save alternative models for reference
          setAlternativeUrdfModels(availableModels);

          // Create model options for the selection modal
          const modelOptions: UrdfFileModel[] = availableModels.map((path) => {
            const fileName = path.split("/").pop() || "";
            const modelName = fileName.replace(/\.urdf$/i, "");
            return {
              path,
              blobUrl: newUrdfBlobUrls[path],
              name: modelName,
            };
          });

          setUrdfModelOptions(modelOptions);

          // If there's only one model, use it directly
          if (availableModels.length === 1) {
            // Extract model name from the URDF file
            const fileName = availableModels[0].split("/").pop() || "";
            const modelName = fileName.replace(/\.urdf$/i, "");
            console.log(`📄 Using model: ${modelName} (${fileName})`);

            // Use the blob URL instead of the file path
            const blobUrl = newUrdfBlobUrls[availableModels[0]];
            if (blobUrl) {
              console.log(`🔗 Using blob URL for URDF: ${blobUrl}`);
              urdfProcessor.loadUrdf(blobUrl);

              // Immediately update model state
              setIsDefaultModel(false);

              // Process the URDF file for parsing
              if (files[availableModels[0]]) {
                console.log(
                  "📄 Reading URDF content for edge function parsing..."
                );

                // Show a toast notification that we're parsing the URDF
                const parsingToast = toast.loading("Analyzing URDF model...", {
                  description: "Extracting robot information",
                  duration: 10000, // Long duration since we'll dismiss it manually
                });

                try {
                  const urdfContent = await readUrdfFileContent(
                    files[availableModels[0]]
                  );

                  console.log(
                    `📏 URDF content read, length: ${urdfContent.length} characters`
                  );

                  // Store the URDF content in state
                  setUrdfContent(urdfContent);

                  // Dismiss the parsing toast
                  toast.dismiss(parsingToast);
                } catch (parseError) {
                  console.error("❌ Error parsing URDF:", parseError);
                  toast.dismiss(parsingToast);
                  toast.error("Error analyzing URDF", {
                    description: `Error: ${
                      parseError instanceof Error
                        ? parseError.message
                        : String(parseError)
                    }`,
                    duration: 3000,
                  });

                  // Still notify callbacks without parsed data
                  notifyUrdfCallbacks({
                    hasUrdf: true,
                    modelName,
                  });
                }
              } else {
                console.error(
                  "❌ Could not find file for URDF model:",
                  availableModels[0]
                );
                console.log("📦 Available files:", Object.keys(files));

                // Still notify callbacks without parsed data
                notifyUrdfCallbacks({
                  hasUrdf: true,
                  modelName,
                });
              }
            } else {
              console.warn(
                `⚠️ No blob URL found for ${availableModels[0]}, using path directly`
              );
              urdfProcessor.loadUrdf(availableModels[0]);

              // Update the state even without a blob URL
              setIsDefaultModel(false);
            }
          } else {
            // Multiple URDF files found, show selection modal
            console.log(
              "📋 Multiple URDF files found, showing selection modal"
            );
            setIsSelectionModalOpen(true);

            // Notify that URDF files are available but selection is needed
            notifyUrdfCallbacks({
              hasUrdf: true,
              modelName: "Multiple models available",
            });
          }
        } else {
          console.warn(
            "❌ No URDF models found in dropped files or no processor available"
          );
          notifyUrdfCallbacks({ hasUrdf: false });

          // Reset to default model when no URDF files are found
          resetToDefaultModel();

          toast.error("No URDF file found", {
            description: "Please upload a folder containing a .urdf file.",
            duration: 3000,
          });
        }
      } catch (error) {
        console.error("❌ Error processing URDF files:", error);
        toast.error("Error processing files", {
          description: `Error: ${
            error instanceof Error ? error.message : String(error)
          }`,
          duration: 3000,
        });

        // Reset to default model on error
        resetToDefaultModel();
      }
    },
    [notifyUrdfCallbacks, urdfBlobUrls, urdfProcessor, resetToDefaultModel]
  );

  // Clean up blob URLs when component unmounts
  React.useEffect(() => {
    return () => {
      Object.values(urdfBlobUrls).forEach(URL.revokeObjectURL);
    };
  }, [urdfBlobUrls]);

  // Create the context value
  const contextValue: UrdfContextType = {
    urdfProcessor,
    registerUrdfProcessor,
    onUrdfDetected,
    processUrdfFiles,
    urdfBlobUrls,
    alternativeUrdfModels,
    isSelectionModalOpen,
    setIsSelectionModalOpen,
    urdfModelOptions,
    selectUrdfModel,

    // New properties for centralized robot data management
    isDefaultModel,
    setIsDefaultModel,
    resetToDefaultModel,
    urdfContent,
  };

  return (
    <UrdfContext.Provider value={contextValue}>{children}</UrdfContext.Provider>
  );
};