brk/website_next/ask/storage.js

const INDEX_KEY = "bitview.ask.chats.v2";
const CHAT_KEY_PREFIX = "bitview.ask.chat.v2.";
const LEGACY_KEY = "bitview.ask.v1";
const NEW_CHAT_TITLE = "New chat";
const CHART_UNITS = new Set([
  "addresses",
  "blocks",
  "btc",
  "number",
  "percent",
  "utxos",
  "usd",
]);
const CHART_VIEWS = new Set(["line", "area", "stacked", "bar", "dots"]);
const CHART_SCALES = new Set(["linear", "log"]);
const CHART_COLORS = new Set([
  "orange",
  "white",
  "gray",
  "sky",
  "cyan",
  "teal",
  "yellow",
  "avocado",
  "amber",
  "green",
  "emerald",
  "lime",
  "rose",
  "pink",
  "fuchsia",
  "purple",
  "violet",
  "indigo",
  "blue",
  "red",
]);

/**
 * @typedef {Object} StoredMessage
 * @property {"user" | "assistant"} role
 * @property {string} content
 * @property {number} [elapsedMs]
 * @property {StoredResponseStep[]} [steps]
 * @property {StoredArtifact[]} [artifacts]
 * @property {string} [capability]
 * @property {string[]} [metricPaths]
 * @property {ApiContext} [apiContext]
 * @property {SourceContext[]} [sourceContext]
 * @property {KnowledgeContext} [knowledgeContext]
 *
 * @typedef {Object} ApiContext
 * @property {string} key
 * @property {Record<string, string | number | boolean | (string | number | boolean)[]>} arguments
 * @property {string[]} [fields]
 *
 * @typedef {Object} SourceContext
 * @property {string} revision
 * @property {string} path
 * @property {number} startLine
 * @property {number} [endLine]
 * @property {string} content
 *
 * @typedef {Object} KnowledgeContext
 * @property {string} title
 * @property {string} description
 *
 * @typedef {Object} StoredResponseStep
 * @property {string} label
 * @property {number} elapsedMs
 *
 * @typedef {Object} ChartSeriesSpec
 * @property {string} path
 * @property {string} label
 * @property {"orange" | "white" | "gray" | "sky" | "cyan" | "teal" | "yellow" | "avocado" | "amber" | "green" | "emerald" | "lime" | "rose" | "pink" | "fuchsia" | "purple" | "violet" | "indigo" | "blue" | "red"} color
 *
 * @typedef {Object} ChartSpec
 * @property {string} title
 * @property {"addresses" | "blocks" | "btc" | "number" | "percent" | "utxos" | "usd"} unit
 * @property {"line" | "area" | "stacked" | "bar" | "dots"} view
 * @property {"linear" | "log"} scale
 * @property {ChartSeriesSpec[]} series
 *
 * @typedef {Object} ChartArtifact
 * @property {"chart"} type
 * @property {string} id
 * @property {ChartSpec} chart
 *
 * @typedef {ChartArtifact} StoredArtifact
 *
 * @typedef {Object} ChatMeta
 * @property {string} id
 * @property {string} title
 * @property {number} createdAt
 * @property {number} updatedAt
 * @property {number} messageCount
 *
 * @typedef {Object} StoredChat
 * @property {string} id
 * @property {string} title
 * @property {number} createdAt
 * @property {number} updatedAt
 * @property {number} messageCount
 * @property {StoredMessage[]} messages
 * @property {string} memory
 * @property {number} compactedCount
 *
 * @typedef {Object} ChatIndex
 * @property {string} activeChatId
 * @property {ChatMeta[]} chats
 */

/** @param {unknown} value @returns {StoredArtifact | undefined} */
function readArtifact(value) {
  if (!value || typeof value !== "object") return undefined;
  const artifact = /** @type {Record<string, any>} */ (value);
  const chart = artifact.chart;
  if (
    artifact.type !== "chart" ||
    typeof artifact.id !== "string" ||
    !chart ||
    typeof chart !== "object" ||
    typeof chart.title !== "string" ||
    typeof chart.unit !== "string" ||
    typeof chart.view !== "string" ||
    typeof chart.scale !== "string" ||
    !CHART_UNITS.has(chart.unit) ||
    !CHART_VIEWS.has(chart.view) ||
    !CHART_SCALES.has(chart.scale) ||
    !Array.isArray(chart.series)
  ) return undefined;

  const series = chart.series.filter(
    (/** @type {any} */ item) =>
      item &&
      typeof item.path === "string" &&
      typeof item.label === "string" &&
      typeof item.color === "string" &&
      CHART_COLORS.has(item.color),
  );
  if (!series.length) return undefined;

  return /** @type {StoredArtifact} */ ({
    type: "chart",
    id: artifact.id,
    chart: {
      title: chart.title,
      unit: chart.unit,
      view: chart.view,
      scale: chart.scale,
      series,
    },
  });
}

/** @param {unknown} value @returns {StoredResponseStep | undefined} */
function readResponseStep(value) {
  if (!value || typeof value !== "object") return undefined;

  const step = /** @type {Record<string, unknown>} */ (value);
  if (
    typeof step.label !== "string" ||
    typeof step.elapsedMs !== "number" ||
    !Number.isFinite(step.elapsedMs) ||
    step.elapsedMs < 0
  ) return undefined;

  return { label: step.label, elapsedMs: step.elapsedMs };
}

/** @param {unknown} value @returns {ApiContext | undefined} */
function readApiContext(value) {
  if (!value || typeof value !== "object") return undefined;
  const context = /** @type {Record<string, unknown>} */ (value);
  if (typeof context.key !== "string" || !context.key.startsWith("GET /")) return undefined;
  const rawArguments = context.arguments;
  if (!rawArguments || typeof rawArguments !== "object" || Array.isArray(rawArguments)) {
    return { key: context.key, arguments: {} };
  }
  const arguments_ = Object.fromEntries(
    Object.entries(rawArguments)
      .filter(([key, item]) =>
        key.length <= 64 &&
        (
          typeof item === "string" ||
          typeof item === "number" ||
          typeof item === "boolean" ||
          Array.isArray(item) &&
            item.every((part) =>
              typeof part === "string" ||
              typeof part === "number" ||
              typeof part === "boolean"
            )
        )
      )
      .slice(0, 16),
  );
  const fields = Array.isArray(context.fields)
    ? context.fields
      .filter((field) => typeof field === "string" && field.length <= 256)
      .slice(0, 12)
    : [];
  return /** @type {ApiContext} */ ({
    key: context.key,
    arguments: arguments_,
    ...(fields.length ? { fields } : {}),
  });
}

/** @param {unknown} value @returns {SourceContext | undefined} */
function readSourceContext(value) {
  if (!value || typeof value !== "object") return undefined;
  const context = /** @type {Record<string, unknown>} */ (value);
  if (
    typeof context.revision !== "string" ||
    typeof context.path !== "string" ||
    typeof context.startLine !== "number" ||
    !Number.isInteger(context.startLine) ||
    context.startLine < 1 ||
    typeof context.content !== "string"
  ) return undefined;
  const endLine =
    typeof context.endLine === "number" &&
    Number.isInteger(context.endLine) &&
    context.endLine >= context.startLine
      ? context.endLine
      : undefined;
  return {
    revision: context.revision.slice(0, 64),
    path: context.path.slice(0, 512),
    startLine: context.startLine,
    ...(endLine === undefined ? {} : { endLine }),
    content: context.content.slice(0, 1_000),
  };
}

/** @param {unknown} value @returns {KnowledgeContext | undefined} */
function readKnowledgeContext(value) {
  if (!value || typeof value !== "object") return undefined;
  const context = /** @type {Record<string, unknown>} */ (value);
  if (
    typeof context.title !== "string" ||
    !context.title.trim() ||
    typeof context.description !== "string" ||
    !context.description.trim()
  ) return undefined;
  return {
    title: context.title.trim().slice(0, 160),
    description: context.description.trim().slice(0, 1_500),
  };
}

/** @param {unknown} value @returns {StoredMessage | undefined} */
function readMessage(value) {
  if (!value || typeof value !== "object") return undefined;

  const message = /** @type {Record<string, unknown>} */ (value);
  if (
    (message.role !== "user" && message.role !== "assistant") ||
    typeof message.content !== "string"
  ) return undefined;

  const artifacts = Array.isArray(message.artifacts)
    ? /** @type {StoredArtifact[]} */ (
        message.artifacts.map(readArtifact).filter(Boolean)
      )
    : [];
  const elapsedMs =
    typeof message.elapsedMs === "number" &&
    Number.isFinite(message.elapsedMs) &&
    message.elapsedMs >= 0
      ? message.elapsedMs
      : undefined;
  const steps = Array.isArray(message.steps)
    ? /** @type {StoredResponseStep[]} */ (
        message.steps.map(readResponseStep).filter(Boolean)
      )
    : [];
  const rawMetricPaths = message.metricPaths;
  const capability =
    typeof message.capability === "string" &&
      /^[a-z][a-z0-9_]{0,63}$/.test(message.capability)
      ? message.capability
      : undefined;
  const hasMetricPaths = Array.isArray(rawMetricPaths);
  const metricPaths = hasMetricPaths
    ? [...new Set(/** @type {string[]} */ (rawMetricPaths.filter(
        (/** @type {unknown} */ path) => typeof path === "string" && path.trim(),
      )))].slice(0, 6)
    : [];
  const apiContext = readApiContext(message.apiContext);
  const sourceContext = Array.isArray(message.sourceContext)
    ? /** @type {SourceContext[]} */ (
        message.sourceContext.map(readSourceContext).filter(Boolean).slice(0, 3)
      )
    : [];
  const knowledgeContext = readKnowledgeContext(message.knowledgeContext);
  return {
    role: message.role,
    content: message.content,
    ...(elapsedMs !== undefined ? { elapsedMs } : {}),
    ...(steps.length ? { steps } : {}),
    ...(artifacts.length ? { artifacts } : {}),
    ...(capability ? { capability } : {}),
    ...(hasMetricPaths ? { metricPaths } : {}),
    ...(apiContext ? { apiContext } : {}),
    ...(sourceContext.length ? { sourceContext } : {}),
    ...(knowledgeContext ? { knowledgeContext } : {}),
  };
}

/** @param {string} question */
function titleFromQuestion(question) {
  const title = question.replace(/\s+/g, " ").trim();
  return title.length > 44 ? `${title.slice(0, 43)}…` : title;
}

/** @returns {ChatMeta} */
function createMeta() {
  const now = Date.now();
  return {
    id: crypto.randomUUID(),
    title: NEW_CHAT_TITLE,
    createdAt: now,
    updatedAt: now,
    messageCount: 0,
  };
}

/** @param {ChatMeta} meta */
function emptyChat(meta) {
  return { ...meta, messages: [], memory: "", compactedCount: 0 };
}

/** @param {string} id */
function chatKey(id) {
  return `${CHAT_KEY_PREFIX}${id}`;
}

/** @param {ChatIndex} index */
function writeIndex(index) {
  localStorage.setItem(INDEX_KEY, JSON.stringify(index));
}

/** @param {StoredChat} chat */
function writeChat(chat) {
  localStorage.setItem(
    chatKey(chat.id),
    JSON.stringify({
      messages: chat.messages,
      memory: chat.memory,
      compactedCount: chat.compactedCount,
    }),
  );
}

/** @returns {ChatIndex | undefined} */
function readIndex() {
  const value = localStorage.getItem(INDEX_KEY);
  if (!value) return undefined;

  const index = /** @type {ChatIndex} */ (JSON.parse(value));
  let migrated = false;
  const chats = index.chats.map((meta) => {
    if (Number.isInteger(meta.messageCount)) return meta;

    migrated = true;
    return { ...meta, messageCount: readChat(meta).messages.length };
  });

  if (!migrated) return index;

  const next = { ...index, chats };
  writeIndex(next);
  return next;
}

/** @param {ChatMeta} meta */
function readChat(meta) {
  const value = localStorage.getItem(chatKey(meta.id));
  if (!value) return emptyChat(meta);

  const stored = JSON.parse(value);
  const messages = Array.isArray(stored.messages)
    ? stored.messages.map(readMessage).filter(Boolean)
    : [];
  const compactedCount = Math.min(
    Number(stored.compactedCount) || 0,
    messages.length,
  );

  return {
    ...meta,
    messageCount: messages.length,
    messages,
    memory: typeof stored.memory === "string" ? stored.memory : "",
    compactedCount,
  };
}

/** @returns {ChatIndex} */
function initialize() {
  const meta = createMeta();
  const legacy = localStorage.getItem(LEGACY_KEY);
  const messages = /** @type {StoredMessage[]} */ (
    legacy ? JSON.parse(legacy).map(readMessage).filter(Boolean) : []
  );

  if (messages.length) {
    meta.messageCount = messages.length;
    const firstQuestion = messages.find((message) => message.role === "user");
    if (firstQuestion) meta.title = titleFromQuestion(firstQuestion.content);
  }

  const index = { activeChatId: meta.id, chats: [meta] };
  writeChat({ ...emptyChat(meta), messages });
  writeIndex(index);
  localStorage.removeItem(LEGACY_KEY);
  return index;
}

function load() {
  const index = readIndex() ?? initialize();
  const activeMeta =
    index.chats.find((chat) => chat.id === index.activeChatId) ?? index.chats[0];

  return {
    chats: index.chats,
    activeChat: readChat(activeMeta),
  };
}

function create() {
  const index = readIndex() ?? initialize();
  const meta = createMeta();
  const next = { activeChatId: meta.id, chats: [meta, ...index.chats] };

  writeChat(emptyChat(meta));
  writeIndex(next);
  return { chats: next.chats, activeChat: emptyChat(meta) };
}

/** @param {string} id */
function select(id) {
  const index = readIndex() ?? initialize();
  const meta = index.chats.find((chat) => chat.id === id);
  if (!meta) throw new Error("Chat not found");

  writeIndex({ ...index, activeChatId: id });
  return { chats: index.chats, activeChat: readChat(meta) };
}

/** @param {StoredChat} chat */
function save(chat) {
  const index = readIndex() ?? initialize();
  const now = Date.now();
  const firstQuestion = chat.messages.find((message) => message.role === "user");
  const title =
    chat.title === NEW_CHAT_TITLE && firstQuestion
      ? titleFromQuestion(firstQuestion.content)
      : chat.title;
  const saved = {
    ...chat,
    title,
    updatedAt: now,
    messageCount: chat.messages.length,
  };
  const meta = {
    id: saved.id,
    title: saved.title,
    createdAt: saved.createdAt,
    updatedAt: saved.updatedAt,
    messageCount: saved.messages.length,
  };
  const chats = [meta, ...index.chats.filter((item) => item.id !== chat.id)];

  writeChat(saved);
  writeIndex({ activeChatId: saved.id, chats });
  return { chats, activeChat: saved };
}

/**
 * @param {string} id
 * @param {{ messageCount: number, previousCompactedCount: number, previousMemory: string, memory: string, compactedCount: number }} update
 */
function saveMemory(id, update) {
  const index = readIndex() ?? initialize();
  const meta = index.chats.find((chat) => chat.id === id);
  if (!meta) return undefined;

  const chat = readChat(meta);
  if (
    chat.messages.length !== update.messageCount ||
    chat.compactedCount !== update.previousCompactedCount ||
    chat.memory !== update.previousMemory
  ) return undefined;

  const saved = {
    ...chat,
    memory: update.memory,
    compactedCount: update.compactedCount,
  };
  writeChat(saved);
  return saved;
}

/** @param {string} id */
function remove(id) {
  const index = readIndex() ?? initialize();
  const chats = index.chats.filter((chat) => chat.id !== id);

  if (!chats.length) {
    const next = createMeta();
    writeChat(emptyChat(next));
    writeIndex({ activeChatId: next.id, chats: [next] });
  } else {
    const activeChatId =
      index.activeChatId === id ? chats[0].id : index.activeChatId;
    writeIndex({ activeChatId, chats });
  }

  localStorage.removeItem(chatKey(id));
  return load();
}

export const askStorage = /** @type {const} */ ({
  load,
  create,
  select,
  save,
  saveMemory,
  remove,
});