mule-image/frontend/src/components/timeline/Timeline.tsx

import { useRef, useEffect, useMemo, useState } from 'react'
import { useVirtualizer } from '@tanstack/react-virtual'
import { format, parseISO } from 'date-fns'
import { usePhotoStore } from '../../store/photoStore'
import { useFilterStore } from '../../store/filterStore'
import { PhotoThumbnail } from './PhotoThumbnail'
import { usePhotosQuery } from '../../hooks/usePhotosQuery'
import { useActiveHeapMembers } from '../../hooks/useActiveHeapMembersQuery'
import type { Photo } from '../../types/photo'

// Layout constants for the grid + grouped headers.
const THUMBNAIL_SIZE = 200
const GAP = 4
const PADDING = 16
const HEADER_HEIGHT = 36

interface PhotoCell {
  photo: Photo
  globalIndex: number
}

type TimelineItem =
  | { type: 'header'; key: string; label: string; height: number }
  | { type: 'row'; key: string; cells: PhotoCell[]; height: number }

/**
 * Build groups by month label when sorted by a date field. For non-temporal
 * sorts (filename / file_size / rating) we return a single un-headered group.
 */
function buildItems(
  photos: Photo[],
  columns: number,
  sortBy: string
): TimelineItem[] {
  if (photos.length === 0) return []

  const isDateSort = sortBy === 'taken_at' || sortBy === 'added_at'
  const items: TimelineItem[] = []

  // Helper: split a flat array of cells into rows of `columns` cells.
  const pushRowsForGroup = (groupKey: string, cells: PhotoCell[]) => {
    for (let i = 0; i < cells.length; i += columns) {
      const slice = cells.slice(i, i + columns)
      items.push({
        type: 'row',
        key: `${groupKey}::row::${i}`,
        cells: slice,
        height: THUMBNAIL_SIZE + GAP,
      })
    }
  }

  if (!isDateSort) {
    // No grouping — one row stream.
    const cells: PhotoCell[] = photos.map((photo, globalIndex) => ({
      photo,
      globalIndex,
    }))
    pushRowsForGroup('all', cells)
    return items
  }

  // Walk photos in order, breaking into groups whenever the month label changes.
  let currentLabel: string | null = null
  let bucket: PhotoCell[] = []
  let bucketIndex = 0

  const flushBucket = () => {
    if (bucket.length === 0 || currentLabel === null) return
    items.push({
      type: 'header',
      key: `header::${bucketIndex}::${currentLabel}`,
      label: currentLabel,
      height: HEADER_HEIGHT,
    })
    pushRowsForGroup(`${bucketIndex}::${currentLabel}`, bucket)
    bucketIndex++
    bucket = []
  }

  photos.forEach((photo, globalIndex) => {
    const dateStr =
      sortBy === 'taken_at'
        ? photo.taken_at
        : (photo as any).added_at ?? photo.taken_at
    let label: string
    if (dateStr) {
      try {
        label = format(parseISO(dateStr), 'MMMM yyyy')
      } catch {
        label = 'Unknown date'
      }
    } else {
      label = 'Unknown date'
    }
    if (label !== currentLabel) {
      flushBucket()
      currentLabel = label
    }
    bucket.push({ photo, globalIndex })
  })
  flushBucket()

  return items
}

export function Timeline() {
  const parentRef = useRef<HTMLDivElement>(null)
  const [containerWidth, setContainerWidth] = useState(0)

  const {
    selectedPhotos,
    lastSelectedIndex,
    rangeStartIndex,
    selectPhoto,
    togglePhotoSelection,
    clearSelection,
    openPreview,
  } = usePhotoStore()

  const sortBy = useFilterStore((s) => s.sortBy)

  // Calculate number of columns based on container width.
  const columns = useMemo(() => {
    if (containerWidth === 0) return 4
    return Math.max(
      1,
      Math.floor((containerWidth - PADDING * 2) / (THUMBNAIL_SIZE + GAP))
    )
  }, [containerWidth])

  // Shared photos query — both Timeline and PreviewView use the same hook so
  // they share one cache entry, regardless of filter state.
  const { data: photos = [], isLoading } = usePhotosQuery()

  // Membership in the active heap (for the basket affordance). Subscribed
  // once at this level so we don't have hundreds of thumbnails each
  // subscribing to the same query.
  const { memberIds: activeHeapMembers } = useActiveHeapMembers()

  // Build the flat virtualizer items: a mix of date-group headers and rows
  // of photos. Headers only appear when sorted by a date field.
  const items = useMemo(
    () => buildItems(photos, columns, sortBy),
    [photos, columns, sortBy]
  )

  // Pre-computed offset of every header in the virtualizer's coordinate
  // space, used to drive the sticky-header overlay below.
  const headerOffsets = useMemo(() => {
    const result: { offset: number; label: string }[] = []
    let cumulative = 0
    for (const item of items) {
      if (item.type === 'header') {
        result.push({ offset: cumulative, label: item.label })
      }
      cumulative += item.height
    }
    return result
  }, [items])

  // Range-selection helper. Operates on the global photos array, not on
  // virtualizer items.
  const selectRange = (endIndex: number) => {
    const startIndex = rangeStartIndex ?? lastSelectedIndex ?? 0
    const minIndex = Math.min(startIndex, endIndex)
    const maxIndex = Math.max(startIndex, endIndex)
    for (let i = minIndex; i <= maxIndex; i++) {
      if (i < photos.length && !selectedPhotos.includes(photos[i].id)) {
        togglePhotoSelection(photos[i].id, i)
      }
    }
  }

  // Virtual scrolling setup with per-item heights.
  const virtualizer = useVirtualizer({
    count: items.length,
    getScrollElement: () => parentRef.current,
    estimateSize: (index) => items[index]?.height ?? THUMBNAIL_SIZE,
    overscan: 5,
  })

  // Re-measure when items change (column count, group structure).
  useEffect(() => {
    virtualizer.measure()
  }, [items, virtualizer])

  // Track scroll position so we can show the current group label as a
  // pinned overlay at the top of the scroll container. The virtualizer's
  // items use transform translateY (so CSS position: sticky doesn't work
  // on the inline headers); the overlay sidesteps that by living outside
  // the virtualizer's positioned children.
  const [scrollTop, setScrollTop] = useState(0)
  useEffect(() => {
    const el = parentRef.current
    if (!el) return
    const onScroll = () => setScrollTop(el.scrollTop)
    el.addEventListener('scroll', onScroll, { passive: true })
    return () => el.removeEventListener('scroll', onScroll)
  }, [])

  // Find the latest header whose start <= scrollTop. That's the label of
  // the group containing whatever is currently at the top of the viewport.
  const stickyLabel = useMemo(() => {
    if (headerOffsets.length === 0) return null
    let current: string | null = null
    for (const h of headerOffsets) {
      if (h.offset <= scrollTop) current = h.label
      else break
    }
    return current
  }, [headerOffsets, scrollTop])

  // Measure container width on mount and resize.
  useEffect(() => {
    const measureWidth = () => {
      if (parentRef.current) {
        setContainerWidth(parentRef.current.clientWidth)
      }
    }
    measureWidth()
    window.addEventListener('resize', measureWidth)
    return () => window.removeEventListener('resize', measureWidth)
  }, [])

  // Handle keyboard shortcuts for photo navigation. Operates on the flat
  // photos array, so it ignores grouping.
  useEffect(() => {
    const handleKeyDown = (e: KeyboardEvent) => {
      if (photos.length === 0) return
      const target = e.target as HTMLElement | null
      if (target && (target.tagName === 'INPUT' || target.tagName === 'TEXTAREA')) {
        return
      }

      const currentIndex = lastSelectedIndex ?? -1

      switch (e.key) {
        case 'ArrowUp':
          e.preventDefault()
          if (currentIndex > columns - 1) {
            const newIndex = currentIndex - columns
            if (e.shiftKey) selectRange(newIndex)
            else selectPhoto(photos[newIndex].id, newIndex)
          }
          break
        case 'ArrowDown':
          e.preventDefault()
          if (currentIndex < photos.length - columns) {
            const newIndex = Math.min(currentIndex + columns, photos.length - 1)
            if (e.shiftKey) selectRange(newIndex)
            else selectPhoto(photos[newIndex].id, newIndex)
          }
          break
        case 'ArrowLeft':
          e.preventDefault()
          if (currentIndex > 0) {
            const newIndex = currentIndex - 1
            if (e.shiftKey) selectRange(newIndex)
            else selectPhoto(photos[newIndex].id, newIndex)
          }
          break
        case 'ArrowRight':
          e.preventDefault()
          if (currentIndex < photos.length - 1) {
            const newIndex = currentIndex + 1
            if (e.shiftKey) selectRange(newIndex)
            else selectPhoto(photos[newIndex].id, newIndex)
          }
          break
        case 'a':
          if (e.ctrlKey || e.metaKey) {
            e.preventDefault()
            photos.forEach((photo, index) => {
              if (!selectedPhotos.includes(photo.id)) {
                togglePhotoSelection(photo.id, index)
              }
            })
          }
          break
        case 'Escape':
          e.preventDefault()
          clearSelection()
          break
      }
    }

    window.addEventListener('keydown', handleKeyDown)
    return () => window.removeEventListener('keydown', handleKeyDown)
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [photos, selectedPhotos, lastSelectedIndex, columns])

  if (isLoading) {
    return (
      <div className="flex items-center justify-center h-full">
        <div className="text-text-muted">Loading photos...</div>
      </div>
    )
  }

  if (photos.length === 0) {
    return (
      <div className="flex items-center justify-center h-full">
        <div className="text-text-muted">No photos found. Add a source folder to get started.</div>
      </div>
    )
  }

  return (
    <div className="relative h-full">
      {/* Sticky group-header overlay. Lives outside the virtualizer's
       *  positioned children so it isn't affected by translateY transforms.
       *  Updates as the user scrolls past month boundaries. */}
      {stickyLabel && (
        <div className="pointer-events-none absolute left-0 right-0 top-0 z-20 border-b border-border bg-bg/90 px-4 py-1 backdrop-blur-sm">
          <h3 className="text-sm font-semibold uppercase tracking-wide text-text">
            {stickyLabel}
          </h3>
        </div>
      )}

      <div
        ref={parentRef}
        className="h-full overflow-auto bg-bg"
        style={{ padding: `${PADDING}px` }}
      >
      <div
        style={{
          height: `${virtualizer.getTotalSize()}px`,
          width: '100%',
          position: 'relative',
        }}
      >
        {virtualizer.getVirtualItems().map((virtualItem) => {
          const item = items[virtualItem.index]
          if (!item) return null

          if (item.type === 'header') {
            return (
              <div
                key={virtualItem.key}
                style={{
                  position: 'absolute',
                  top: 0,
                  left: 0,
                  width: '100%',
                  height: `${virtualItem.size}px`,
                  transform: `translateY(${virtualItem.start}px)`,
                }}
                className="flex items-end pb-1"
              >
                <h3 className="text-sm font-semibold uppercase tracking-wide text-text-muted">
                  {item.label}
                </h3>
              </div>
            )
          }

          // row
          return (
            <div
              key={virtualItem.key}
              style={{
                position: 'absolute',
                top: 0,
                left: 0,
                width: '100%',
                height: `${virtualItem.size}px`,
                transform: `translateY(${virtualItem.start}px)`,
              }}
            >
              <div className="flex" style={{ gap: `${GAP}px` }}>
                {item.cells.map(({ photo, globalIndex }) => (
                  <PhotoThumbnail
                    key={photo.id}
                    photo={photo}
                    size={THUMBNAIL_SIZE}
                    isSelected={selectedPhotos.includes(photo.id)}
                    isInActiveHeap={activeHeapMembers.has(photo.id)}
                    onClick={(e) => {
                      if (e.shiftKey && lastSelectedIndex !== null) {
                        selectRange(globalIndex)
                      } else if (e.ctrlKey || e.metaKey) {
                        togglePhotoSelection(photo.id, globalIndex)
                      } else {
                        selectPhoto(photo.id, globalIndex)
                      }
                    }}
                    onDoubleClick={() => openPreview(photo.id)}
                  />
                ))}
              </div>
            </div>
          )
        })}
      </div>
      </div>
    </div>
  )
}