maternal-app/maternal-web/hooks/useStreamingChat.ts

import { useState, useCallback } from 'react';
import apiClient from '@/lib/api/client';

export interface StreamChunk {
  type: 'token' | 'metadata' | 'done' | 'error';
  content?: string;
  metadata?: any;
  error?: string;
}

export interface ChatMessageDto {
  message: string;
  conversationId?: string;
  language?: string;
}

/**
 * Hook for streaming AI chat responses
 * Uses Server-Sent Events (SSE) for real-time token streaming
 */
export function useStreamingChat() {
  const [isStreaming, setIsStreaming] = useState(false);
  const [error, setError] = useState<string | null>(null);

  const streamMessage = useCallback(
    async (
      chatDto: ChatMessageDto,
      onChunk: (chunk: StreamChunk) => void,
      onComplete?: () => void,
      onError?: (error: string) => void
    ) => {
      setIsStreaming(true);
      setError(null);

      try {
        const response = await fetch(`${process.env.NEXT_PUBLIC_API_URL}/api/v1/ai/chat/stream`, {
          method: 'POST',
          headers: {
            'Content-Type': 'application/json',
            // Add auth token if available
            ...(typeof window !== 'undefined' && localStorage.getItem('accessToken')
              ? { Authorization: `Bearer ${localStorage.getItem('accessToken')}` }
              : {}),
          },
          body: JSON.stringify(chatDto),
        });

        if (!response.ok) {
          throw new Error(`HTTP error! status: ${response.status}`);
        }

        const reader = response.body?.getReader();
        if (!reader) {
          throw new Error('Response body is not readable');
        }

        const decoder = new TextDecoder();
        let buffer = '';

        while (true) {
          const { done, value } = await reader.read();

          if (done) {
            break;
          }

          // Decode the chunk and add to buffer
          buffer += decoder.decode(value, { stream: true });

          // Split by newlines to process complete SSE events
          const lines = buffer.split('\n');

          // Keep the last incomplete line in the buffer
          buffer = lines.pop() || '';

          for (const line of lines) {
            const trimmed = line.trim();

            // Skip empty lines
            if (!trimmed) {
              continue;
            }

            // Parse SSE data format
            if (trimmed.startsWith('data: ')) {
              const data = trimmed.substring(6);

              try {
                const chunk: StreamChunk = JSON.parse(data);

                // Emit the chunk
                onChunk(chunk);

                // Check for completion
                if (chunk.type === 'done') {
                  setIsStreaming(false);
                  if (onComplete) {
                    onComplete();
                  }
                  return;
                }

                // Check for errors
                if (chunk.type === 'error') {
                  setIsStreaming(false);
                  const errorMsg = chunk.error || 'Streaming error occurred';
                  setError(errorMsg);
                  if (onError) {
                    onError(errorMsg);
                  }
                  return;
                }
              } catch (parseError) {
                console.error('Failed to parse SSE chunk:', parseError);
              }
            }
          }
        }
      } catch (err) {
        const errorMsg = err instanceof Error ? err.message : 'Streaming failed';
        setError(errorMsg);
        setIsStreaming(false);
        if (onError) {
          onError(errorMsg);
        }
      }
    },
    []
  );

  return {
    streamMessage,
    isStreaming,
    error,
  };
}