road-monitoring-ui/src/components/video/useVideoAnnotationPlayback.ts

'use client';

import { useQueryClient } from '@tanstack/react-query';
import {
  RefObject,
  useCallback,
  useEffect,
  useMemo,
  useRef,
  useState,
} from 'react';
import type { MediaPlayerInstance } from '@vidstack/react';

import { useVideoAnnotationFramesQuery } from '@/app/(modules)/results/hooks/useVideoResults';
import { videoKeys } from '@/app/(modules)/results/queries/videoKeys';
import { videoService } from '@/services/api';
import type {
  DetectionResultLog,
  VideoAnnotationFrameDetection,
  VideoAnnotationFramesParams,
  VideoAnnotationFramesResponse,
} from '@/types';

const ANNOTATION_WINDOW_MS = 30000;
const ANNOTATION_PREFETCH_THRESHOLD_MS = 5000;

const buildWindowParams = (
  startTimeMs: number,
): VideoAnnotationFramesParams => ({
  start_time_ms: startTimeMs,
  end_time_ms: startTimeMs + ANNOTATION_WINDOW_MS,
});

const getWindowStartMs = (currentTimeMs: number) =>
  Math.floor(Math.max(currentTimeMs, 0) / ANNOTATION_WINDOW_MS) *
  ANNOTATION_WINDOW_MS;

interface UseVideoAnnotationPlaybackParams {
  videoId: string | undefined;
  logs: DetectionResultLog[];
  fps: number;
  videoRef: RefObject<MediaPlayerInstance | null>;
}

export function useVideoAnnotationPlayback({
  videoId,
  logs,
  fps,
  videoRef,
}: UseVideoAnnotationPlaybackParams) {
  const queryClient = useQueryClient();
  const shouldPauseAfterSeekRef = useRef(false);
  const visibleDetectionIdsRef = useRef('');
  const [currentWindowStartMs, setCurrentWindowStartMs] = useState(0);
  const [activeLog, setActiveLog] = useState<DetectionResultLog | undefined>();
  const [visibleDetections, setVisibleDetections] = useState<
    VideoAnnotationFrameDetection[]
  >([]);

  const sortedLogs = useMemo(
    () =>
      [...logs].sort(
        (a, b) => a.frame.timestamp_seconds - b.frame.timestamp_seconds,
      ),
    [logs],
  );
  const currentWindowParams = useMemo(
    () => buildWindowParams(currentWindowStartMs),
    [currentWindowStartMs],
  );
  const annotationFramesQuery = useVideoAnnotationFramesQuery(
    videoId,
    currentWindowParams,
  );

  useEffect(() => {
    setActiveLog((current) => {
      if (current && sortedLogs.some((log) => log.id === current.id)) {
        return current;
      }

      return sortedLogs[0];
    });
  }, [sortedLogs]);

  const prefetchWindow = useCallback(
    (startTimeMs: number) => {
      if (!videoId || startTimeMs < 0) return;

      const params = buildWindowParams(startTimeMs);

      void queryClient.prefetchQuery({
        queryKey: videoKeys.annotationFrames(videoId, params),
        queryFn: () => videoService.getVideoAnnotationFrames(videoId, params),
        staleTime: 1000 * 60 * 5,
        gcTime: 1000 * 60 * 30,
      });
    },
    [queryClient, videoId],
  );

  const getWindowData = useCallback(
    (startTimeMs: number) => {
      if (!videoId) return undefined;

      const params = buildWindowParams(startTimeMs);

      return (
        queryClient.getQueryData<VideoAnnotationFramesResponse>(
          videoKeys.annotationFrames(videoId, params),
        ) ??
        (startTimeMs === currentWindowStartMs
          ? annotationFramesQuery.data
          : undefined)
      );
    },
    [annotationFramesQuery.data, currentWindowStartMs, queryClient, videoId],
  );

  const updateActiveLog = useCallback(
    (currentTimeSeconds: number) => {
      if (sortedLogs.length === 0) {
        setActiveLog(undefined);
        return;
      }

      const currentLog = sortedLogs.findLast(
        (log) => log.frame.timestamp_seconds <= currentTimeSeconds,
      );
      const nextActiveLog = currentLog ?? sortedLogs[0];

      setActiveLog((previous) =>
        previous?.id === nextActiveLog.id ? previous : nextActiveLog,
      );
    },
    [sortedLogs],
  );

  const getVisibleDetections = useCallback(
    (currentTimeSeconds: number) => {
      const currentTimeMs = currentTimeSeconds * 1000;
      const targetWindowStartMs = getWindowStartMs(currentTimeMs);

      if (targetWindowStartMs !== currentWindowStartMs) {
        setCurrentWindowStartMs(targetWindowStartMs);
      }

      const windowData = getWindowData(targetWindowStartMs);

      if (!windowData) {
        return [];
      }

      if (
        windowData.window.has_more_after &&
        currentTimeMs >=
          windowData.window.end_time_ms - ANNOTATION_PREFETCH_THRESHOLD_MS
      ) {
        prefetchWindow(windowData.window.end_time_ms);
      }

      const frameDurationMs = fps > 0 ? 1000 / fps : 0;
      const currentFrame = windowData.frames.find((frame) => {
        const frameEndMs =
          frame.end_timestamp_ms ?? frame.timestamp_ms + frameDurationMs;

        return (
          currentTimeMs >= frame.timestamp_ms && currentTimeMs < frameEndMs
        );
      });

      return currentFrame?.detections ?? [];
    },
    [currentWindowStartMs, fps, getWindowData, prefetchWindow],
  );

  const updateVisibleDetections = useCallback(
    (currentTimeSeconds: number) => {
      const nextDetections = getVisibleDetections(currentTimeSeconds);
      const nextIds = nextDetections.map((detection) => detection.id).join('|');

      if (nextIds === visibleDetectionIdsRef.current) return;

      visibleDetectionIdsRef.current = nextIds;
      setVisibleDetections(nextDetections);
    },
    [getVisibleDetections],
  );

  const handleSeek = useCallback(
    (log: DetectionResultLog) => {
      const video = videoRef.current;
      if (!video) return;

      const targetTimeSeconds = log.frame.timestamp_seconds;
      const targetWindowStartMs = getWindowStartMs(targetTimeSeconds * 1000);

      prefetchWindow(targetWindowStartMs);
      setCurrentWindowStartMs(targetWindowStartMs);
      shouldPauseAfterSeekRef.current = true;
      video.pause();
      video.currentTime = targetTimeSeconds;
      setActiveLog(log);
      updateVisibleDetections(targetTimeSeconds);
    },
    [prefetchWindow, updateVisibleDetections, videoRef],
  );

  const handleSeeked = useCallback(() => {
    const video = videoRef.current;
    if (!video || !shouldPauseAfterSeekRef.current) return;

    shouldPauseAfterSeekRef.current = false;
    video.pause();
  }, [videoRef]);

  const handleTimeUpdate = useCallback(() => {
    const video = videoRef.current;
    if (!video) return;

    updateActiveLog(video.currentTime);
  }, [updateActiveLog, videoRef]);

  const handleVideoFrame = useCallback(
    (mediaTime: number) => {
      updateActiveLog(mediaTime);
      updateVisibleDetections(mediaTime);
    },
    [updateActiveLog, updateVisibleDetections],
  );

  return {
    activeLog,
    annotationFramesQuery,
    handleSeek,
    handleSeeked,
    handleTimeUpdate,
    handleVideoFrame,
    sortedLogs,
    visibleDetections,
  };
}