lib/chatStore.ts

"use client";

import { create } from "zustand";
import { persist, createJSONStorage } from "zustand/middleware";
import type { AskSmartRequest, AskSmartResponse } from "./types";

export type ChatTurn = {
  id: string;
  question: string;
  response?: AskSmartResponse;
  error?: { status: number; message: string };
  pending?: boolean;
  timestamp: number;
};

export type Conversation = {
  id: string;
  title: string;
  turns: ChatTurn[];
  createdAt: number;
  updatedAt: number;
  /** When set, follow-up questions go to /ask scoped to this doc instead of /ask_smart. */
  pinnedDocId?: string | null;
  pinnedDocName?: string | null;
};

export type Settings = Required<
  Omit<AskSmartRequest, "question" | "history" | "stream">
>;

export const DEFAULT_SETTINGS: Settings = {
  top_k: 1,
  max_new_tokens: 200,
  similarity_threshold: 0.45,
  rerank_threshold: 0.6,
  anchor_threshold: 0.2,
  use_grounding: true,
  repetition_penalty: 1.15,
  no_repeat_ngram_size: 4,
  scaler: 1.0,
  bias_scaler: 1.0,
};

type ChatStore = {
  conversations: Record<string, Conversation>;
  activeId: string | null;
  settings: Settings;
  sidebarOpen: boolean;

  // sync state
  hydrated: boolean;
  syncStatus: "idle" | "saving" | "error";
  lastSavedAt: number;

  // selectors
  active: () => Conversation | null;
  list: () => Conversation[];

  // mutations
  newConversation: () => string;
  newConversationPinned: (doc_id: string, doc_name: string) => string;
  openConversation: (id: string) => void;
  deleteConversation: (id: string) => void;
  renameConversation: (id: string, title: string) => void;
  pinDoc: (id: string, doc_id: string | null, doc_name?: string | null) => void;
  appendTurn: (turn: ChatTurn) => void;
  patchTurn: (id: string, patch: Partial<ChatTurn>) => void;
  clearActive: () => void;

  setSettings: (next: Settings) => void;
  resetSettings: () => void;
  toggleSidebar: () => void;
  setSidebarOpen: (open: boolean) => void;

  // sync helpers (used by useBackendSync hook)
  _setHydrated: (v: boolean) => void;
  _setSyncStatus: (v: "idle" | "saving" | "error") => void;
  _hydrateFromRemote: (payload: {
    conversations: Record<string, Conversation>;
    activeId: string | null;
  }) => void;
};

const newId = () =>
  (typeof crypto !== "undefined" && crypto.randomUUID
    ? crypto.randomUUID()
    : Math.random().toString(36).slice(2, 12));

/**
 * Build a multi-turn history payload for the backend from a conversation.
 * Excludes pending and errored turns; caps to last `max` Q→A pairs.
 */
export function buildHistory(
  conv: Conversation | null | undefined,
  max = 5
): { question: string; answer: string }[] {
  if (!conv) return [];
  const pairs: { question: string; answer: string }[] = [];
  for (const t of conv.turns) {
    if (t.pending || t.error || !t.response) continue;
    pairs.push({ question: t.question, answer: t.response.answer });
  }
  return pairs.slice(-max);
}

const seedConversation = (): Conversation => ({
  id: newId(),
  title: "New conversation",
  turns: [],
  createdAt: Date.now(),
  updatedAt: Date.now(),
});

export const useChatStore = create<ChatStore>()(
  persist(
    (set, get) => ({
      conversations: {},
      activeId: null,
      settings: DEFAULT_SETTINGS,
      sidebarOpen: true,
      hydrated: false,
      syncStatus: "idle",
      lastSavedAt: 0,

      active: () => {
        const { activeId, conversations } = get();
        return activeId ? conversations[activeId] ?? null : null;
      },

      list: () =>
        Object.values(get().conversations).sort(
          (a, b) => b.updatedAt - a.updatedAt
        ),

      newConversation: () => {
        const c = seedConversation();
        set((s) => ({
          conversations: { ...s.conversations, [c.id]: c },
          activeId: c.id,
        }));
        return c.id;
      },

      newConversationPinned: (doc_id, doc_name) => {
        const c: Conversation = {
          ...seedConversation(),
          title: doc_name ? `Scoped: ${doc_name}` : "Scoped conversation",
          pinnedDocId: doc_id,
          pinnedDocName: doc_name ?? null,
        };
        set((s) => ({
          conversations: { ...s.conversations, [c.id]: c },
          activeId: c.id,
        }));
        return c.id;
      },

      pinDoc: (id, doc_id, doc_name) =>
        set((s) => {
          const conv = s.conversations[id];
          if (!conv) return s;
          return {
            conversations: {
              ...s.conversations,
              [id]: {
                ...conv,
                pinnedDocId: doc_id,
                pinnedDocName: doc_name ?? conv.pinnedDocName ?? null,
                updatedAt: Date.now(),
              },
            },
          };
        }),

      openConversation: (id) => set({ activeId: id }),

      deleteConversation: (id) =>
        set((s) => {
          const next = { ...s.conversations };
          delete next[id];
          const remainingIds = Object.keys(next);
          return {
            conversations: next,
            activeId:
              s.activeId === id
                ? remainingIds[0] ?? null
                : s.activeId,
          };
        }),

      renameConversation: (id, title) =>
        set((s) => ({
          conversations: {
            ...s.conversations,
            [id]: s.conversations[id]
              ? { ...s.conversations[id], title, updatedAt: Date.now() }
              : s.conversations[id],
          },
        })),

      appendTurn: (turn) =>
        set((s) => {
          let { activeId } = s;
          let conversations = s.conversations;
          if (!activeId || !conversations[activeId]) {
            const c = seedConversation();
            activeId = c.id;
            conversations = { ...conversations, [c.id]: c };
          }
          const conv = conversations[activeId];
          const isFirstTurn = conv.turns.length === 0;
          const updated: Conversation = {
            ...conv,
            title: isFirstTurn
              ? turn.question.slice(0, 60).trim() || conv.title
              : conv.title,
            turns: [...conv.turns, turn],
            updatedAt: Date.now(),
          };
          return {
            activeId,
            conversations: { ...conversations, [activeId]: updated },
          };
        }),

      patchTurn: (id, patch) =>
        set((s) => {
          if (!s.activeId) return s;
          const conv = s.conversations[s.activeId];
          if (!conv) return s;
          const turns = conv.turns.map((t) =>
            t.id === id ? { ...t, ...patch } : t
          );
          return {
            conversations: {
              ...s.conversations,
              [s.activeId]: { ...conv, turns, updatedAt: Date.now() },
            },
          };
        }),

      clearActive: () =>
        set((s) => {
          if (!s.activeId) return s;
          const conv = s.conversations[s.activeId];
          if (!conv) return s;
          return {
            conversations: {
              ...s.conversations,
              [s.activeId]: { ...conv, turns: [], updatedAt: Date.now() },
            },
          };
        }),

      setSettings: (next) => set({ settings: next }),
      resetSettings: () => set({ settings: DEFAULT_SETTINGS }),
      toggleSidebar: () => set((s) => ({ sidebarOpen: !s.sidebarOpen })),
      setSidebarOpen: (open) => set({ sidebarOpen: open }),

      // ── sync helpers ─────────────────────────────────────────────
      _setHydrated: (v) => set({ hydrated: v }),
      _setSyncStatus: (v) => set({ syncStatus: v }),
      _hydrateFromRemote: ({ conversations, activeId }) =>
        set((s) => {
          // Merge remote into local: remote wins per-id (server is the
          // source of truth on hydrate). If remote is empty but local has
          // pending in-flight turns, keep local — avoids losing a turn
          // that's mid-flight when hydration races with the request.
          const remoteIds = Object.keys(conversations);
          if (remoteIds.length === 0) {
            // Server is empty — keep local; the next save will populate it.
            return { hydrated: true };
          }
          // Preserve any locally-pending turns that aren't yet on server.
          const merged: Record<string, Conversation> = { ...conversations };
          for (const [id, local] of Object.entries(s.conversations)) {
            if (!merged[id]) {
              merged[id] = local;
              continue;
            }
            const localPending = local.turns.filter((t) => t.pending);
            if (localPending.length === 0) continue;
            const serverTurnIds = new Set(merged[id].turns.map((t) => t.id));
            const extras = localPending.filter(
              (t) => !serverTurnIds.has(t.id)
            );
            if (extras.length > 0) {
              merged[id] = {
                ...merged[id],
                turns: [...merged[id].turns, ...extras],
              };
            }
          }
          const nextActive =
            activeId && merged[activeId]
              ? activeId
              : s.activeId && merged[s.activeId]
              ? s.activeId
              : Object.keys(merged)[0] ?? null;
          return {
            conversations: merged,
            activeId: nextActive,
            hydrated: true,
          };
        }),
    }),
    {
      name: "etiya-d2l-chat",
      storage: createJSONStorage(() => localStorage),
      partialize: (state) => ({
        conversations: state.conversations,
        activeId: state.activeId,
        settings: state.settings,
        sidebarOpen: state.sidebarOpen,
      }),
    }
  )
);