mono/packages/ui/src/contexts/LayoutContext.tsx

import React, { createContext, useContext, useState, useEffect, ReactNode } from 'react';
import { UnifiedLayoutManager, PageLayout, WidgetInstance, LayoutContainer } from '@/lib/unifiedLayoutManager';
import { widgetRegistry } from '@/lib/widgetRegistry';

interface LayoutContextType {
  // Generic page management
  loadPageLayout: (pageId: string, defaultName?: string) => Promise<void>;
  getLoadedPageLayout: (pageId: string) => PageLayout | null;
  clearPageLayout: (pageId: string) => Promise<void>;

  // Generic page actions
  addWidgetToPage: (pageId: string, containerId: string, widgetId: string, targetColumn?: number) => Promise<WidgetInstance>;
  removeWidgetFromPage: (pageId: string, widgetInstanceId: string) => Promise<void>;
  moveWidgetInPage: (pageId: string, widgetInstanceId: string, direction: 'up' | 'down' | 'left' | 'right') => Promise<void>;
  updatePageContainerColumns: (pageId: string, containerId: string, columns: number) => Promise<void>;
  updatePageContainerSettings: (pageId: string, containerId: string, settings: Partial<LayoutContainer['settings']>) => Promise<void>;
  addPageContainer: (pageId: string, parentContainerId?: string) => Promise<LayoutContainer>;
  removePageContainer: (pageId: string, containerId: string) => Promise<void>;
  movePageContainer: (pageId: string, containerId: string, direction: 'up' | 'down') => Promise<void>;
  updateWidgetProps: (pageId: string, widgetInstanceId: string, props: Record<string, any>) => Promise<void>;
  renameWidget: (pageId: string, widgetInstanceId: string, newId: string) => Promise<boolean>;
  exportPageLayout: (pageId: string) => Promise<string>;
  importPageLayout: (pageId: string, jsonData: string) => Promise<PageLayout>;

  // Manual save
  saveToApi: () => Promise<boolean>;

  // State
  isLoading: boolean;
  loadedPages: Map<string, PageLayout>;
}

const LayoutContext = createContext<LayoutContextType | undefined>(undefined);

interface LayoutProviderProps {
  children: ReactNode;
}

export const LayoutProvider: React.FC<LayoutProviderProps> = ({ children }) => {
  const [loadedPages, setLoadedPages] = useState<Map<string, PageLayout>>(new Map());
  const [isLoading, setIsLoading] = useState(true);

  // Initialize layouts on mount
  useEffect(() => {
    const initializeLayouts = async () => {
      try {
        // Get valid widget IDs from registry
        const validWidgetIds = new Set(widgetRegistry.getAll().map(w => w.metadata.id));

        // Only cleanup if we have widgets registered
        if (validWidgetIds.size > 0) {
          // Clean up all known pages
          const knownPages = ['playground-layout', 'dashboard-layout', 'profile-layout'];
          await Promise.all(knownPages.map(pageId =>
            UnifiedLayoutManager.cleanupInvalidWidgets(pageId, validWidgetIds)
          ));
        }
      } catch (error) {
        console.error('Failed to initialize layouts:', error);
      } finally {
        setIsLoading(false);
      }
    };

    initializeLayouts();
  }, []);

  useEffect(() => {

  }, [loadedPages]);

  const loadPageLayout = async (pageId: string, defaultName?: string) => {
    // Only load if not already cached
    if (!loadedPages.has(pageId)) {
      try {
        const layout = await UnifiedLayoutManager.getPageLayout(pageId, defaultName);
        setLoadedPages(prev => new Map(prev).set(pageId, layout));
      } catch (error) {
        console.error(`Failed to load page layout ${pageId}:`, error);
      }
    }
  };

  const getLoadedPageLayout = (pageId: string): PageLayout | null => {
    return loadedPages.get(pageId) || null;
  };

  const clearPageLayout = async (pageId: string) => {
    try {
      const currentLayout = loadedPages.get(pageId);
      if (!currentLayout) {
        throw new Error(`Layout for page ${pageId} not loaded`);
      }

      // Create a fresh empty layout with one empty container
      const clearedLayout: PageLayout = {
        id: pageId,
        name: currentLayout.name,
        containers: [
          {
            id: UnifiedLayoutManager.generateContainerId(),
            type: 'container',
            columns: 1,
            gap: 16,
            widgets: [],
            children: [],
            order: 0
          }
        ],
        createdAt: currentLayout.createdAt,
        updatedAt: Date.now()
      };

      // Update the in-memory cache
      setLoadedPages(prev => new Map(prev).set(pageId, clearedLayout));

      // Save to localStorage cache
      await saveLayoutToCache(pageId);
    } catch (error) {
      console.error(`Failed to clear page layout ${pageId}:`, error);
      throw error;
    }
  };

  const saveLayoutToCache = async (pageId: string) => {
    const currentLayout = loadedPages.get(pageId);
    if (!currentLayout) {
      console.error(`Cannot save page ${pageId}: layout not loaded`);
      return;
    }

    // Save the specific page to database
    await UnifiedLayoutManager.savePageLayout(currentLayout);

    // Force re-render of the specific page with deep clone
    // Create a deep clone to ensure React detects the change
    const clonedLayout = JSON.parse(JSON.stringify(currentLayout));
    setLoadedPages(prev => new Map(prev).set(pageId, clonedLayout));
  };

  const addWidgetToPage = async (pageId: string, containerId: string, widgetId: string, targetColumn?: number): Promise<WidgetInstance> => {
    try {
      const currentLayout = loadedPages.get(pageId);
      if (!currentLayout) {
        throw new Error(`Layout for page ${pageId} not loaded`);
      }

      const widget = UnifiedLayoutManager.addWidgetToContainer(currentLayout, containerId, widgetId, targetColumn);
      currentLayout.updatedAt = Date.now();

      await saveLayoutToCache(pageId);

      return widget;
    } catch (error) {
      console.error(`Failed to add widget to page ${pageId}:`, error);
      throw error;
    }
  };

  const removeWidgetFromPage = async (pageId: string, widgetInstanceId: string) => {
    try {
      const currentLayout = loadedPages.get(pageId);
      if (!currentLayout) {
        throw new Error(`Layout for page ${pageId} not loaded`);
      }

      const removed = UnifiedLayoutManager.removeWidgetFromContainer(currentLayout, widgetInstanceId);
      if (!removed) {
        throw new Error(`Widget ${widgetInstanceId} not found`);
      }

      currentLayout.updatedAt = Date.now();

      await saveLayoutToCache(pageId);
    } catch (error) {
      console.error(`Failed to remove widget from page ${pageId}:`, error);
      throw error;
    }
  };

  const moveWidgetInPage = async (pageId: string, widgetInstanceId: string, direction: 'up' | 'down' | 'left' | 'right') => {
    try {
      const currentLayout = loadedPages.get(pageId);
      if (!currentLayout) {
        throw new Error(`Layout for page ${pageId} not loaded`);
      }

      const moved = UnifiedLayoutManager.moveWidgetInContainer(currentLayout, widgetInstanceId, direction);
      if (!moved) {
        throw new Error(`Failed to move widget ${widgetInstanceId}`);
      }

      currentLayout.updatedAt = Date.now();

      await saveLayoutToCache(pageId);
    } catch (error) {
      console.error(`Failed to move widget in page ${pageId}:`, error);
      throw error;
    }
  };

  const updatePageContainerColumns = async (pageId: string, containerId: string, columns: number) => {
    try {
      const currentLayout = loadedPages.get(pageId);
      if (!currentLayout) {
        throw new Error(`Layout for page ${pageId} not loaded`);
      }

      const updated = UnifiedLayoutManager.updateContainerColumns(currentLayout, containerId, columns);
      if (!updated) {
        throw new Error(`Container ${containerId} not found`);
      }

      currentLayout.updatedAt = Date.now();

      await saveLayoutToCache(pageId);
    } catch (error) {
      console.error(`Failed to update container columns in page ${pageId}:`, error);
      throw error;
    }
  };

  const updatePageContainerSettings = async (pageId: string, containerId: string, settings: Partial<LayoutContainer['settings']>) => {
    try {
      const currentLayout = loadedPages.get(pageId);
      if (!currentLayout) {
        throw new Error(`Layout for page ${pageId} not loaded`);
      }

      const updated = UnifiedLayoutManager.updateContainerSettings(currentLayout, containerId, settings);
      if (!updated) {
        throw new Error(`Container ${containerId} not found`);
      }

      currentLayout.updatedAt = Date.now();

      await saveLayoutToCache(pageId);
    } catch (error) {
      console.error(`Failed to update container settings in page ${pageId}:`, error);
      throw error;
    }
  };

  const addPageContainer = async (pageId: string, parentContainerId?: string): Promise<LayoutContainer> => {
    try {
      const currentLayout = loadedPages.get(pageId);
      if (!currentLayout) {
        throw new Error(`Layout for page ${pageId} not loaded`);
      }

      const container = UnifiedLayoutManager.addContainer(currentLayout, parentContainerId);
      currentLayout.updatedAt = Date.now();

      await saveLayoutToCache(pageId);

      return container;
    } catch (error) {
      console.error(`Failed to add container to page ${pageId}:`, error);
      throw error;
    }
  };

  const removePageContainer = async (pageId: string, containerId: string) => {
    try {
      const currentLayout = loadedPages.get(pageId);
      if (!currentLayout) {
        throw new Error(`Layout for page ${pageId} not loaded`);
      }

      // For extension slots, allow removing the last container (to effectively remove the canvas)
      const isExtensionSlot = pageId.includes('-slot-');

      const removed = UnifiedLayoutManager.removeContainer(currentLayout, containerId, isExtensionSlot);
      if (!removed) {
        throw new Error(`Container ${containerId} not found or cannot be removed`);
      }

      currentLayout.updatedAt = Date.now();

      await saveLayoutToCache(pageId);
    } catch (error) {
      console.error(`Failed to remove container from page ${pageId}:`, error);
      throw error;
    }
  };

  const movePageContainer = async (pageId: string, containerId: string, direction: 'up' | 'down') => {
    try {
      const currentLayout = loadedPages.get(pageId);
      if (!currentLayout) {
        throw new Error(`Layout for page ${pageId} not loaded`);
      }

      const moved = UnifiedLayoutManager.moveRootContainer(currentLayout, containerId, direction);
      if (!moved) {
        // This can fail gracefully if the container is at the top/bottom, so no error needed.
        return;
      }

      currentLayout.updatedAt = Date.now();

      await saveLayoutToCache(pageId);
    } catch (error) {
      console.error(`Failed to move container in page ${pageId}:`, error);
      throw error;
    }
  };

  const updateWidgetProps = async (pageId: string, widgetInstanceId: string, props: Record<string, any>) => {
    try {
      const currentLayout = loadedPages.get(pageId);
      if (!currentLayout) {
        throw new Error(`Layout for page ${pageId} not loaded`);
      }

      const updated = UnifiedLayoutManager.updateWidgetProps(currentLayout, widgetInstanceId, props);
      if (!updated) {
        throw new Error(`Widget ${widgetInstanceId} not found`);
      }

      currentLayout.updatedAt = Date.now();

      await saveLayoutToCache(pageId);
    } catch (error) {
      console.error(`Failed to update widget props in page ${pageId}:`, error);
      throw error;
    }
  };

  const renameWidget = async (pageId: string, widgetInstanceId: string, newId: string): Promise<boolean> => {
    try {
      const currentLayout = loadedPages.get(pageId);
      if (!currentLayout) {
        throw new Error(`Layout for page ${pageId} not loaded`);
      }

      const success = UnifiedLayoutManager.renameWidget(currentLayout, widgetInstanceId, newId);
      if (!success) {
        return false;
      }

      currentLayout.updatedAt = Date.now();

      await saveLayoutToCache(pageId);
      return true;
    } catch (error) {
      console.error(`Failed to rename widget in page ${pageId}:`, error);
      throw error;
    }
  };

  const exportPageLayout = async (pageId: string): Promise<string> => {
    try {
      // Export directly from memory state to ensure consistency with UI
      const currentLayout = loadedPages.get(pageId);
      if (currentLayout) {
        return JSON.stringify(currentLayout, null, 2);
      }
      // Fallback to ULM if not loaded (though it should be)
      return await UnifiedLayoutManager.exportPageLayout(pageId);
    } catch (error) {
      console.error(`Failed to export page layout ${pageId}:`, error);
      throw error;
    }
  };

  const importPageLayout = async (pageId: string, jsonData: string): Promise<PageLayout> => {
    try {
      console.log(`[LayoutContext] Importing page layout for ${pageId}...`);
      const layout = await UnifiedLayoutManager.importPageLayout(pageId, jsonData);
      console.log('[LayoutContext] Layout imported successfully from ULM:', layout);
      // Directly update the state with the newly imported layout
      setLoadedPages(prev => {
        const newPages = new Map(prev);
        newPages.set(pageId, layout);
        console.log('[LayoutContext] Updating loadedPages state with new layout for pageId:', pageId, newPages);
        return newPages;
      });
      return layout;
    } catch (error) {
      console.error(`[LayoutContext] Failed to import page layout ${pageId}:`, error);
      throw error;
    }
  };

  const saveToApi = async (): Promise<boolean> => {
    try {
      // Save all loaded pages to database
      let allSaved = true;
      for (const [pageId, layout] of loadedPages.entries()) {
        try {
          await UnifiedLayoutManager.savePageLayout(layout);
        } catch (error) {
          console.error(`Failed to save page ${pageId}:`, error);
          allSaved = false;
        }
      }
      return allSaved;
    } catch (error) {
      console.error('Failed to save layouts to database:', error);
      return false;
    }
  };

  const value: LayoutContextType = {
    loadPageLayout,
    getLoadedPageLayout,
    clearPageLayout,
    addWidgetToPage,
    removeWidgetFromPage,
    moveWidgetInPage,
    updatePageContainerColumns,
    updatePageContainerSettings,
    addPageContainer,
    removePageContainer,
    movePageContainer,
    updateWidgetProps,
    exportPageLayout,
    importPageLayout,
    saveToApi,
    isLoading,
    loadedPages,
    renameWidget,
  };

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

export const useLayout = (): LayoutContextType => {
  const context = useContext(LayoutContext);
  if (context === undefined) {
    throw new Error('useLayout must be used within a LayoutProvider');
  }
  return context;
};