tubearr/src/frontend/src/contexts/DownloadProgressContext.tsx

import { createContext, useContext, useCallback, useRef, type ReactNode } from 'react';
import { useQueryClient, type QueryClient } from '@tanstack/react-query';
import { useSyncExternalStore } from 'react';
import { useWebSocket } from '../hooks/useWebSocket';
import { contentKeys } from '../api/hooks/useContent';
import type { ContentItem } from '@shared/types/index';
import type { PaginatedResponse } from '@shared/types/api';

// ── Types ──

export interface ProgressInfo {
  percent: number;
  speed: string;
  eta: string;
}

interface DownloadProgressEvent {
  type: 'download:progress';
  contentItemId: number;
  percent: number;
  speed: string;
  eta: string;
}

interface DownloadCompleteEvent {
  type: 'download:complete';
  contentItemId: number;
}

interface DownloadFailedEvent {
  type: 'download:failed';
  contentItemId: number;
  error: string;
}

type DownloadEvent = DownloadProgressEvent | DownloadCompleteEvent | DownloadFailedEvent;

// ── Scan Event Types ──

interface ScanStartedEvent {
  type: 'scan:started';
  channelId: number;
  channelName: string;
}

interface ScanItemDiscoveredEvent {
  type: 'scan:item-discovered';
  channelId: number;
  channelName: string;
  item: ContentItem;
}

interface ScanCompleteEvent {
  type: 'scan:complete';
  channelId: number;
  channelName: string;
  newItems: number;
  totalFetched: number;
}

interface ScanErrorEvent {
  type: 'scan:error';
  channelId: number;
  channelName: string;
  error: string;
}

type ScanEvent = ScanStartedEvent | ScanItemDiscoveredEvent | ScanCompleteEvent | ScanErrorEvent;

// ── Download Progress Store (external to React for zero unnecessary re-renders) ──

class ProgressStore {
  private _map = new Map<number, ProgressInfo>();
  private _listeners = new Set<() => void>();

  subscribe = (listener: () => void) => {
    this._listeners.add(listener);
    return () => this._listeners.delete(listener);
  };

  getSnapshot = () => this._map;

  set(id: number, info: ProgressInfo) {
    this._map = new Map(this._map);
    this._map.set(id, info);
    this._notify();
  }

  delete(id: number) {
    if (!this._map.has(id)) return;
    this._map = new Map(this._map);
    this._map.delete(id);
    this._notify();
  }

  private _notify() {
    for (const listener of this._listeners) listener();
  }
}

// ── Scan Progress Store ──

export interface ScanProgress {
  scanning: boolean;
  newItemCount: number;
}

class ScanStore {
  private _map = new Map<number, ScanProgress>();
  private _listeners = new Set<() => void>();

  subscribe = (listener: () => void) => {
    this._listeners.add(listener);
    return () => this._listeners.delete(listener);
  };

  getSnapshot = () => this._map;

  startScan(channelId: number) {
    this._map = new Map(this._map);
    this._map.set(channelId, { scanning: true, newItemCount: 0 });
    this._notify();
  }

  incrementItems(channelId: number) {
    this._map = new Map(this._map);
    const current = this._map.get(channelId) ?? { scanning: true, newItemCount: 0 };
    this._map.set(channelId, { ...current, newItemCount: current.newItemCount + 1 });
    this._notify();
  }

  completeScan(channelId: number) {
    this._map = new Map(this._map);
    const current = this._map.get(channelId);
    if (current) {
      this._map.set(channelId, { scanning: false, newItemCount: current.newItemCount });
    }
    this._notify();
  }

  clearScan(channelId: number) {
    if (!this._map.has(channelId)) return;
    this._map = new Map(this._map);
    this._map.delete(channelId);
    this._notify();
  }

  private _notify() {
    for (const listener of this._listeners) listener();
  }
}

// ── Context ──

interface DownloadProgressContextValue {
  /** Get progress for a specific content item. Returns undefined if not downloading. */
  getProgress: (contentItemId: number) => ProgressInfo | undefined;
  /** Whether the WebSocket is connected */
  isConnected: boolean;
  /** Subscribe to scan store changes */
  scanStoreSubscribe: (listener: () => void) => () => void;
  /** Get scan store snapshot */
  scanStoreGetSnapshot: () => Map<number, ScanProgress>;
}

const DownloadProgressContext = createContext<DownloadProgressContextValue | null>(null);

// ── Provider ──

export function DownloadProgressProvider({ children }: { children: ReactNode }) {
  const queryClient = useQueryClient();
  const storeRef = useRef(new ProgressStore());
  const store = storeRef.current;
  const scanStoreRef = useRef(new ScanStore());
  const scanStore = scanStoreRef.current;

  // Subscribe to the store with useSyncExternalStore for optimal re-renders
  const progressMap = useSyncExternalStore(store.subscribe, store.getSnapshot);

  const handleMessage = useCallback(
    (data: unknown) => {
      const event = data as DownloadEvent | ScanEvent;
      if (!event?.type) return;

      switch (event.type) {
        case 'download:progress':
          store.set(event.contentItemId, {
            percent: event.percent,
            speed: event.speed,
            eta: event.eta,
          });
          break;

        case 'download:complete':
          store.delete(event.contentItemId);
          // Invalidate content queries so the UI refreshes with updated status
          queryClient.invalidateQueries({ queryKey: ['content'] });
          queryClient.invalidateQueries({ queryKey: ['queue'] });
          break;

        case 'download:failed':
          store.delete(event.contentItemId);
          // Invalidate to show updated status (failed)
          queryClient.invalidateQueries({ queryKey: ['content'] });
          queryClient.invalidateQueries({ queryKey: ['queue'] });
          break;

        case 'scan:started':
          scanStore.startScan(event.channelId);
          break;

        case 'scan:item-discovered':
          scanStore.incrementItems(event.channelId);
          injectContentItemIntoCache(queryClient, event.channelId, event.item);
          break;

        case 'scan:complete':
          scanStore.completeScan(event.channelId);
          // Safety net: reconcile any missed items
          queryClient.invalidateQueries({
            queryKey: contentKeys.byChannel(event.channelId),
          });
          break;

        case 'scan:error':
          scanStore.completeScan(event.channelId);
          break;
      }
    },
    [store, scanStore, queryClient],
  );

  const { isConnected } = useWebSocket({ onMessage: handleMessage });

  const getProgress = useCallback(
    (contentItemId: number): ProgressInfo | undefined => {
      return progressMap.get(contentItemId);
    },
    [progressMap],
  );

  return (
    <DownloadProgressContext.Provider
      value={{
        getProgress,
        isConnected,
        scanStoreSubscribe: scanStore.subscribe,
        scanStoreGetSnapshot: scanStore.getSnapshot,
      }}
    >
      {children}
    </DownloadProgressContext.Provider>
  );
}

// ── Hook ──

/**
 * Get download progress for a specific content item.
 * Returns undefined when the item is not actively downloading via WebSocket.
 */
export function useDownloadProgress(contentItemId: number): ProgressInfo | undefined {
  const context = useContext(DownloadProgressContext);
  if (!context) {
    throw new Error('useDownloadProgress must be used within a DownloadProgressProvider');
  }
  return context.getProgress(contentItemId);
}

/**
 * Get the WebSocket connection status.
 */
export function useDownloadProgressConnection(): boolean {
  const context = useContext(DownloadProgressContext);
  return context?.isConnected ?? false;
}

// ── Scan Progress Hook ──

/**
 * Get scan progress for a specific channel.
 * Returns `{ scanning, newItemCount }` from the scan store via useSyncExternalStore.
 * Only re-renders components that use this hook when the scan store changes.
 */
export function useScanProgress(channelId: number): ScanProgress {
  const context = useContext(DownloadProgressContext);
  if (!context) {
    throw new Error('useScanProgress must be used within a DownloadProgressProvider');
  }
  const scanMap = useSyncExternalStore(
    context.scanStoreSubscribe,
    context.scanStoreGetSnapshot,
  );
  return scanMap.get(channelId) ?? { scanning: false, newItemCount: 0 };
}

// ── Cache Injection Helper ──

/**
 * Inject a newly discovered content item into all matching TanStack Query caches
 * for the given channel. Prepends the item to page 1 queries and increments pagination counts.
 */
function injectContentItemIntoCache(
  queryClient: QueryClient,
  channelId: number,
  item: ContentItem,
) {
  queryClient.setQueriesData<PaginatedResponse<ContentItem>>(
    { queryKey: contentKeys.byChannel(channelId) },
    (oldData) => {
      if (!oldData?.data) return oldData;
      // Avoid duplicates
      if (oldData.data.some((existing) => existing.id === item.id)) return oldData;
      return {
        ...oldData,
        data: [item, ...oldData.data],
        pagination: {
          ...oldData.pagination,
          totalItems: oldData.pagination.totalItems + 1,
          totalPages: Math.ceil(
            (oldData.pagination.totalItems + 1) / oldData.pagination.pageSize,
          ),
        },
      };
    },
  );
}