Map Image Overlays and State Restoration

.env

@@ -1 +1,2 @@
 DATABASE_URL="postgresql://user:password@host:port/db"
+NEXT_PUBLIC_MAPBOX_ACCESS_TOKEN=pk.eyJ1IjoiZHVtbXkiLCJhIjoiY2p1bW15In0.ZHVtbXk=

app/search/[id]/page.tsx

@@ -59,6 +59,17 @@ export default async function SearchPage({ params }: SearchPageProps) {
     };
   });
 
+  // Extract initial map data from the latest 'data' message
+  const latestDataMessage = [...dbMessages].reverse().find(m => m.role === 'data');
+  let initialMapData = undefined;
+  if (latestDataMessage) {
+    try {
+      initialMapData = JSON.parse(latestDataMessage.content);
+    } catch (e) {
+      console.error('Failed to parse initial map data:', e);
+    }
+  }
+
   return (
     <AI
       initialAIState={{
@@ -68,7 +79,7 @@ export default async function SearchPage({ params }: SearchPageProps) {
                              // If needed for styling or other logic, it can be set.
       }}
     >
-      <MapDataProvider>
+      <MapDataProvider initialData={initialMapData}>
         <Chat id={id} />
       </MapDataProvider>
     </AI>

components/chat.tsx

@@ -19,6 +19,7 @@ import { MapDataProvider, useMapData } from './map/map-data-context'; // Add thi
 import { updateDrawingContext } from '@/lib/actions/chat'; // Import the server action
 import dynamic from 'next/dynamic'
 import { HeaderSearchButton } from './header-search-button'
+import { debounce } from 'lodash'
 
 type ChatProps = {
   id?: string // This is the chatId
@@ -89,21 +90,29 @@ export function Chat({ id }: ChatProps) {
     }
   }, [isSubmitting])
 
-  // useEffect to call the server action when drawnFeatures changes
+  // Debounced version of updateDrawingContext
+  const debouncedUpdateDrawingContext = useRef(
+    debounce((chatId: string, contextData: any) => {
+      console.log('Chat.tsx: calling debounced updateDrawingContext');
+      updateDrawingContext(chatId, contextData);
+    }, 2000)
+  ).current;
+
+  // useEffect to call the server action when map data changes
   useEffect(() => {
-    if (id && mapData.drawnFeatures && mapData.cameraState) {
-      console.log('Chat.tsx: drawnFeatures changed, calling updateDrawingContext', mapData.drawnFeatures);
-      updateDrawingContext(id, {
-        drawnFeatures: mapData.drawnFeatures,
+    if (id && (mapData.drawnFeatures || mapData.imageOverlays) && mapData.cameraState) {
+      debouncedUpdateDrawingContext(id, {
+        drawnFeatures: mapData.drawnFeatures || [],
+        imageOverlays: mapData.imageOverlays || [],
         cameraState: mapData.cameraState,
       });
     }
-  }, [id, mapData.drawnFeatures, mapData.cameraState]);
+  }, [id, mapData.drawnFeatures, mapData.imageOverlays, mapData.cameraState, debouncedUpdateDrawingContext]);
 
   // Mobile layout
   if (isMobile) {
     return (
-      <MapDataProvider> {/* Add Provider */}
+      <>
         <HeaderSearchButton />
         <div className="mobile-layout-container">
           <div className="mobile-map-section">
@@ -157,13 +166,13 @@ export function Chat({ id }: ChatProps) {
           )}
         </div>
         </div>
-      </MapDataProvider>
+      </>
     );
   }
 
   // Desktop layout
   return (
-    <MapDataProvider> {/* Add Provider */}
+    <>
       <HeaderSearchButton />
       <div className="flex justify-start items-start">
         {/* This is the new div for scrolling */}
@@ -221,6 +230,6 @@ export function Chat({ id }: ChatProps) {
           {activeView ? <SettingsView /> : <MapProvider />}
         </div>
       </div>
-    </MapDataProvider>
+    </>
   );
 }

components/map/image-overlay-layer.tsx

@@ -0,0 +1,167 @@
+'use client'
+
+import { useEffect, useRef, useCallback } from 'react'
+import mapboxgl from 'mapbox-gl'
+import { useMap } from './map-context'
+import { useMapData, ImageOverlay } from './map-data-context'
+import { useMapToggle, MapToggleEnum } from '../map-toggle-context'
+import { X } from 'lucide-react'
+import { Button } from '../ui/button'
+import { createRoot } from 'react-dom/client'
+
+interface ImageOverlayLayerProps {
+  overlay: ImageOverlay;
+}
+
+export function ImageOverlayLayer({ overlay }: ImageOverlayLayerProps) {
+  const { map } = useMap()
+  const { setMapData } = useMapData()
+  const { mapType } = useMapToggle()
+  const isDrawingMode = mapType === MapToggleEnum.DrawingMode
+  const markersRef = useRef<mapboxgl.Marker[]>([])
+  const deleteMarkerRef = useRef<mapboxgl.Marker | null>(null)
+  const deleteRootRef = useRef<any>(null)
+
+  const sourceId = `image-source-${overlay.id}`
+  const layerId = `image-layer-${overlay.id}`
+
+  // Update overlay coordinates in global state
+  const updateCoordinates = useCallback((index: number, lngLat: [number, number]) => {
+    setMapData(prev => {
+      const overlays = prev.imageOverlays || []
+      const newOverlays = overlays.map(o => {
+        if (o.id === overlay.id) {
+          const newCoords = [...o.coordinates] as [[number, number], [number, number], [number, number], [number, number]]
+          newCoords[index] = lngLat
+          return { ...o, coordinates: newCoords }
+        }
+        return o
+      })
+      return { ...prev, imageOverlays: newOverlays }
+    })
+  }, [overlay.id, setMapData])
+
+  const removeOverlay = useCallback(() => {
+    setMapData(prev => ({
+      ...prev,
+      imageOverlays: (prev.imageOverlays || []).filter(o => o.id !== overlay.id)
+    }))
+  }, [overlay.id, setMapData])
+
+  useEffect(() => {
+    if (!map) return
+
+    const onMapLoad = () => {
+      if (!map.getSource(sourceId)) {
+        map.addSource(sourceId, {
+          type: 'image',
+          url: overlay.url,
+          coordinates: overlay.coordinates
+        })
+
+        map.addLayer({
+          id: layerId,
+          type: 'raster',
+          source: sourceId,
+          paint: {
+            'raster-opacity': overlay.opacity || 0.7
+          }
+        })
+      }
+    }
+
+    if (map.isStyleLoaded()) {
+      onMapLoad()
+    } else {
+      map.on('load', onMapLoad)
+    }
+
+    return () => {
+      if (map.getLayer(layerId)) map.removeLayer(layerId)
+      if (map.getSource(sourceId)) map.removeSource(sourceId)
+    }
+  }, [map, sourceId, layerId, overlay.url]) // coordinates handled separately
+
+  // Sync coordinates with map source
+  useEffect(() => {
+    if (!map) return
+    const source = map.getSource(sourceId) as mapboxgl.ImageSource
+    if (source) {
+      source.setCoordinates(overlay.coordinates)
+    }
+  }, [map, overlay.coordinates, sourceId])
+
+  // Sync opacity
+  useEffect(() => {
+    if (!map) return
+    if (map.getLayer(layerId)) {
+      map.setPaintProperty(layerId, 'raster-opacity', overlay.opacity || 0.7)
+    }
+  }, [map, overlay.opacity, layerId])
+
+  // Draggable markers for corners
+  useEffect(() => {
+    if (!map) return
+
+    // Clean up existing markers
+    markersRef.current.forEach(m => m.remove())
+    markersRef.current = []
+    if (deleteMarkerRef.current) {
+      deleteMarkerRef.current.remove()
+      deleteMarkerRef.current = null
+    }
+
+    if (isDrawingMode) {
+      // Add 4 corner markers
+      overlay.coordinates.forEach((coord, index) => {
+        const el = document.createElement('div')
+        el.className = 'w-4 h-4 bg-white border-2 border-primary rounded-full cursor-move shadow-md'
+
+        const marker = new mapboxgl.Marker({
+          element: el,
+          draggable: true
+        })
+          .setLngLat(coord)
+          .addTo(map)
+
+        marker.on('drag', () => {
+          const newLngLat = marker.getLngLat()
+          updateCoordinates(index, [newLngLat.lng, newLngLat.lat])
+        })
+
+        markersRef.current.push(marker)
+      })
+
+      // Add delete button near top-right corner
+      const deleteEl = document.createElement('div')
+      const root = createRoot(deleteEl)
+      deleteRootRef.current = root
+      root.render(
+        <Button
+          variant="destructive"
+          size="icon"
+          className="w-6 h-6 rounded-full shadow-lg"
+          onClick={removeOverlay}
+        >
+          <X size={14} />
+        </Button>
+      )
+
+      const deleteMarker = new mapboxgl.Marker({
+        element: deleteEl,
+        offset: [0, -20]
+      })
+        .setLngLat(overlay.coordinates[1]) // Near Top-right
+        .addTo(map)
+
+      deleteMarkerRef.current = deleteMarker
+    }
+
+    return () => {
+      markersRef.current.forEach(m => m.remove())
+      if (deleteMarkerRef.current) deleteMarkerRef.current.remove()
+    }
+  }, [map, isDrawingMode, overlay.coordinates, updateCoordinates, removeOverlay])
+
+  return null
+}

components/map/map-data-context.tsx

@@ -12,6 +12,13 @@ export interface CameraState {
   heading?: number;
 }
 
+export interface ImageOverlay {
+  id: string;
+  url: string;
+  coordinates: [[number, number], [number, number], [number, number], [number, number]]; // [NW, NE, SE, SW]
+  opacity: number;
+}
+
 export interface MapData {
   targetPosition?: { lat: number; lng: number } | null; // For flying to a location
   cameraState?: CameraState; // For saving camera state
@@ -29,6 +36,7 @@ export interface MapData {
     longitude: number;
     title?: string;
   }>;
+  imageOverlays?: ImageOverlay[];
 }
 
 interface MapDataContextType {
@@ -38,8 +46,8 @@ interface MapDataContextType {
 
 const MapDataContext = createContext<MapDataContextType | undefined>(undefined);
 
-export const MapDataProvider: React.FC<{ children: ReactNode }> = ({ children }) => {
-  const [mapData, setMapData] = useState<MapData>({ drawnFeatures: [], markers: [] });
+export const MapDataProvider: React.FC<{ children: ReactNode; initialData?: MapData }> = ({ children, initialData }) => {
+  const [mapData, setMapData] = useState<MapData>(initialData || { drawnFeatures: [], markers: [], imageOverlays: [] });
 
   return (
     <MapDataContext.Provider value={{ mapData, setMapData }}>

components/map/mapbox-map.tsx

@@ -13,6 +13,7 @@ import { useMapToggle, MapToggleEnum } from '../map-toggle-context'
 import { useMapData } from './map-data-context'; // Add this import
 import { useMapLoading } from '../map-loading-context'; // Import useMapLoading
 import { useMap } from './map-context'
+import { ImageOverlayLayer } from './image-overlay-layer'
 
 mapboxgl.accessToken = process.env.NEXT_PUBLIC_MAPBOX_ACCESS_TOKEN as string;
 
@@ -569,6 +570,9 @@ export const Mapbox: React.FC<{ position?: { latitude: number; longitude: number
         onMouseUp={handleMouseUp}
         onMouseLeave={handleMouseUp} // Clear timer if mouse leaves container while pressed
       />
+      {mapData.imageOverlays?.map(overlay => (
+        <ImageOverlayLayer key={overlay.id} overlay={overlay} />
+      ))}
     </div>
   )
 }