Briven/apps/studio/components/ui/Charts/ComposedChartHandler.tsx

import dayjs from 'dayjs'
import { List, Loader2 } from 'lucide-react'
import { useRouter } from 'next/router'
import React, { PropsWithChildren, useEffect, useMemo, useRef, useState } from 'react'
import { Card, cn, WarningIcon } from 'ui'

import type { ChartHighlightAction } from './ChartHighlightActions'
import type { ChartData } from './Charts.types'
import { ComposedChart } from './ComposedChart'
import { MultiAttribute } from './ComposedChart.utils'
import { useChartHighlight } from './useChartHighlight'
import Panel from '@/components/ui/Panel'
import { AnalyticsInterval, DataPoint } from '@/data/analytics/constants'
import { useInfraMonitoringQueries } from '@/data/analytics/infra-monitoring-queries'
import { InfraMonitoringAttribute } from '@/data/analytics/infra-monitoring-query'
import { useProjectDailyStatsQueries } from '@/data/analytics/project-daily-stats-queries'
import { ProjectDailyStatsAttribute } from '@/data/analytics/project-daily-stats-query'
import type { UpdateDateRange } from '@/pages/project/[ref]/observability/database'
import { useDatabaseSelectorStateSnapshot } from '@/state/database-selector'

export interface ComposedChartHandlerProps {
  id?: string
  label: string
  attributes: MultiAttribute[]
  startDate: string
  endDate: string
  interval?: string
  customDateFormat?: string
  defaultChartStyle?: 'bar' | 'line' | 'stackedAreaLine'
  hideChartType?: boolean
  data?: ChartData | DataPoint[]
  isLoading?: boolean
  format?: string
  highlightedValue?: string | number
  className?: string
  showTooltip?: boolean
  showLegend?: boolean
  showTotal?: boolean
  showMaxValue?: boolean
  normalizeVisibleStackToPercent?: boolean
  updateDateRange?: UpdateDateRange
  valuePrecision?: number
  isVisible?: boolean
  docsUrl?: string
  hide?: boolean
  syncId?: string
  YAxisProps?: {
    width?: number
    tickFormatter?: (value: number) => string
    domain?: [number | string, number | string]
    allowDataOverflow?: boolean
  }
}

/**
 * Wrapper component that handles intersection observer logic for lazy loading
 */
const LazyChartWrapper = ({ children }: PropsWithChildren) => {
  const [isVisible, setIsVisible] = useState(false)
  const ref = useRef<HTMLDivElement>(null)

  useEffect(() => {
    const observer = new IntersectionObserver(
      ([entry]) => {
        if (entry.isIntersecting) {
          setIsVisible(true)
          observer.disconnect()
        }
      },
      {
        rootMargin: '150px 0px', // Start loading before the component enters viewport
        threshold: 0,
      }
    )

    const currentRef = ref.current
    if (currentRef) {
      observer.observe(currentRef)
    }

    return () => {
      if (currentRef) {
        observer.unobserve(currentRef)
      }
    }
  }, [])

  return (
    <div ref={ref}>
      {React.cloneElement(children as React.ReactElement<{ isVisible: boolean }>, { isVisible })}
    </div>
  )
}

/**
 * Controls chart display state. Optionally fetches static chart data if data is not provided.
 *
 * If the `data` prop is provided, it will disable automatic chart data fetching and pass the data directly to the chart render.
 * - loading state can also be provided through the `isLoading` prop, to display loading placeholders. Ignored if `data` key not provided.
 * - if `isLoading=true` and `data` is `undefined`, loading error message will be shown.
 *
 * Provided data must be in the expected chart format.
 */
const ComposedChartHandler = ({
  label,
  attributes,
  startDate,
  endDate,
  interval,
  customDateFormat,
  children = null,
  defaultChartStyle = 'bar',
  hideChartType = false,
  data,
  isLoading,
  format,
  highlightedValue,
  className,
  showTooltip,
  showLegend,
  showMaxValue,
  showTotal,
  updateDateRange,
  valuePrecision,
  isVisible = true,
  id,
  syncId,
  ...otherProps
}: PropsWithChildren<ComposedChartHandlerProps>) => {
  const router = useRouter()
  const { ref } = router.query

  const state = useDatabaseSelectorStateSnapshot()
  const [chartStyle, setChartStyle] = useState<string>(defaultChartStyle)
  const chartHighlight = useChartHighlight()

  const databaseIdentifier = state.selectedDatabaseId

  const attributeQueries = useAttributeQueries(
    attributes,
    ref,
    startDate,
    endDate,
    interval as AnalyticsInterval,
    databaseIdentifier,
    Array.isArray(data) ? undefined : data,
    isVisible
  )

  const combinedData = useMemo(() => {
    if (data) return Array.isArray(data) ? data : data.data

    const isLoading = attributeQueries.some((query: any) => query.isLoading)
    if (isLoading) return undefined

    const hasError = attributeQueries.some((query: any) => !query.data)
    if (hasError) return undefined

    const timestamps = new Set<string>()
    attributeQueries.forEach((query: any) => {
      query.data?.data?.forEach((point: any) => {
        if (point?.period_start) {
          timestamps.add(point.period_start)
        }
      })
    })

    const referenceLineQueries = attributeQueries.filter(
      (_, index) => attributes[index].provider === 'reference-line'
    )

    const combined = Array.from(timestamps)
      .sort()
      .map((timestamp) => {
        const point: any = { timestamp }

        attributes.forEach((attr, index) => {
          if (!attr) return

          if (attr.customValue !== undefined) {
            point[attr.attribute] = attr.customValue
            return
          }

          if (attr.provider === 'reference-line') return

          const queryData = attributeQueries[index]?.data?.data
          const matchingPoint = queryData?.find((p: any) => p.period_start === timestamp)
          let value = matchingPoint?.[attr.attribute] ?? 0

          if (attr.manipulateValue && typeof attr.manipulateValue === 'function') {
            const numericValue = typeof value === 'number' ? value : Number(value) || 0
            value = attr.manipulateValue(numericValue)
          }

          point[attr.attribute] = value
        })

        referenceLineQueries.forEach((query: any) => {
          const attr = query.data.attribute
          const value = query.data.total
          point[attr] = value
        })

        const formattedDataPoint: DataPoint =
          !('period_start' in point) && 'timestamp' in point
            ? { ...point, period_start: dayjs.utc(point.timestamp).unix() * 1000 }
            : point

        return formattedDataPoint
      })

    return combined as DataPoint[]
  }, [data, attributeQueries, attributes])

  const loading = isLoading || attributeQueries.some((query: any) => query.isLoading)

  const _highlightedValue = useMemo(() => {
    if (highlightedValue !== undefined) return highlightedValue

    const firstAttr = attributes[0]
    const firstQuery = attributeQueries[0]
    const firstData = firstQuery?.data

    if (!firstData) return undefined

    const shouldHighlightMaxValue =
      firstAttr.provider === 'daily-stats' &&
      !firstAttr.attribute.includes('ingress') &&
      !firstAttr.attribute.includes('egress') &&
      'maximum' in firstData

    const shouldHighlightTotalGroupedValue = 'totalGrouped' in firstData

    return shouldHighlightMaxValue
      ? firstData.maximum
      : firstAttr.provider === 'daily-stats'
        ? firstData.total
        : shouldHighlightTotalGroupedValue
          ? firstData.totalGrouped?.[firstAttr.attribute as keyof typeof firstData.totalGrouped]
          : (firstData.data[firstData.data.length - 1] as any)?.[firstAttr.attribute]
  }, [highlightedValue, attributes, attributeQueries])

  const highlightActions: ChartHighlightAction[] = useMemo(() => {
    return [
      {
        id: 'open-logs',
        label: 'Open in Postgres Logs',
        icon: <List size={12} />,
        onSelect: ({ start, end }) => {
          const projectRef = ref as string
          if (!projectRef) return
          const url = `/project/${projectRef}/logs/postgres-logs?its=${start}&ite=${end}`
          router.push(url)
        },
      },
    ]
  }, [ref])

  if (loading) {
    return (
      <Card
        className={cn(
          'flex min-h-[280px] w-full flex-col items-center justify-center gap-y-2',
          className
        )}
      >
        <Loader2 size={18} className="animate-spin text-border-strong" />
        <p className="text-xs text-foreground-lighter">Loading data for {label}</p>
      </Card>
    )
  }

  if (!combinedData) {
    return (
      <div className="flex h-52 w-full flex-col items-center justify-center gap-y-2">
        <WarningIcon />
        <p className="text-xs text-foreground-lighter">Unable to load data for {label}</p>
      </div>
    )
  }

  return (
    <Panel
      noMargin
      noHideOverflow
      className={cn('relative w-full scroll-mt-16', className)}
      wrapWithLoading={false}
      id={id ?? label.toLowerCase().replaceAll(' ', '-')}
    >
      <Panel.Content className="flex flex-col gap-4">
        <div className="absolute right-6 z-50 flex justify-between scroll-mt-16">{children}</div>
        <ComposedChart
          attributes={attributes}
          data={combinedData as DataPoint[]}
          format={format}
          // [Joshen] This is where it's messing up
          xAxisKey="period_start"
          yAxisKey={attributes[0].attribute}
          highlightedValue={_highlightedValue}
          title={label}
          customDateFormat={customDateFormat}
          chartHighlight={chartHighlight}
          chartStyle={chartStyle}
          showTooltip={showTooltip}
          showLegend={showLegend}
          showTotal={showTotal}
          showMaxValue={showMaxValue}
          onChartStyleChange={setChartStyle}
          updateDateRange={updateDateRange}
          valuePrecision={valuePrecision}
          hideChartType={hideChartType}
          syncId={syncId}
          highlightActions={highlightActions}
          {...otherProps}
        />
      </Panel.Content>
    </Panel>
  )
}

const useAttributeQueries = (
  attributes: MultiAttribute[],
  ref: string | string[] | undefined,
  startDate: string,
  endDate: string,
  interval: AnalyticsInterval,
  databaseIdentifier: string | undefined,
  data: ChartData | undefined,
  isVisible: boolean
) => {
  const infraAttributes = attributes
    .filter((attr) => attr?.provider === 'infra-monitoring')
    .map((attr) => attr.attribute as InfraMonitoringAttribute)
  const dailyStatsAttributes = attributes
    .filter((attr) => attr?.provider === 'daily-stats')
    .map((attr) => attr.attribute as ProjectDailyStatsAttribute)
  const referenceLines = attributes.filter((attr) => attr?.provider === 'reference-line')

  const infraQueries = useInfraMonitoringQueries(
    infraAttributes,
    ref,
    startDate,
    endDate,
    interval,
    databaseIdentifier,
    data,
    isVisible
  )
  const dailyStatsQueries = useProjectDailyStatsQueries(
    dailyStatsAttributes,
    ref,
    startDate,
    endDate,
    data,
    isVisible
  )

  const referenceLineQueries = referenceLines.map((line) => {
    let value = line.value ?? line.customValue ?? 0

    return {
      data: {
        data: [],
        attribute: line.attribute,
        total: value,
        maximum: value,
        totalGrouped: { [line.attribute]: value },
      },
      isLoading: false,
      isError: false,
    }
  })

  return [...infraQueries, ...dailyStatsQueries, ...referenceLineQueries]
}

export function LazyComposedChartHandler(props: ComposedChartHandlerProps) {
  if (props.hide) return null

  return (
    <LazyChartWrapper>
      <ComposedChartHandler {...props} />
    </LazyChartWrapper>
  )
}