import { useCallback, useRef, useState } from 'react';
import type { RoomWsClient } from '@/lib/room-ws-client';

export interface AiStreamChunk {
  type: string;
  content: string;
  seq?: number;
}

export interface ActiveAiStream {
  message_id: string;
  display_name: string;
}

/**
 * Hook managing AI streaming state: streaming chunks, active stream indicator,
 * and stream cancellation. Separated from the main room context to reduce
 * the God component size (~1583 lines → ~300).
 */
export function useAiStreaming(clientRef: React.MutableRefObject<RoomWsClient | null>) {
  const [streamingChunks, setStreamingChunks] = useState<Map<string, AiStreamChunk[]>>(new Map());
  const [activeAiStream, setActiveAiStream] = useState<ActiveAiStream | null>(null);
  // Ref to latest chunks so done handler reads current state (setState is async)
  const chunksRef = useRef<Map<string, AiStreamChunk[]>>(new Map());

  const clearStreamingState = useCallback((msgId: string) => {
    setStreamingChunks(prev => { prev.delete(msgId); return new Map(prev); });
    chunksRef.current.delete(msgId);
    setActiveAiStream(null);
  }, []);

  const insertChunk = useCallback((
    messageId: string,
    chunkType: string | undefined,
    content: string,
    seq: number | undefined,
  ) => {
    setStreamingChunks(prev => {
      const next = new Map(prev);
      const existing: AiStreamChunk[] = next.get(messageId) ?? [];
      const s = seq ?? existing.length;
      const newChunk: AiStreamChunk = { type: chunkType ?? 'answer', content, seq: s };
      const insertIdx = existing.findIndex(c => c.seq != null && c.seq > s);
      next.set(messageId,
        insertIdx === -1
          ? [...existing, newChunk]
          : [...existing.slice(0, insertIdx), newChunk, ...existing.slice(insertIdx)]
      );
      chunksRef.current = new Map(next);
      return next;
    });
  }, []);

  const getOrderedChunks = useCallback((msgId: string): AiStreamChunk[] => {
    return chunksRef.current.get(msgId) ?? [];
  }, []);

  const cancelAiStream = useCallback(async () => {
    const client = clientRef.current;
    if (!client) return false;
    const roomId = client.getSubscribedRooms().values().next().value;
    if (!roomId) return false;
    return client.cancelAiStream(roomId as string);
  }, [clientRef]);

  return {
    streamingChunks,
    activeAiStream,
    setActiveAiStream,
    clearStreamingState,
    insertChunk,
    getOrderedChunks,
    cancelAiStream,
    chunksRef,
  };
}