mnote/wolai-frontend/src/lib/ai-agent/tools/builtins/doc/docServerTools.ts

import type { Json } from "@/types/supabase";

export type DocToolContext = {
  documentId: string;
  userId: string;
  /**
   * 来自前端的“最新文档快照”（优先使用，避免覆盖用户尚未落盘的编辑）。
   * 允许是 blocks 数组，或 { blocks } 结构。
   */
  baseBlocks?: unknown;
};

type SupabaseRouteClient = {
  from: (table: string) => {
    select: (columns: string) => SupabaseQuery;
    update: (values: Record<string, unknown>) => SupabaseUpdateQuery;
  };
};

type SupabaseQuery = {
  eq: (column: string, value: unknown) => SupabaseQuery;
  single: () => Promise<{ data: unknown; error: unknown }>;
};

type SupabaseUpdateQuery = {
  eq: (column: string, value: unknown) => SupabaseUpdateQuery;
};

export type DocSupabaseClient = SupabaseRouteClient;

type DocBlockSummary = {
  id: string;
  type: string;
  text: string;
  depth: number;
  childCount: number;
};

type DocBlockSpec = {
  type: "paragraph" | "heading";
  text: string;
  level?: number;
};

const isRecord = (v: unknown): v is Record<string, unknown> =>
  typeof v === "object" && v !== null && !Array.isArray(v);

const getValue = (obj: unknown, key: string): unknown => (isRecord(obj) ? obj[key] : undefined);

const normalizeBlocks = (content: unknown): unknown[] => {
  if (Array.isArray(content)) return content;
  const blocks = getValue(content, "blocks");
  if (Array.isArray(blocks)) return blocks;
  return [];
};

const extractInlineText = (block: unknown): string => {
  const content = getValue(block, "content");
  const nodes = Array.isArray(content) ? content : [];
  const pieces: string[] = [];
  for (const n of nodes) {
    const t = getValue(n, "text");
    if (typeof t === "string") pieces.push(t);
  }
  return pieces.join("").trim();
};

const walkSummaries = (rootBlocks: unknown[], maxNodes: number): DocBlockSummary[] => {
  const list: DocBlockSummary[] = [];
  const queue: Array<{ block: unknown; depth: number }> = rootBlocks.map((b) => ({ block: b, depth: 0 }));
  while (queue.length > 0 && list.length < maxNodes) {
    const item = queue.shift();
    if (!item) break;
    const { block, depth } = item;
    const id = String(getValue(block, "id") ?? "").trim();
    const type = String(getValue(block, "type") ?? "").trim();
    const childrenRaw = getValue(block, "children");
    const children = Array.isArray(childrenRaw) ? childrenRaw : [];
    const text = extractInlineText(block);
    if (id) {
      list.push({ id, type: type || "unknown", text, depth, childCount: children.length });
    }
    children.forEach((c) => queue.push({ block: c, depth: depth + 1 }));
  }
  return list;
};

const findContainerById = (
  blocks: unknown[],
  targetId: string,
): { container: unknown[]; index: number } | null => {
  const id = String(targetId || "").trim();
  if (!id) return null;
  for (let i = 0; i < blocks.length; i += 1) {
    const b = blocks[i];
    const bid = String(getValue(b, "id") ?? "").trim();
    if (bid === id) return { container: blocks, index: i };
    const childrenRaw = getValue(b, "children");
    const children = Array.isArray(childrenRaw) ? childrenRaw : [];
    const found = findContainerById(children, id);
    if (found) return found;
  }
  return null;
};

const createTextContent = (text: string): unknown[] => [{ type: "text", text }];

const generateId = () => {
  if (typeof crypto !== "undefined" && "randomUUID" in crypto) {
    return crypto.randomUUID();
  }
  return `bn_${Math.random().toString(36).slice(2, 10)}`;
};

const buildBlockFromSpec = (spec: DocBlockSpec): Record<string, unknown> => {
  const type = spec.type;
  const text = String(spec.text ?? "").trim();
  const base: Record<string, unknown> = {
    id: generateId(),
    type,
    props: {},
    content: createTextContent(text),
    children: [],
  };
  if (type === "heading") {
    const level = Number(spec.level ?? 2);
    base.props = { level: Math.max(1, Math.min(5, Number.isFinite(level) ? Math.floor(level) : 2)) };
  }
  return base;
};

const loadDocumentBlocks = async (supabase: DocSupabaseClient, ctx: DocToolContext) => {
  const base = normalizeBlocks(ctx.baseBlocks);
  if (base.length > 0) {
    return { blocks: base, source: "client" as const };
  }

  const { data, error } = await supabase
    .from("documents")
    .select("id,content")
    .eq("id", ctx.documentId)
    .eq("user_id", ctx.userId)
    .single();
  if (error) {
    throw new Error(isRecord(error) && typeof error.message === "string" ? error.message : "读取文档失败");
  }
  const content = isRecord(data) ? data.content : null;
  const blocks = normalizeBlocks(content);
  return { blocks, source: "db" as const };
};

const saveDocumentBlocks = async (supabase: DocSupabaseClient, ctx: DocToolContext, blocks: unknown[]) => {
  const resp = (await (supabase
    .from("documents")
    .update({ content: blocks as unknown as Json })
    .eq("id", ctx.documentId)
    .eq("user_id", ctx.userId) as unknown as Promise<{ error: unknown }>)) ?? { error: null };
  const error = resp.error;
  if (error) {
    throw new Error(isRecord(error) && typeof error.message === "string" ? error.message : "保存文档失败");
  }
};

export const createDocServerTools = (args: {
  supabase: DocSupabaseClient;
  ctx: DocToolContext;
  allowedToolIds: Set<string>;
}) => {
  const run = async (toolId: string, toolArgs: Record<string, unknown>) => {
    if (!args.allowedToolIds.has(toolId)) {
      throw new Error(`工具未被允许：${toolId}`);
    }

    if (toolId === "doc_get") {
      const maxNodesRaw = Number(toolArgs.maxBlocks ?? 80);
      const maxBlocks = Math.max(10, Math.min(240, Number.isFinite(maxNodesRaw) ? Math.floor(maxNodesRaw) : 80));
      const { blocks, source } = await loadDocumentBlocks(args.supabase, args.ctx);
      const summary = walkSummaries(blocks, maxBlocks);
      return { ok: true, source, totalTopLevelBlocks: blocks.length, blocks: summary };
    }

    if (toolId === "doc_find") {
      const query = String(toolArgs.query ?? "").trim();
      if (!query) throw new Error("缺少 query");
      const maxRaw = Number(toolArgs.maxResults ?? 8);
      const maxResults = Math.max(1, Math.min(30, Number.isFinite(maxRaw) ? Math.floor(maxRaw) : 8));
      const { blocks, source } = await loadDocumentBlocks(args.supabase, args.ctx);
      const summary = walkSummaries(blocks, 400);
      const q = query.toLowerCase();
      const hits = summary.filter((x) => x.text.toLowerCase().includes(q)).slice(0, maxResults);
      return { ok: true, source, query, results: hits };
    }

    if (toolId === "doc_insert_blocks") {
      const afterBlockId = String(toolArgs.afterBlockId ?? "").trim();
      const beforeBlockId = String(toolArgs.beforeBlockId ?? "").trim();
      const specsRaw = toolArgs.blocks;
      if (!Array.isArray(specsRaw) || specsRaw.length === 0) throw new Error("缺少 blocks");
      if (specsRaw.length > 20) throw new Error("blocks 过多（最多 20）");

      const specs: DocBlockSpec[] = specsRaw.map((x) => {
        const t = isRecord(x) ? String(x.type ?? "paragraph") : "paragraph";
        const text = isRecord(x) ? String(x.text ?? "") : "";
        const level = isRecord(x) ? Number(x.level ?? 2) : 2;
        return { type: t === "heading" ? "heading" : "paragraph", text, level };
      });

      const created = specs.map(buildBlockFromSpec);

      const { blocks, source } = await loadDocumentBlocks(args.supabase, args.ctx);
      const targetId = beforeBlockId || afterBlockId;
      const found = targetId ? findContainerById(blocks, targetId) : null;
      if (targetId && !found) {
        throw new Error(`未找到 blockId：${targetId}`);
      }

      if (!found) {
        blocks.push(...created);
      } else {
        const insertAt = beforeBlockId ? found.index : found.index + 1;
        found.container.splice(insertAt, 0, ...created);
      }

      await saveDocumentBlocks(args.supabase, args.ctx, blocks);
      return {
        ok: true,
        source,
        inserted: created.map((b) => String(b.id ?? "")),
        data: blocks,
      };
    }

    if (toolId === "doc_replace_range") {
      const blockId = String(toolArgs.blockId ?? "").trim();
      const text = String(toolArgs.text ?? "").trim();
      if (!blockId) throw new Error("缺少 blockId");
      if (!text) throw new Error("缺少 text");
      const modeRaw = String(toolArgs.mode ?? "replace").trim();
      const mode = modeRaw === "append" || modeRaw === "prepend" ? modeRaw : "replace";

      const { blocks, source } = await loadDocumentBlocks(args.supabase, args.ctx);
      const found = findContainerById(blocks, blockId);
      if (!found) throw new Error(`未找到 blockId：${blockId}`);
      const block = found.container[found.index];
      if (!isRecord(block)) throw new Error(`block 数据异常：${blockId}`);
      const prevText = extractInlineText(block);
      const nextText = mode === "append" ? `${prevText}${text}` : mode === "prepend" ? `${text}${prevText}` : text;
      found.container[found.index] = { ...block, content: createTextContent(nextText) };

      await saveDocumentBlocks(args.supabase, args.ctx, blocks);
      return { ok: true, source, blockId, mode, data: blocks };
    }

    throw new Error(`未知工具：${toolId}`);
  };

  return { run };
};