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

@@ -89,21 +89,22 @@ export function Chat({ id }: ChatProps) {
     }
   }, [isSubmitting])
 
-  // useEffect to call the server action when drawnFeatures changes
+  // 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);
+    if (id && (mapData.drawnFeatures || mapData.imageOverlays) && mapData.cameraState) {
+      console.log('Chat.tsx: map data changed, calling updateDrawingContext');
       updateDrawingContext(id, {
-        drawnFeatures: mapData.drawnFeatures,
+        drawnFeatures: mapData.drawnFeatures || [],
+        imageOverlays: mapData.imageOverlays || [],
         cameraState: mapData.cameraState,
       });
     }
-  }, [id, mapData.drawnFeatures, mapData.cameraState]);
+  }, [id, mapData.drawnFeatures, mapData.imageOverlays, mapData.cameraState]);
 
   // Mobile layout
   if (isMobile) {
     return (
-      <MapDataProvider> {/* Add Provider */}
+      <>
         <HeaderSearchButton />
         <div className="mobile-layout-container">
           <div className="mobile-map-section">
@@ -157,13 +158,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 +222,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>
   )
 }

components/user-message.tsx

@@ -1,6 +1,14 @@
+'use client'
+
 import React from 'react'
 import Image from 'next/image'
 import { ChatShare } from './chat-share'
+import { Button } from './ui/button'
+import { MapPin } from 'lucide-react'
+import { useMap } from './map/map-context'
+import { useMapData } from './map/map-data-context'
+import { useMapToggle, MapToggleEnum } from './map-toggle-context'
+import { nanoid } from 'nanoid'
 
 type UserMessageContentPart =
   | { type: 'text'; text: string }
@@ -17,6 +25,9 @@ export const UserMessage: React.FC<UserMessageProps> = ({
   chatId,
   showShare = false
 }) => {
+  const { map } = useMap()
+  const { setMapData } = useMapData()
+  const { setMapType } = useMapToggle()
   const enableShare = process.env.ENABLE_SHARE === 'true'
 
   // Normalize content to an array
@@ -31,18 +42,62 @@ export const UserMessage: React.FC<UserMessageProps> = ({
     (part): part is { type: 'image'; image: string } => part.type === 'image'
   )?.image
 
+  const handlePlaceOnMap = () => {
+    if (!map || !imagePart) return
+
+    const bounds = map.getBounds()
+    if (!bounds) return
+    const north = bounds.getNorth()
+    const south = bounds.getSouth()
+    const east = bounds.getEast()
+    const west = bounds.getWest()
+
+    const latStep = (north - south) / 4
+    const lngStep = (east - west) / 4
+
+    const nw: [number, number] = [west + lngStep, north - latStep]
+    const ne: [number, number] = [east - lngStep, north - latStep]
+    const se: [number, number] = [east - lngStep, south + latStep]
+    const sw: [number, number] = [west + lngStep, south + latStep]
+
+    const newOverlay = {
+      id: nanoid(),
+      url: imagePart,
+      coordinates: [nw, ne, se, sw] as [[number, number], [number, number], [number, number], [number, number]],
+      opacity: 0.7
+    }
+
+    setMapData(prev => ({
+      ...prev,
+      imageOverlays: [...(prev.imageOverlays || []), newOverlay]
+    }))
+
+    setMapType(MapToggleEnum.DrawingMode)
+  }
+
   return (
     <div className="flex items-start w-full space-x-3 mt-2">
       <div className="flex-1 space-y-2">
         {imagePart && (
-          <div className="p-2 border rounded-lg bg-muted w-fit">
+          <div className="group relative p-2 border rounded-lg bg-muted w-fit">
             <Image
               src={imagePart}
               alt="attachment"
               width={300}
               height={300}
               className="max-w-xs max-h-64 rounded-md object-contain"
             />
+            <div className="absolute top-4 right-4 opacity-0 group-hover:opacity-100 transition-opacity">
+              <Button
+                variant="secondary"
+                size="sm"
+                className="h-8 px-2 shadow-md"
+                onClick={handlePlaceOnMap}
+              >
+                <MapPin size={14} className="mr-1" />
+                Place on Map
+              </Button>
+            </div>
           </div>
         )}
         {textPart && <div className="text-xl break-words">{textPart}</div>}

lib/actions/chat.ts

@@ -162,7 +162,7 @@ export async function saveChat(chat: OldChatType, userId: string): Promise<strin
 //   return null;
 // }
 
-export async function updateDrawingContext(chatId: string, contextData: { drawnFeatures: any[], cameraState: any }) {
+export async function updateDrawingContext(chatId: string, contextData: { drawnFeatures?: any[], imageOverlays?: any[], cameraState: any }) {
   'use server';
   console.log('[Action] updateDrawingContext called for chatId:', chatId);