src/lib/store.ts

import { create } from "zustand";
import { createJSONStorage, persist } from "zustand/middleware";
import type { ChatResponse, SearchHit, SearchResponse } from "./api";
import { runChat, runSearch, streamChat } from "./api";

export const QUERY_MAX_CHARS = 1200;
const MAX_ANSWER_CHARS = 12000;
const MAX_SAVED_QUERIES = 12;
const TOP_K_MIN = 3;
const TOP_K_MAX = 20;
const CHAT_TOP_K_MAX = 10;

type Status = "idle" | "searching" | "chatting";

type SearchState = {
  query: string;
  topK: number;
  hybrid: boolean;
  includeImages: boolean;
  results: SearchHit[];
  answer: string;
  citations: SearchHit[];
  status: Status;
  error: string | null;
  savedQueries: string[];
  selectedHit: SearchHit | null;
  streamController: AbortController | null;
  activeRequestId: number;
  setQuery: (value: string) => void;
  setTopK: (value: number) => void;
  toggleHybrid: () => void;
  toggleImages: () => void;
  setError: (value: string | null) => void;
  saveQuery: () => void;
  removeSavedQuery: (value: string) => void;
  setSelectedHit: (hit: SearchHit) => void;
  clearSelectedHit: () => void;
  cancelStream: () => void;
  clearResults: () => void;
  runSearch: () => Promise<SearchResponse | null>;
  runChat: () => Promise<ChatResponse | null>;
  runChatStream: () => Promise<void>;
};

export const useSearchStore = create<SearchState>()(
  persist(
    (set, get) => ({
      query: "",
      topK: 8,
      hybrid: true,
      includeImages: true,
      results: [],
      answer: "",
      citations: [],
      status: "idle",
      error: null,
      savedQueries: [],
      selectedHit: null,
      streamController: null,
      activeRequestId: 0,
      setQuery: (value) => {
        const trimmed = value.slice(0, QUERY_MAX_CHARS);
        set({ query: trimmed });
      },
      setTopK: (value) => {
        const next = Math.min(TOP_K_MAX, Math.max(TOP_K_MIN, value));
        set({ topK: next });
      },
      toggleHybrid: () => set((state) => ({ hybrid: !state.hybrid })),
      toggleImages: () => set((state) => ({ includeImages: !state.includeImages })),
      setError: (value) => set({ error: value }),
      saveQuery: () => {
        const query = get().query.trim();
        if (!query) return;
        set((state) => {
          const existing = state.savedQueries.filter((item) => item !== query);
          const next = [query, ...existing].slice(0, MAX_SAVED_QUERIES);
          return { savedQueries: next };
        });
      },
      removeSavedQuery: (value) =>
        set((state) => ({
          savedQueries: state.savedQueries.filter((item) => item !== value),
        })),
      setSelectedHit: (hit) => set({ selectedHit: hit }),
      clearSelectedHit: () => set({ selectedHit: null }),
      cancelStream: () => {
        const controller = get().streamController;
        if (controller) {
          controller.abort();
        }
        set({ streamController: null, status: "idle" });
      },
      clearResults: () => set({ results: [], citations: [], answer: "", error: null }),
      runSearch: async () => {
        const { query, topK, hybrid, includeImages, activeRequestId } = get();
        const trimmedQuery = query.trim().slice(0, QUERY_MAX_CHARS);
        if (!trimmedQuery) return null;
        const requestId = activeRequestId + 1;
        set({
          query: trimmedQuery,
          status: "searching",
          error: null,
          answer: "",
          citations: [],
          activeRequestId: requestId,
        });
        try {
          const response = await runSearch({
            query: trimmedQuery,
            top_k: topK,
            include_images: includeImages,
            hybrid,
          });
          if (get().activeRequestId !== requestId) return null;
          set({ results: response.results, status: "idle" });
          return response;
        } catch (err) {
          if (get().activeRequestId !== requestId) return null;
          set({
            status: "idle",
            error: err instanceof Error ? err.message : String(err),
          });
          return null;
        }
      },
      runChat: async () => {
        const { query, topK, hybrid, includeImages, activeRequestId } = get();
        const trimmedQuery = query.trim().slice(0, QUERY_MAX_CHARS);
        if (!trimmedQuery) return null;
        const requestId = activeRequestId + 1;
        set({
          query: trimmedQuery,
          status: "chatting",
          error: null,
          answer: "",
          activeRequestId: requestId,
        });
        try {
          const response = await runChat({
            query: trimmedQuery,
            top_k: Math.min(topK, CHAT_TOP_K_MAX),
            include_images: includeImages,
            hybrid,
          });
          if (get().activeRequestId !== requestId) return null;
          set({
            answer: response.answer.slice(0, MAX_ANSWER_CHARS),
            citations: response.citations,
            status: "idle",
          });
          return response;
        } catch (err) {
          if (get().activeRequestId !== requestId) return null;
          set({
            status: "idle",
            error: err instanceof Error ? err.message : String(err),
          });
          return null;
        }
      },
      runChatStream: async () => {
        const { query, topK, hybrid, includeImages, activeRequestId } = get();
        const trimmedQuery = query.trim().slice(0, QUERY_MAX_CHARS);
        if (!trimmedQuery) return;
        const existing = get().streamController;
        if (existing) {
          existing.abort();
        }
        const controller = new AbortController();
        const requestId = activeRequestId + 1;
        set({
          status: "chatting",
          error: null,
          answer: "",
          citations: [],
          streamController: controller,
          activeRequestId: requestId,
          query: trimmedQuery,
        });
        try {
          await streamChat(
            {
              query: trimmedQuery,
              top_k: Math.min(topK, CHAT_TOP_K_MAX),
              include_images: includeImages,
              hybrid,
            },
            (event) => {
              if (get().activeRequestId !== requestId) return;
              if (get().streamController !== controller) return;
              if (event.event === "delta") {
                const text = (event.data as { text?: string }).text ?? "";
                if (text) {
                  set((state) => ({
                    answer: (state.answer + text).slice(0, MAX_ANSWER_CHARS),
                  }));
                }
              }
              if (event.event === "citations") {
                const payload = event.data as { citations?: SearchHit[] };
                if (Array.isArray(payload.citations)) {
                  set({ citations: payload.citations });
                }
              }
              if (event.event === "done") {
                set({ status: "idle", streamController: null });
              }
              if (event.event === "error") {
                const payload = event.data as { message?: string };
                set({
                  status: "idle",
                  error: payload.message ?? "Streaming error.",
                  streamController: null,
                });
              }
            },
            controller.signal
          );
        } catch (err) {
          if (get().activeRequestId !== requestId) return;
          if (err instanceof Error && err.name === "AbortError") {
            set({ status: "idle", streamController: null });
            return;
          }
          set({
            status: "idle",
            error: err instanceof Error ? err.message : String(err),
            streamController: null,
          });
        }
      },
    }),
    {
      name: "layra-query-store",
      storage:
        typeof window === "undefined"
          ? undefined
          : createJSONStorage(() => localStorage),
      partialize: (state) => ({ savedQueries: state.savedQueries }),
    }
  )
);