"use client"

import { useState, useEffect, useRef, useMemo, useCallback, type ReactNode } from "react"
import { useLocation } from "wouter"
import { tabForPath, pathForTab, splitPath, joinPath } from "../lib/routes"
import type {
  MetricData,
  OtelHistoryMessage,
  OtelMessage,
  LogEntry,
  TraceData,
  ResourceSpan,
  ScopeSpan,
  Span,
} from "../lib/types"
import { ObservabilityContext, type TimeRangeFilter } from "./ObservabilityContextValue"
import { getServiceName, formatAttributeValue, isSpanError } from "../lib/otel-utils"
import { addFifoSet, appendFifo, partitionMessages, mergeLogsByTraceId, STREAM_LIMITS } from "../lib/ingest"
import { authFetch, tokens } from "../lib/auth"

interface ObservabilityProviderProps {
  children: ReactNode
}

export function ObservabilityProvider({
  children,
}: ObservabilityProviderProps) {
  const [logsByTraceId, setLogsByTraceId] = useState<
    Record<string, LogEntry[]>
  >({})
  const [traces, setTraces] = useState<TraceData[]>([])
  const [logs, setLogs] = useState<LogEntry[]>([])
  const [metrics, setMetrics] = useState<MetricData[]>([])

  const [isConnected, setIsConnected] = useState(false)
  const wsRef = useRef<WebSocket | null>(null)
  const reconnectTimeoutRef = useRef<number>(null)

  // --- Real Streaming / Pause Buffer State ---
  const [isStreaming, setIsStreaming] = useState(true)
  const isStreamingRef = useRef(true)
  const bufferedMessagesRef = useRef<OtelMessage[]>([])
  const [bufferedCount, setBufferedCount] = useState(0)

  // --- Issues live update counter ---
  const [issuesVersion, setIssuesVersion] = useState(0)

  // --- Latency filter (milliseconds) ---
  const [minDurationMs, setMinDurationMs] = useState<number>(0)

  // --- Persistence Keys ---
  const SERVICES_STORAGE_KEY = "otel_selected_services"
  const ATTRIBUTES_STORAGE_KEY = "otel_excluded_attributes"
  const SHOW_ERRORS_STORAGE_KEY = "otel_show_errors_only"
  const SEARCH_QUERY_STORAGE_KEY = "otel_search_query"
  const TIME_WINDOW_STORAGE_KEY = "otel_time_window"

  // --- Time Window Filter ---
  const [timeWindow, setTimeWindow] = useState<"15m" | "1h" | "6h" | "24h" | "all">(() => {
    if (typeof window === "undefined") return "15m"
    return (localStorage.getItem(TIME_WINDOW_STORAGE_KEY) as any) || "15m"
  })

  // --- Selected Time Period Filter (e.g. from clicking waveform) ---
  const [selectedTimeRange, setSelectedTimeRange] = useState<TimeRangeFilter | null>(null)

  // --- Filtering State ---
  const defaultExcludedAttributes = useCallback((): Record<string, Set<string>> => ({
    "user_agent.original": new Set([
      "kube-probe/1.33",
      "kube-probe/1.31",
      "Uptime-Kuma/1.23.11",
    ]),
    "http.target": new Set(["/healthz", "/readyz", "/livez"]),
    "http.host": new Set(["git.k6n.net"]),
    "http.user_agent": new Set(["Uptime-Kuma/1.23.11"]),
  }), [])

  const [selectedServices, setSelectedServices] = useState<Set<string>>(() => {
    if (typeof window === "undefined") return new Set()
    const saved = localStorage.getItem(SERVICES_STORAGE_KEY)
    return saved ? new Set(JSON.parse(saved)) : new Set()
  })

  const [excludedAttributes, setExcludedAttributes] = useState<Record<string, Set<string>>>(() => {
    if (typeof window === "undefined") return defaultExcludedAttributes()
    const saved = localStorage.getItem(ATTRIBUTES_STORAGE_KEY)
    if (!saved) return defaultExcludedAttributes()
    try {
      const parsed = JSON.parse(saved)
      const restored: Record<string, Set<string>> = {}
      Object.entries(parsed).forEach(([key, values]) => {
        restored[key] = new Set(values as string[])
      })
      return restored
    } catch (e) {
      console.error("Failed to restore excluded attributes", e)
      return defaultExcludedAttributes()
    }
  })

  const [showErrorsOnly, setShowErrorsOnly] = useState<boolean>(() => {
    if (typeof window === "undefined") return false
    const saved = localStorage.getItem(SHOW_ERRORS_STORAGE_KEY)
    return saved === "true"
  })

  const [searchQuery, setSearchQuery] = useState<string>(() => {
    if (typeof window === "undefined") return ""
    return localStorage.getItem(SEARCH_QUERY_STORAGE_KEY) || ""
  })

  // --- Persistence Effect ---
  useEffect(() => {
    if (typeof window === "undefined") return
    localStorage.setItem(SERVICES_STORAGE_KEY, JSON.stringify(Array.from(selectedServices)))

    const serializeable: Record<string, string[]> = {}
    Object.entries(excludedAttributes).forEach(([key, values]) => {
      serializeable[key] = Array.from(values)
    })
    localStorage.setItem(ATTRIBUTES_STORAGE_KEY, JSON.stringify(serializeable))
    localStorage.setItem(SHOW_ERRORS_STORAGE_KEY, showErrorsOnly.toString())
    localStorage.setItem(SEARCH_QUERY_STORAGE_KEY, searchQuery)
    localStorage.setItem(TIME_WINDOW_STORAGE_KEY, timeWindow)
  }, [selectedServices, excludedAttributes, showErrorsOnly, searchQuery, timeWindow])


  // --- Selection State ---
  const [activeTraceId, setActiveTraceId] = useState<string | null>(null)
  const [activeSpanId, setActiveSpanId] = useState<string | null>(null)
  // Tab state lives in the URL (wouter): /issues, /traces, ... so deep
  // links survive the OIDC login redirect and views are shareable. The
  // state mirror below keeps every existing setActiveTab call site working;
  // it follows the URL on navigation and pushes the URL on programmatic
  // setActiveTab calls.
  const [location, navigate] = useLocation()
  const [activeTab, setActiveTab] = useState<string>(
    () => tabForPath(window.location.pathname) ?? "tab0",
  )

  // External navigation (back/forward, manual URL edit, login redirect):
  // adopt the tab encoded in the new location.
  useEffect(() => {
    const tab = tabForPath(location)
    if (tab && tab !== activeTab) setActiveTab(tab)
    // eslint-disable-next-line react-hooks/exhaustive-deps -- activeTab is intentionally read stale here; see comment above
  }, [location])

  const setActiveTabWithRoute = useCallback(
    (tab: string) => {
      setActiveTab(tab)
      // Already inside this tab (possibly on a deep link): keep the sub-route.
      if (tabForPath(window.location.pathname) === tab) return
      const path = pathForTab(tab)
      if (path) navigate(path)
    },
    [navigate],
  )

  // Trace selection <-> /traces/<traceId>[/<spanId>]. The URL is adopted when
  // it names a trace; a bare /traces never clears the selection, because other
  // views set the trace first and switch tabs second.
  useEffect(() => {
    const { base, rest } = splitPath(location)
    if (base !== "/traces" || !rest[0]) return
    setActiveTraceId((prev) => (prev === rest[0] ? prev : rest[0]))
    setActiveSpanId((prev) => (prev === (rest[1] ?? null) ? prev : (rest[1] ?? null)))
    // eslint-disable-next-line react-hooks/exhaustive-deps -- only follow URL changes
  }, [location])

  useEffect(() => {
    if (activeTab !== "tab1" || !activeTraceId) return
    const target = joinPath("/traces", activeSpanId ? [activeTraceId, activeSpanId] : [activeTraceId])
    if (target !== window.location.pathname) navigate(target)
    // eslint-disable-next-line react-hooks/exhaustive-deps -- location is read imperatively
  }, [activeTab, activeTraceId, activeSpanId, navigate])
  const [filtersExpanded, setFiltersExpanded] = useState<boolean>(false)
  // Trace ids we already have spans for, so selecting a trace that fell
  // outside the initial history window (or arrived before this session
  // connected) triggers an on-demand GET /api/traces/{traceId} instead of
  // silently showing nothing.
  const [knownTraceIds, setKnownTraceIds] = useState<Set<string>>(new Set())

  const toggleService = useCallback((service: string) => {
    setSelectedServices(prev => {
      const next = new Set(prev)
      if (next.has(service)) {
        next.delete(service)
      } else {
        next.add(service)
      }
      return next
    })
  }, [])

  const toggleExcludedValue = useCallback((key: string, value: string) => {
    setExcludedAttributes((prev) => {
      const currentSet = prev[key] ?? new Set<string>()
      const nextSet = new Set(currentSet)
      if (nextSet.has(value)) {
        nextSet.delete(value)
      } else {
        nextSet.add(value)
      }
      const next = { ...prev }
      if (nextSet.size === 0) {
        delete next[key]
      } else {
        next[key] = nextSet
      }
      return next
    })
  }, [])

  const removeExcludedValue = useCallback((key: string, value: string) => {
    setExcludedAttributes((prev) => {
      const current = prev[key]
      if (!current || !current.has(value)) {
        return prev
      }
      const nextSet = new Set(current)
      nextSet.delete(value)
      const next = { ...prev }
      if (nextSet.size === 0) {
        delete next[key]
      } else {
        next[key] = nextSet
      }
      return next
    })
  }, [])

  const clearAttributeKey = useCallback((key: string) => {
    setExcludedAttributes((prev) => {
      if (!prev[key]) {
        return prev
      }
      const next = { ...prev }
      delete next[key]
      return next
    })
  }, [])

  const resetFilters = useCallback(() => {
    setSelectedServices(new Set())
    setExcludedAttributes(defaultExcludedAttributes())
    setSearchQuery("")
    setSelectedTimeRange(null)
  }, [defaultExcludedAttributes])

  // Flush buffered WebSocket stream entries when resuming
  const flushBufferedStream = useCallback(() => {
    const queued = bufferedMessagesRef.current
    bufferedMessagesRef.current = []
    setBufferedCount(0)
    if (queued.length > 0) {
      const partitioned = partitionMessages(queued)
      if (partitioned.traces.length) {
        setTraces((prev) => appendFifo(prev, partitioned.traces, STREAM_LIMITS.traces))
      }
      if (partitioned.metrics.length) {
        setMetrics((prev) => appendFifo(prev, partitioned.metrics, STREAM_LIMITS.metrics))
      }
      if (partitioned.logs.length) {
        setLogs((prev) => appendFifo(prev, partitioned.logs, STREAM_LIMITS.logs))
        setLogsByTraceId((prev) => mergeLogsByTraceId(prev, partitioned.logsByTraceId))
      }
      if (partitioned.traceIds.size) {
        setKnownTraceIds((prev) => addFifoSet(prev, partitioned.traceIds, STREAM_LIMITS.knownTraceIds))
      }
    }
  }, [])

  useEffect(() => {
    isStreamingRef.current = isStreaming
    if (isStreaming && bufferedMessagesRef.current.length > 0) {
      flushBufferedStream()
    }
  }, [isStreaming, flushBufferedStream])

  // Compute cutoff timestamp in nanoseconds for timeWindow
  const cutoffNano = useMemo(() => {
    if (timeWindow === "all") return 0n
    const now = Date.now()
    const windowMs = {
      "15m": 15 * 60 * 1000,
      "1h": 60 * 60 * 1000,
      "6h": 6 * 60 * 60 * 1000,
      "24h": 24 * 60 * 60 * 1000,
      "all": 0,
    }[timeWindow]
    return BigInt(now - windowMs) * 1_000_000n
  }, [timeWindow])

  // Discover all unique services across all signals
  const uniqueServices = useMemo(() => {
    const services = new Set<string>()

    // Check logs
    logs.forEach(log => {
      services.add(getServiceName(log.resource?.attributes ?? []))
      ;(log.scopeLogs ?? []).forEach(sl => (sl.logRecords ?? []).forEach(record => {
        services.add(getServiceName(log.resource?.attributes ?? [], record.attributes ?? []))
      }))
    })

    // Check traces
    traces.forEach(trace => {
      trace.resourceSpans?.forEach(resourceSpan => {
        services.add(getServiceName(resourceSpan.resource?.attributes ?? []))
        resourceSpan.scopeSpans?.forEach(ss => {
          ss.spans?.forEach(s => {
            services.add(getServiceName(resourceSpan.resource?.attributes ?? [], s.attributes ?? []))
          })
        })
      })
    })

    // Check metrics
    metrics.forEach(metric => {
      metric.resourceMetrics?.forEach(rm => {
        services.add(getServiceName(rm.resource?.attributes ?? []))
      })
    })

    return Array.from(services).sort().filter(s => s !== "unknown-service")
  }, [logs, traces, metrics])

  // --- Base Filtered Data (Global Filters: Service, TimeWindow, Search, Excluded Attributes) ---
  const unfocusedLogs = useMemo(() => {
    return logs.map(log => {
      const filteredScopeLogs = (log.scopeLogs ?? []).map(sl => {
        const filteredRecords = (sl.logRecords ?? []).filter(record => {
          const serviceName = getServiceName(log.resource?.attributes ?? [], record.attributes ?? [])
          // Filter by service
          if (selectedServices.size > 0 && !selectedServices.has(serviceName)) return false

          // Filter by time window
          if (cutoffNano > 0n) {
            const tNano = BigInt(record.timeUnixNano || "0")
            if (tNano > 0n && tNano < cutoffNano) return false
          }

          // Filter by attributes
          const recordAttrs = Object.fromEntries((record.attributes ?? []).map(a => [a.key, formatAttributeValue(a.value)]))
          const isExcluded = Object.entries(excludedAttributes).some(([key, values]) => {
            const val = recordAttrs[key]
            return val && values.has(val)
          })
          if (isExcluded) return false

          if (showErrorsOnly) {
            return (record.severityNumber ?? 0) >= 13 // 13 is Warn, 17 is Error
          }

          if (searchQuery) {
            const query = searchQuery.toLowerCase()
            const matchInBody = record.body?.stringValue?.toLowerCase()?.includes(query) ?? false
            const matchInAttrs = record.attributes?.some(a =>
              a.key.toLowerCase()?.includes(query) ||
              formatAttributeValue(a.value).toLowerCase()?.includes(query)
            ) ?? false
            const matchInServiceName = serviceName.toLowerCase()?.includes(query) ?? false
            const matchInTrace = (record.traceId || "").toLowerCase().includes(query)
            if (!matchInBody && !matchInAttrs && !matchInServiceName && !matchInTrace) return false
          }

          return true
        })
        if (filteredRecords.length === 0) return null
        return { ...sl, logRecords: filteredRecords }
      }).filter(Boolean)

      if (filteredScopeLogs.length === 0) return null
      return { ...log, scopeLogs: filteredScopeLogs as any }
    }).filter(Boolean) as LogEntry[]
  }, [logs, selectedServices, excludedAttributes, showErrorsOnly, searchQuery, cutoffNano])

  const unfocusedTraces = useMemo(() => {
    return traces.map(trace => {
      const filteredResourceSpans = (trace.resourceSpans?.map(rs => {
        const serviceName = getServiceName(rs.resource?.attributes ?? [])
        // Filter by service
        if (selectedServices.size > 0 && !selectedServices.has(serviceName)) return null

        const filteredScopeSpans = (rs.scopeSpans?.map(ss => {
          const filteredSpans = ss.spans?.filter(span => {
            // Filter by time window
            if (cutoffNano > 0n) {
              const tNano = BigInt(span.startTimeUnixNano || "0")
              if (tNano > 0n && tNano < cutoffNano) return false
            }

            // Filter by latency threshold (minDurationMs)
            if (minDurationMs > 0) {
              const start = Number(span.startTimeUnixNano || "0")
              const end = Number(span.endTimeUnixNano || span.startTimeUnixNano || "0")
              const durMs = (end - start) / 1_000_000
              if (durMs < minDurationMs) return false
            }

            // 1. Excluded attributes
            const spanAttrs = Object.fromEntries(span.attributes?.map(a => [a.key, formatAttributeValue(a.value)]) ?? [])
            const isExcluded = Object.entries(excludedAttributes).some(([key, values]) => {
              const val = spanAttrs[key]
              return val && values.has(val)
            })
            if (isExcluded) return false

            // 2. Search query
            if (searchQuery) {
              const query = searchQuery.toLowerCase()
              const matchInName = (span.name || "").toLowerCase()?.includes(query) ?? false
              const matchInTrace = (span.traceId || "").toLowerCase().includes(query)
              const matchInAttrs = span.attributes?.some(a =>
                a.key.toLowerCase()?.includes(query) ||
                formatAttributeValue(a.value).toLowerCase()?.includes(query)
              ) ?? false
              if (!matchInName && !matchInAttrs && !matchInTrace) return false
            }

            return true
          })
          if (!filteredSpans || filteredSpans.length === 0) return null
          return { ...ss, spans: filteredSpans }
        }).filter(Boolean) as ScopeSpan[]) || []

        if (!filteredScopeSpans || filteredScopeSpans.length === 0) return null
        return { ...rs, scopeSpans: filteredScopeSpans }
      }).filter(Boolean) as ResourceSpan[]) || []

      if (!filteredResourceSpans || filteredResourceSpans.length === 0) return null

      if (showErrorsOnly) {
        const hasError = filteredResourceSpans?.some((rs: ResourceSpan) =>
          rs?.scopeSpans?.some((ss: ScopeSpan) =>
            ss?.spans?.some((span: Span) => isSpanError(span))
          )
        )
        if (!hasError) return null
      }

      return { ...trace, resourceSpans: filteredResourceSpans }
    }).filter(Boolean) as TraceData[]
  }, [traces, selectedServices, excludedAttributes, showErrorsOnly, searchQuery, cutoffNano, minDurationMs])

  // --- Selected Time Period Filtering (e.g. from waveform section click) ---
  const filteredLogs = useMemo(() => {
    if (!selectedTimeRange) return unfocusedLogs
    const { startMs, endMs } = selectedTimeRange
    return unfocusedLogs.map(log => {
      const filteredScopeLogs = (log.scopeLogs ?? []).map(sl => {
        const filteredRecords = (sl.logRecords ?? []).filter(record => {
          const tMs = Number(BigInt(record.timeUnixNano || "0") / 1_000_000n)
          return tMs >= startMs && tMs <= endMs
        })
        if (filteredRecords.length === 0) return null
        return { ...sl, logRecords: filteredRecords }
      }).filter(Boolean)

      if (filteredScopeLogs.length === 0) return null
      return { ...log, scopeLogs: filteredScopeLogs as any }
    }).filter(Boolean) as LogEntry[]
  }, [unfocusedLogs, selectedTimeRange])

  const filteredTraces = useMemo(() => {
    if (!selectedTimeRange) return unfocusedTraces
    const { startMs, endMs } = selectedTimeRange
    return unfocusedTraces.filter(trace => {
      return (trace.resourceSpans ?? []).some(rs =>
        (rs.scopeSpans ?? []).some(ss =>
          (ss.spans ?? []).some(span => {
            const spanStartMs = Number(BigInt(span.startTimeUnixNano || "0") / 1_000_000n)
            const spanEndMs = Number(BigInt(span.endTimeUnixNano || span.startTimeUnixNano || "0") / 1_000_000n)
            return spanStartMs <= endMs && spanEndMs >= startMs
          })
        )
      )
    })
  }, [unfocusedTraces, selectedTimeRange])

  const logCountBySpanId = useMemo(() => {
    const counts: Record<string, number> = {}
    filteredLogs.forEach(log => {
      ;(log.scopeLogs ?? []).forEach(sl => {
        ;(sl.logRecords ?? []).forEach(record => {
          if (record.spanId) {
            counts[record.spanId] = (counts[record.spanId] || 0) + 1
          }
        })
      })
    })
    return counts
  }, [filteredLogs])

  const filteredMetrics = useMemo(() => {
    return metrics.map(data => {
      const filteredResourceMetrics = data.resourceMetrics?.map(rm => {
        const serviceName = rm.resource?.attributes?.find(attr => attr.key === "service.name")?.value.stringValue || "unknown-service"
        // Filter by service
        if (selectedServices.size > 0 && !selectedServices.has(serviceName)) return null

        const filteredScopeMetrics = rm.scopeMetrics?.map(sm => {
          const filteredMetricsList = sm.metrics?.filter(_m => {
            return true
          })
          if (filteredMetricsList?.length === 0) return null
          return { ...sm, metrics: filteredMetricsList }
        }).filter(Boolean)

        if (filteredScopeMetrics?.length === 0) return null
        return { ...rm, scopeMetrics: filteredScopeMetrics }
      }).filter(Boolean)

      if (filteredResourceMetrics?.length === 0) return null
      return { ...data, resourceMetrics: filteredResourceMetrics }
    }).filter(Boolean) as MetricData[]
  }, [metrics, selectedServices])

  // --- Real Service Health & Offender Metrics ---
  const serviceHealth = useMemo(() => {
    const health: Record<string, {
      name: string;
      totalSignals: number;
      errorCount: number;
      spansCount: number;
      logsCount: number;
      hasErrors: boolean;
    }> = {}

    const getOrInit = (srv: string) => {
      if (!health[srv]) {
        health[srv] = {
          name: srv,
          totalSignals: 0,
          errorCount: 0,
          spansCount: 0,
          logsCount: 0,
          hasErrors: false,
        }
      }
      return health[srv]
    }

    filteredLogs.forEach((log) => {
      (log.scopeLogs ?? []).forEach((sl) => {
        (sl.logRecords ?? []).forEach((r) => {
          const sName = getServiceName(log.resource?.attributes ?? [], r.attributes ?? [])
          if (sName && sName !== "unknown-service") {
            const h = getOrInit(sName)
            h.totalSignals++
            h.logsCount++
            if ((r.severityNumber ?? 0) >= 17) {
              h.errorCount++
              h.hasErrors = true
            }
          }
        })
      })
    })

    filteredTraces.forEach((trace) => {
      trace.resourceSpans?.forEach((rs) => {
        rs.scopeSpans?.forEach((ss) => {
          ss.spans?.forEach((s) => {
            const sName = getServiceName(rs.resource?.attributes ?? [], s.attributes ?? [])
            if (sName && sName !== "unknown-service") {
              const h = getOrInit(sName)
              h.totalSignals++
              h.spansCount++
              if (isSpanError(s)) {
                h.errorCount++
                h.hasErrors = true
              }
            }
          })
        })
      })
    })

    return health
  }, [filteredLogs, filteredTraces])


  useEffect(() => {
    // Do not start protected history/WebSocket traffic until a session exists.
    // The auth callback reloads the app after storing tokens, so this is also
    // sufficient for the normal login flow.
    if (!tokens()) return

    const fetchHistory = async () => {
      try {
        // Fetch each signal type separately with its own generous limit,
        // rather than one combined /api/history call - otherwise a noisy
        // signal type (metrics in particular) can crowd the others out of
        // the initial window.
        const [tracesRes, logsRes, metricsRes] = await Promise.all([
          authFetch("/api/history/traces?limit=300"),
          authFetch("/api/history/logs?limit=300"),
          authFetch("/api/history/metrics?limit=150"),
        ])
        const [tracesData, logsData, metricsData] = await Promise.all([
          tracesRes.json() as Promise<OtelHistoryMessage[] | null>,
          logsRes.json() as Promise<OtelHistoryMessage[] | null>,
          metricsRes.json() as Promise<OtelHistoryMessage[] | null>,
        ])

        const partitioned = partitionMessages([
          ...(tracesData ?? []),
          ...(logsData ?? []),
          ...(metricsData ?? []),
        ])

        console.log(
          "Fetched historical data:",
          partitioned.traces.length, "traces,",
          partitioned.logs.length, "logs,",
          partitioned.metrics.length, "metrics",
        )

        setTraces(partitioned.traces.slice(-STREAM_LIMITS.traces))
        setMetrics(partitioned.metrics.slice(-STREAM_LIMITS.metrics))
        setLogs(partitioned.logs.slice(-STREAM_LIMITS.logs))
        setLogsByTraceId(mergeLogsByTraceId({}, partitioned.logsByTraceId))
        setKnownTraceIds(addFifoSet(new Set(), partitioned.traceIds, STREAM_LIMITS.knownTraceIds))
      } catch (error) {
        console.error("Error fetching history:", error)
      }
    }

    const connect = () => {
      const protocol = window.location.protocol === "https:" ? "wss:" : "ws:"
      const t = tokens()
      const tokenQuery = t?.accessToken ? `?access_token=${encodeURIComponent(t.accessToken)}` : ""
      const wsUrl = `${protocol}//${window.location.host}/ws${tokenQuery}`
      console.log(`Connecting to WebSocket at ${wsUrl}`)

      const ws = new WebSocket(wsUrl)
      wsRef.current = ws

      ws.onopen = () => {
        console.log("WebSocket connected")
        setIsConnected(true)
      }

      ws.onmessage = (event) => {
        try {
          const data = JSON.parse(event.data)

          // Server issues notification broadcast
          if (data?.type === "issues") {
            setIssuesVersion((v) => v + 1)
            return
          }

          // If stream is paused, buffer the incoming signal
          if (!isStreamingRef.current) {
            bufferedMessagesRef.current = appendFifo(
              bufferedMessagesRef.current,
              [data as OtelMessage],
              STREAM_LIMITS.bufferedMessages,
            )
            setBufferedCount(bufferedMessagesRef.current.length)
            return
          }

          const partitioned = partitionMessages([data as OtelMessage])

          if (partitioned.traces.length) {
            setTraces((prev) => appendFifo(prev, partitioned.traces, STREAM_LIMITS.traces))
          }
          if (partitioned.metrics.length) {
            setMetrics((prev) => appendFifo(prev, partitioned.metrics, STREAM_LIMITS.metrics))
          }
          if (partitioned.logs.length) {
            setLogs((prev) => appendFifo(prev, partitioned.logs, STREAM_LIMITS.logs))
            setLogsByTraceId((prev) => mergeLogsByTraceId(prev, partitioned.logsByTraceId))
          }
          if (partitioned.traceIds.size) {
            setKnownTraceIds((prev) => addFifoSet(prev, partitioned.traceIds, STREAM_LIMITS.knownTraceIds))
          }
        } catch (error) {
          console.error("Error parsing WebSocket message:", error)
        }
      }

      ws.onerror = (error) => {
        console.error("WebSocket error:", error)
        setIsConnected(false)
      }

      ws.onclose = () => {
        console.log("WebSocket disconnected, reconnecting in 5s...")
        setIsConnected(false)

        // Reconnect after 5 seconds
        reconnectTimeoutRef.current = setTimeout(() => {
          connect()
        }, 5000)
      }
    }
    fetchHistory().then(() => connect())

    return () => {
      if (reconnectTimeoutRef.current) {
        clearTimeout(reconnectTimeoutRef.current)
      }
      if (wsRef.current) {
        wsRef.current.close()
      }
    }
  }, [flushBufferedStream])

  // On-demand trace fetch: if the selected trace id isn't one we already
  // have spans for (e.g. it arrived before this session connected, or fell
  // outside the initial history window), fetch it directly instead of
  // silently showing an empty trace view.
  useEffect(() => {
    if (!tokens() || !activeTraceId || knownTraceIds.has(activeTraceId)) return

    let cancelled = false
    authFetch(`/api/traces/${encodeURIComponent(activeTraceId)}`)
      .then((res) => {
        if (!res.ok) throw new Error(`HTTP ${res.status}`)
        return res.json() as Promise<OtelHistoryMessage[] | null>
      })
      .then((data) => {
        if (cancelled) return
        const partitioned = partitionMessages(data ?? [])

        if (partitioned.traces.length) {
          setTraces((prev) => appendFifo(prev, partitioned.traces, STREAM_LIMITS.traces))
        }
        if (partitioned.metrics.length) {
          setMetrics((prev) => appendFifo(prev, partitioned.metrics, STREAM_LIMITS.metrics))
        }
        if (partitioned.logs.length) {
          setLogs((prev) => appendFifo(prev, partitioned.logs, STREAM_LIMITS.logs))
          setLogsByTraceId((prev) => mergeLogsByTraceId(prev, partitioned.logsByTraceId))
        }

        setKnownTraceIds((prev) => {
          return addFifoSet(
            prev,
            new Set([...partitioned.traceIds, activeTraceId]),
            STREAM_LIMITS.knownTraceIds,
          )
        })
      })
      .catch((error) => {
        console.error(`Error fetching trace ${activeTraceId}:`, error)
      })

    return () => {
      cancelled = true
    }
  }, [activeTraceId, knownTraceIds])

  return (
    <ObservabilityContext.Provider
      value={{
        logs: filteredLogs,
        metrics: filteredMetrics,
        traces: filteredTraces,
        allTraces: traces,
        unfocusedLogs,
        unfocusedTraces,
        logsByTraceId,
        logCountBySpanId,
        isConnected,

        // Real stream controls
        isStreaming,
        setIsStreaming,
        bufferedCount,
        flushBufferedStream,

        // Real time window & latency
        timeWindow,
        setTimeWindow,
        selectedTimeRange,
        setSelectedTimeRange,
        minDurationMs,
        setMinDurationMs,

        // Discovered service health & issues updates
        serviceHealth,
        issuesVersion,

        selectedServices,
        toggleService,
        setSelectedServices,
        excludedAttributes,
        toggleExcludedValue,
        removeExcludedValue,
        clearAttributeKey,
        resetFilters,
        uniqueServices,
        showErrorsOnly,
        setShowErrorsOnly,
        searchQuery,
        setSearchQuery,
        activeTraceId,
        setActiveTraceId,
        activeSpanId,
        setActiveSpanId,
        activeTab,
        setActiveTab: setActiveTabWithRoute,
        filtersExpanded,
        setFiltersExpanded,
      }}
    >
      {children}
    </ObservabilityContext.Provider>
  )

}