LocalMind/hooks/useLMStudio.ts

import { useCallback, useRef } from 'react';
import { SYSTEM_PROMPT } from '../constants/SystemPrompt';
import * as api from '../services/lmStudioApi';
import { ChatMessage as ApiMessage } from '../services/types';
import { useAppStore } from '../store/useAppStore';
import { WebSearchLogger } from '../services/webSearchLogger';

export function useLMStudio() {
    const serverConfig = useAppStore((s) => s.serverConfig);
    const isConnected = useAppStore((s) => s.isConnected);
    const setIsConnected = useAppStore((s) => s.setIsConnected);
    const availableModels = useAppStore((s) => s.availableModels);
    const setAvailableModels = useAppStore((s) => s.setAvailableModels);
    const activeModelId = useAppStore((s) => s.activeModelId);
    const setActiveModelId = useAppStore((s) => s.setActiveModelId);
    const messages = useAppStore((s) => s.messages);
    const addMessage = useAppStore((s) => s.addMessage);
    const appendToMessage = useAppStore((s) => s.appendToMessage);
    const isStreaming = useAppStore((s) => s.isStreaming);
    const setIsStreaming = useAppStore((s) => s.setIsStreaming);

    const abortRef = useRef<{ abort: () => void } | null>(null);

    const checkConnection = useCallback(async () => {
        const ok = await api.testConnection(serverConfig);
        setIsConnected(ok);
        return ok;
    }, [serverConfig, setIsConnected]);

    const fetchModels = useCallback(async () => {
        try {
            const currentConfig = useAppStore.getState().serverConfig;
            const models = await api.getModels(currentConfig);
            setAvailableModels(models);
            setIsConnected(true);

            if (!activeModelId && models.length > 0) {
                setActiveModelId(models[0].id);
            }
            return models;
        } catch {
            setIsConnected(false);
            setAvailableModels([]);
            return [];
        }
    }, [activeModelId, setAvailableModels, setIsConnected, setActiveModelId]);

    const sendMessage = useCallback(async (text: string) => {
        if (isStreaming) return;

        const pipelineStart = Date.now();
        const store = useAppStore.getState();
        const isWebSearchEnabled = store.isWebSearchEnabled;

        addMessage({ role: 'user', content: text });

        const modelId = store.activeModelId;
        if (!modelId) {
            addMessage({
                role: 'assistant',
                content: '⚠️ No model selected. Please select a model first.',
            });
            return;
        }

        const assistantId = addMessage({ role: 'assistant', content: '' });
        setIsStreaming(true);

        try {
            let searchContext = '';
            let searchSources: Array<{ title: string; url: string }> = [];

            if (isWebSearchEnabled) {
                // ── Step 1: Build search service URL ─────────────────────────────
                // ServerConfig has { host, port } — construct the search service URL
                // from the same host, but on port 3030.
                const rawHost = store.serverConfig.host.replace(/\/+$/, '');
                const hostWithProtocol = rawHost.startsWith('http') ? rawHost : `http://${rawHost}`;
                const searchEndpoint = `${hostWithProtocol}:3030/search-and-extract`;

                WebSearchLogger.pipelineStart(text, searchEndpoint);
                WebSearchLogger.endpointResolved(
                    store.serverConfig.host,
                    store.serverConfig.port,
                    searchEndpoint,
                );

                appendToMessage(assistantId, '<think>Searching the web...</think>\n');

                try {
                    const maxContextLength = 10000;
                    // 60s timeout: SearXNG + 3 páginas em paralelo (max ~8s cada) + parsing
                    // Sequential era ~24s worst case; paralelo cabe em ~10-15s. 60s é margem segura.
                    const SEARCH_TIMEOUT_MS = 60000;
                    const searchController = new AbortController();
                    const searchTimeout = setTimeout(() => {
                        searchController.abort();
                    }, SEARCH_TIMEOUT_MS);

                    WebSearchLogger.requestSent(text, maxContextLength, SEARCH_TIMEOUT_MS);

                    const searchRes = await fetch(searchEndpoint, {
                        method: 'POST',
                        headers: { 'Content-Type': 'application/json' },
                        body: JSON.stringify({
                            query: text,
                            maxContextLength,
                        }),
                        signal: searchController.signal,
                    });

                    clearTimeout(searchTimeout);

                    if (searchRes.ok) {
                        const searchData = await searchRes.json();
                        searchContext = searchData.context ?? '';
                        searchSources = searchData.sources ?? [];

                        WebSearchLogger.contextReceived(
                            searchContext,
                            searchSources,
                            searchData.meta,
                        );
                    } else {
                        WebSearchLogger.requestFailed(searchRes.status, searchRes.statusText);
                    }
                } catch (error) {
                    const isTimeout = error instanceof Error && error.name === 'AbortError';
                    WebSearchLogger.serviceUnreachable(error, isTimeout);
                }

                // Clear the temporary "searching..." thought bubble
                store.updateMessageContent(assistantId, '');
            }

            // ── Step 2: Build API messages ──────────────────────────────────────
            const currentMessages = useAppStore.getState().messages;
            // Exclude the empty assistant placeholder from the history
            const historyToSend = currentMessages.slice(0, currentMessages.length - 1);

            // Warn explicitly when web search was enabled but returned no context
            if (isWebSearchEnabled && !searchContext) {
                console.warn(
                    '[WebSearch/Fallback] ⚠️ Web search was enabled but returned NO context. ' +
                    'Proceeding with LLM response without web data. ' +
                    'Check if local-search-service is running at the correct address.',
                );
            }

            const systemContent = SYSTEM_PROMPT + (searchContext ? `\n\n${searchContext}` : '');

            const apiMessages: ApiMessage[] = [
                { role: 'system', content: systemContent },
                ...historyToSend.map((m) => ({
                    role: m.role,
                    content: m.content,
                })),
            ];

            // ── Step 3: Context Window Analysis ────────────────────────────────
            if (isWebSearchEnabled) {
                const historyChars = historyToSend.reduce((sum, m) => sum + m.content.length, 0);
                WebSearchLogger.contextAnalysis(
                    systemContent.length,
                    historyChars,
                    searchContext.length,
                );
            }

            // ── Step 4: Stream LLM response ────────────────────────────────────
            WebSearchLogger.llmRequestStart(modelId, apiMessages.length);

            let chunkCount = 0;

            abortRef.current = api.chatCompletionStream(
                serverConfig,
                modelId,
                apiMessages,
                (chunk) => {
                    chunkCount++;
                    appendToMessage(assistantId, chunk);
                },
                () => {
                    WebSearchLogger.llmStreamComplete(assistantId, chunkCount);
                    WebSearchLogger.pipelineComplete(text, Date.now() - pipelineStart, isWebSearchEnabled);
                    setIsStreaming(false);
                    abortRef.current = null;
                    useAppStore.getState().saveCurrentChat();
                },
                (error) => {
                    WebSearchLogger.llmError(error);
                    const errorMsg = `\n\n⚠️ Error: ${error.message}`;
                    appendToMessage(assistantId, errorMsg);
                    setIsStreaming(false);
                    abortRef.current = null;
                    useAppStore.getState().saveCurrentChat();
                },
            );
        } catch (err) {
            const errorMsg = `⚠️ Connection failed: ${(err as Error).message}`;
            appendToMessage(assistantId, errorMsg);
            setIsStreaming(false);
            useAppStore.getState().saveCurrentChat();
        }
    }, [isStreaming, serverConfig, addMessage, appendToMessage, setIsStreaming]);

    const regenerateResponse = useCallback(async () => {
        if (isStreaming) return;

        const msgs = useAppStore.getState().messages;
        let lastUserIdx = -1;
        for (let i = msgs.length - 1; i >= 0; i--) {
            if (msgs[i].role === 'user') {
                lastUserIdx = i;
                break;
            }
        }
        if (lastUserIdx === -1) return;

        const lastUserMsg = msgs[lastUserIdx];
        const store = useAppStore.getState();
        store.deleteMessagesAfter(lastUserMsg.id);

        const remaining = useAppStore.getState().messages;
        const apiMessages: ApiMessage[] = [
            { role: 'system', content: SYSTEM_PROMPT },
            ...remaining.map((m) => ({
                role: m.role,
                content: m.content,
            })),
        ];

        const modelId = useAppStore.getState().activeModelId;
        if (!modelId) return;

        const assistantId = addMessage({ role: 'assistant', content: '' });
        setIsStreaming(true);

        try {
            abortRef.current = api.chatCompletionStream(
                serverConfig,
                modelId,
                apiMessages,
                (chunk) => appendToMessage(assistantId, chunk),
                () => {
                    setIsStreaming(false);
                    abortRef.current = null;
                    useAppStore.getState().saveCurrentChat();
                },
                (error) => {
                    appendToMessage(assistantId, `\n\n⚠️ Error: ${error.message}`);
                    setIsStreaming(false);
                    abortRef.current = null;
                    useAppStore.getState().saveCurrentChat();
                },
            );
        } catch (err) {
            appendToMessage(assistantId, `⚠️ Connection failed: ${(err as Error).message}`);
            setIsStreaming(false);
            useAppStore.getState().saveCurrentChat();
        }
    }, [isStreaming, serverConfig, addMessage, appendToMessage, setIsStreaming]);

    const stopStreaming = useCallback(() => {
        if (abortRef.current) {
            abortRef.current.abort();
            abortRef.current = null;
            setIsStreaming(false);
        }
    }, [setIsStreaming]);

    return {
        isConnected,
        isStreaming,
        availableModels,
        activeModelId,
        messages,
        sendMessage,
        stopStreaming,
        fetchModels,
        checkConnection,
        regenerateResponse,
    };
}