TJWaterFrontend_Refine/src/components/chat/chatStorage.ts

import { openDB, type DBSchema } from "idb";

import type {
  BranchGroup,
  ChatSessionRecord,
  ChatSessionSummary,
  ChatStorageMeta,
  LegacyPersistedChatState,
  LoadedChatState,
  Message,
} from "./GlobalChatbox.types";
import {
  cloneBranchGroups,
  cloneMessages,
  createId,
} from "./GlobalChatbox.utils";

const CHAT_DB_NAME = "tjwater-agent-chat";
const CHAT_DB_VERSION = 1;
const SESSION_STORE = "sessions";
const META_STORE = "meta";
const META_KEY = "chat-meta" as const;
const LEGACY_CHAT_STORAGE_KEY = "tjwater_agent_chat_state_v1";

type ChatDB = DBSchema & {
  sessions: {
    key: string;
    value: ChatSessionRecord;
    indexes: {
      "by-updatedAt": number;
    };
  };
  meta: {
    key: string;
    value: ChatStorageMeta;
  };
};

const emptyLoadedChatState = (): LoadedChatState => ({
  storageSessionId: undefined,
  title: undefined,
  messages: [],
  sessionId: undefined,
  branchGroups: [],
});

const sanitizeMessages = (messages: Message[] | undefined) =>
  Array.isArray(messages) ? cloneMessages(messages) : [];

const sanitizeBranchGroups = (branchGroups: BranchGroup[] | undefined) =>
  Array.isArray(branchGroups) ? cloneBranchGroups(branchGroups) : [];

const toLoadedChatState = (session: ChatSessionRecord | undefined): LoadedChatState => {
  if (!session) return emptyLoadedChatState();
  return {
    storageSessionId: session.id,
    title: session.title,
    messages: sanitizeMessages(session.messages),
    sessionId: session.sessionId,
    branchGroups: sanitizeBranchGroups(session.branchGroups),
  };
};

const toSessionSummary = (session: ChatSessionRecord): ChatSessionSummary => ({
  id: session.id,
  title: session.title,
  createdAt: session.createdAt,
  updatedAt: session.updatedAt,
});

const getDb = () =>
  openDB<ChatDB>(CHAT_DB_NAME, CHAT_DB_VERSION, {
    upgrade(db) {
      if (!db.objectStoreNames.contains(SESSION_STORE)) {
        const sessionStore = db.createObjectStore(SESSION_STORE, { keyPath: "id" });
        sessionStore.createIndex("by-updatedAt", "updatedAt");
      }

      if (!db.objectStoreNames.contains(META_STORE)) {
        db.createObjectStore(META_STORE, { keyPath: "key" });
      }
    },
  });

const readLegacyChatState = (): LegacyPersistedChatState | null => {
  if (typeof window === "undefined") return null;

  try {
    const storedRaw = window.localStorage.getItem(LEGACY_CHAT_STORAGE_KEY);
    if (!storedRaw) return null;

    const parsed = JSON.parse(storedRaw) as LegacyPersistedChatState;
    if (!Array.isArray(parsed.messages)) {
      window.localStorage.removeItem(LEGACY_CHAT_STORAGE_KEY);
      return null;
    }

    return {
      messages: sanitizeMessages(parsed.messages),
      sessionId: parsed.sessionId,
      branchGroups: sanitizeBranchGroups(parsed.branchGroups),
    };
  } catch (error) {
    console.error("[GlobalChatbox] Failed to read legacy chat state:", error);
    window.localStorage.removeItem(LEGACY_CHAT_STORAGE_KEY);
    return null;
  }
};

const clearLegacyChatState = () => {
  if (typeof window === "undefined") return;
  window.localStorage.removeItem(LEGACY_CHAT_STORAGE_KEY);
};

const getMeta = async () => {
  const db = await getDb();
  return db.get(META_STORE, META_KEY);
};

const setMeta = async (meta: Omit<ChatStorageMeta, "key">) => {
  const db = await getDb();
  await db.put(META_STORE, {
    key: META_KEY,
    ...meta,
  });
};

const getLatestSession = async () => {
  const db = await getDb();
  const sessions = await db.getAll(SESSION_STORE);
  if (sessions.length === 0) return undefined;
  return sessions.sort((left, right) => right.updatedAt - left.updatedAt)[0];
};

const migrateLegacyLocalStorage = async () => {
  const meta = await getMeta();
  if (meta?.migratedFromLocalStorage) return;

  const legacyState = readLegacyChatState();
  if (!legacyState) {
    await setMeta({
      activeSessionId: meta?.activeSessionId,
      migratedFromLocalStorage: true,
    });
    return;
  }

  const hasContent =
    legacyState.messages.length > 0 ||
    (legacyState.branchGroups?.length ?? 0) > 0 ||
    Boolean(legacyState.sessionId);

  if (!hasContent) {
    clearLegacyChatState();
    await setMeta({
      activeSessionId: undefined,
      migratedFromLocalStorage: true,
    });
    return;
  }

  const now = Date.now();
  const sessionRecord: ChatSessionRecord = {
    id: createId(),
    title: "新对话",
    createdAt: now,
    updatedAt: now,
    sessionId: legacyState.sessionId,
    messages: sanitizeMessages(legacyState.messages),
    branchGroups: sanitizeBranchGroups(legacyState.branchGroups),
  };

  const db = await getDb();
  await db.put(SESSION_STORE, sessionRecord);
  clearLegacyChatState();
  await setMeta({
    activeSessionId: sessionRecord.id,
    migratedFromLocalStorage: true,
  });
};

export const loadActiveChatState = async (): Promise<LoadedChatState> => {
  if (typeof window === "undefined") return emptyLoadedChatState();

  await migrateLegacyLocalStorage();

  const meta = await getMeta();
  const db = await getDb();

  if (meta?.activeSessionId) {
    const activeSession = await db.get(SESSION_STORE, meta.activeSessionId);
    if (activeSession) {
      return toLoadedChatState(activeSession);
    }
  }

  const latestSession = await getLatestSession();
  if (!latestSession) {
    return emptyLoadedChatState();
  }

  await setMeta({
    activeSessionId: latestSession.id,
    migratedFromLocalStorage: meta?.migratedFromLocalStorage ?? true,
  });

  return toLoadedChatState(latestSession);
};

export const saveActiveChatState = async (
  state: LoadedChatState,
): Promise<string | undefined> => {
  if (typeof window === "undefined") return state.storageSessionId;

  const hasContent =
    state.messages.length > 0 ||
    state.branchGroups.length > 0 ||
    Boolean(state.sessionId);

  const db = await getDb();
  const existingSession = state.storageSessionId
    ? await db.get(SESSION_STORE, state.storageSessionId)
    : undefined;
  const meta = await getMeta();

  if (!hasContent) {
    if (state.storageSessionId) {
      await db.delete(SESSION_STORE, state.storageSessionId);
    }
    await setMeta({
      activeSessionId: undefined,
      migratedFromLocalStorage: meta?.migratedFromLocalStorage ?? true,
    });
    return undefined;
  }

  const now = Date.now();
  const storageSessionId = state.storageSessionId ?? createId();
  const preferredTitle = state.title?.trim();
  const finalTitle = preferredTitle || existingSession?.title || "新对话";
  const nextRecord: ChatSessionRecord = {
    id: storageSessionId,
    title: finalTitle,
    createdAt: existingSession?.createdAt ?? now,
    updatedAt: now,
    sessionId: state.sessionId,
    messages: sanitizeMessages(state.messages),
    branchGroups: sanitizeBranchGroups(state.branchGroups),
  };

  await db.put(SESSION_STORE, nextRecord);
  await setMeta({
    activeSessionId: storageSessionId,
    migratedFromLocalStorage: meta?.migratedFromLocalStorage ?? true,
  });

  return storageSessionId;
};

export const listChatSessions = async (): Promise<ChatSessionSummary[]> => {
  if (typeof window === "undefined") return [];

  await migrateLegacyLocalStorage();

  const db = await getDb();
  const sessions = await db.getAll(SESSION_STORE);
  return sessions
    .sort((left, right) => right.updatedAt - left.updatedAt)
    .map(toSessionSummary);
};

export const updateChatSessionTitle = async (
  storageSessionId: string,
  title: string,
): Promise<void> => {
  if (typeof window === "undefined") return;

  const normalizedTitle = title.trim();
  if (!normalizedTitle) return;

  const db = await getDb();
  const session = await db.get(SESSION_STORE, storageSessionId);
  if (!session) return;

  await db.put(SESSION_STORE, {
    ...session,
    title: normalizedTitle,
    updatedAt: Date.now(),
  });
};

export const createEmptyChatSession = async (): Promise<LoadedChatState> => {
  if (typeof window === "undefined") return emptyLoadedChatState();

  await migrateLegacyLocalStorage();

  const now = Date.now();
  const session: ChatSessionRecord = {
    id: createId(),
    title: "新对话",
    createdAt: now,
    updatedAt: now,
    sessionId: undefined,
    messages: [],
    branchGroups: [],
  };

  const db = await getDb();
  await db.put(SESSION_STORE, session);
  const meta = await getMeta();
  await setMeta({
    activeSessionId: session.id,
    migratedFromLocalStorage: meta?.migratedFromLocalStorage ?? true,
  });

  return toLoadedChatState(session);
};

export const loadChatSessionById = async (sessionId: string): Promise<LoadedChatState> => {
  if (typeof window === "undefined") return emptyLoadedChatState();

  await migrateLegacyLocalStorage();

  const db = await getDb();
  const session = await db.get(SESSION_STORE, sessionId);
  if (!session) {
    return emptyLoadedChatState();
  }

  const meta = await getMeta();
  await setMeta({
    activeSessionId: session.id,
    migratedFromLocalStorage: meta?.migratedFromLocalStorage ?? true,
  });

  return toLoadedChatState(session);
};

export const deleteChatSession = async (sessionId: string): Promise<string | undefined> => {
  if (typeof window === "undefined") return undefined;

  const db = await getDb();
  await db.delete(SESSION_STORE, sessionId);

  const remainingSessions = await db.getAll(SESSION_STORE);
  const nextActiveSession = remainingSessions.sort(
    (left, right) => right.updatedAt - left.updatedAt,
  )[0];
  const meta = await getMeta();

  await setMeta({
    activeSessionId: nextActiveSession?.id,
    migratedFromLocalStorage: meta?.migratedFromLocalStorage ?? true,
  });

  return nextActiveSession?.id;
};