gooseek/backend/webui/src/lib/hooks/useChat.ts

'use client';

import { useState, useCallback, useRef } from 'react';
import type { Message, Citation, Widget, StreamEvent, ChatAttachmentInfo } from '../types';
import { streamChat, generateId } from '../api';
import { getLanguageFromStorage } from '../contexts/LanguageContext';
import type { SendOptions, ChatAttachment } from '@/components/ChatInput';

function fileToDataUrl(file: File): Promise<string> {
  return new Promise((resolve, reject) => {
    const reader = new FileReader();
    reader.onload = () => resolve(reader.result as string);
    reader.onerror = reject;
    reader.readAsDataURL(file);
  });
}

interface UseChatOptions {
  onError?: (error: Error) => void;
}

export function useChat(options: UseChatOptions = {}) {
  const [messages, setMessages] = useState<Message[]>([]);
  const [isLoading, setIsLoading] = useState(false);
  const [citations, setCitations] = useState<Citation[]>([]);
  const [widgets, setWidgets] = useState<Widget[]>([]);
  const abortControllerRef = useRef<AbortController | null>(null);

  const sendMessage = useCallback(async (content: string, sendOptions?: SendOptions | 'speed' | 'balanced' | 'quality') => {
    const mode = typeof sendOptions === 'string' ? sendOptions : sendOptions?.mode ?? 'balanced';
    const webSearch = typeof sendOptions === 'object' ? sendOptions.webSearch : false;
    const attachments = typeof sendOptions === 'object' ? sendOptions.attachments : [];
    const locale = getLanguageFromStorage();
    if (!content.trim() || isLoading) return;

    const userMessage: Message = {
      id: generateId(),
      role: 'user',
      content: content.trim(),
      createdAt: new Date(),
    };

    const assistantMessage: Message = {
      id: generateId(),
      role: 'assistant',
      content: '',
      citations: [],
      widgets: [],
      isStreaming: true,
      createdAt: new Date(),
    };

    setMessages((prev) => [...prev, userMessage, assistantMessage]);
    setIsLoading(true);
    setCitations([]);
    setWidgets([]);

    const history: [string, string][] = messages
      .filter((m) => !m.isStreaming)
      .reduce((acc, m, i, arr) => {
        if (m.role === 'user' && arr[i + 1]?.role === 'assistant') {
          acc.push([m.content, arr[i + 1].content]);
        }
        return acc;
      }, [] as [string, string][]);

    const attachmentInfos = await Promise.all(
      attachments.map(async (att) => {
        let dataUrl: string | undefined;
        if (att.type === 'image' || att.file.type.startsWith('image/')) {
          dataUrl = await fileToDataUrl(att.file);
        }
        return {
          name: att.file.name,
          type: att.file.type,
          size: att.file.size,
          dataUrl,
        };
      })
    );

    try {
      const stream = streamChat({
        message: {
          messageId: assistantMessage.id,
          chatId: userMessage.id,
          content: content.trim(),
        },
        optimizationMode: mode,
        history,
        locale,
        webSearch,
        attachments: attachmentInfos.length > 0 ? attachmentInfos : undefined,
      });

      let fullContent = '';
      const collectedCitations: Citation[] = [];
      const collectedWidgets: Widget[] = [];

      for await (const event of stream) {
        switch (event.type) {
          case 'messageStart':
            break;

          case 'message':
            if (event.data && typeof event.data === 'object' && 'content' in event.data) {
              const chunk = (event.data as { content: string }).content;
              fullContent += chunk;
              setMessages((prev) =>
                prev.map((m) =>
                  m.id === assistantMessage.id
                    ? { ...m, content: fullContent }
                    : m
                )
              );
            }
            break;

          case 'block':
            // Handle new block format from agent-svc
            if (event.block && typeof event.block === 'object') {
              const block = event.block as { type: string; data: unknown };
              if (block.type === 'text' && typeof block.data === 'string') {
                fullContent = block.data;
                setMessages((prev) =>
                  prev.map((m) =>
                    m.id === assistantMessage.id
                      ? { ...m, content: fullContent }
                      : m
                  )
                );
              } else if (block.type === 'widget' && block.data) {
                const widget = block.data as Widget;
                collectedWidgets.push(widget);
                setWidgets([...collectedWidgets]);
              }
            }
            break;

          case 'textChunk':
            // Handle streaming text chunks
            if (event.chunk && typeof event.chunk === 'string') {
              fullContent += event.chunk;
              setMessages((prev) =>
                prev.map((m) =>
                  m.id === assistantMessage.id
                    ? { ...m, content: fullContent }
                    : m
                )
              );
            }
            break;

          case 'updateBlock':
            // Handle block updates (final text)
            if (event.patch && Array.isArray(event.patch)) {
              for (const p of event.patch) {
                if (p.op === 'replace' && p.path === '/data' && typeof p.value === 'string') {
                  fullContent = p.value;
                  setMessages((prev) =>
                    prev.map((m) =>
                      m.id === assistantMessage.id
                        ? { ...m, content: fullContent }
                        : m
                    )
                  );
                }
              }
            }
            break;

          case 'sources':
            if (event.data && Array.isArray(event.data)) {
              const sources = event.data as Array<{
                url: string;
                title?: string;
                metadata?: { title?: string };
              }>;
              sources.forEach((source, idx) => {
                const citation: Citation = {
                  index: idx + 1,
                  url: source.url,
                  title: source.title || source.metadata?.title || new URL(source.url).hostname,
                  domain: new URL(source.url).hostname,
                };
                collectedCitations.push(citation);
              });
              setCitations([...collectedCitations]);
            }
            break;

          case 'widget':
            if (event.data && typeof event.data === 'object') {
              const widget = event.data as Widget;
              collectedWidgets.push(widget);
              setWidgets([...collectedWidgets]);
            }
            break;

          case 'researchComplete':
            // Research phase complete, text generation starting
            break;

          case 'messageEnd':
            setMessages((prev) =>
              prev.map((m) =>
                m.id === assistantMessage.id
                  ? {
                      ...m,
                      content: fullContent,
                      citations: collectedCitations,
                      widgets: collectedWidgets,
                      isStreaming: false,
                    }
                  : m
              )
            );
            break;

          case 'error':
            throw new Error(
              event.data && typeof event.data === 'object' && 'message' in event.data
                ? String((event.data as { message: string }).message)
                : 'Unknown error'
            );
        }
      }
    } catch (error) {
      const errorMessage = error instanceof Error ? error.message : 'Ошибка при отправке сообщения';
      setMessages((prev) =>
        prev.map((m) =>
          m.id === assistantMessage.id
            ? { ...m, content: `Ошибка: ${errorMessage}`, isStreaming: false }
            : m
        )
      );
      options.onError?.(error instanceof Error ? error : new Error(errorMessage));
    } finally {
      setIsLoading(false);
    }
  }, [isLoading, messages, options]);

  const stopGeneration = useCallback(() => {
    abortControllerRef.current?.abort();
    setIsLoading(false);
    setMessages((prev) =>
      prev.map((m) => (m.isStreaming ? { ...m, isStreaming: false } : m))
    );
  }, []);

  const clearMessages = useCallback(() => {
    setMessages([]);
    setCitations([]);
    setWidgets([]);
  }, []);

  return {
    messages,
    isLoading,
    citations,
    widgets,
    sendMessage,
    stopGeneration,
    clearMessages,
  };
}