hyde.ts

import { getLoggedInClient, getSupabaseContext } from "./supabaseContext";
import { Result } from "./types";
import normalizePageTitle from "roamjs-components/queries/normalizePageTitle";
import { render as renderToast } from "roamjs-components/components/Toast";
import findDiscourseNode from "./findDiscourseNode";
import { nextApiRoot } from "@repo/utils/execContext";
import { DiscourseNode } from "./getDiscourseNodes";
import { getNodesByType } from "@repo/database/lib/queries";
import getAllReferencesOnPage from "./getAllReferencesOnPage";
import { getGlobalSetting } from "~/components/settings/utils/accessors";
import {
  GLOBAL_KEYS,
  SUGGESTIVE_MODE_KEYS,
} from "~/components/settings/utils/settingKeys";

type ApiEmbeddingResponse = {
  data: Array<{
    embedding: number[];
  }>;
};

type ApiSupabaseResultItem = {
  roam_uid: string;
  text_content: string;
  similarity: number;
};

export type EmbeddingVectorType = number[];

export type CandidateNodeWithEmbedding = Result & {
  type: string;
};

export type SuggestedNode = Result & {
  type: string;
};

export type RelationDetails = {
  relationLabel: string;
  relatedNodeText: string;
  relatedNodeFormat: string;
};

export type NodeSearchResult = {
  object: { uid: string; text: string };
  score: number;
};

type ResultItemMin = { uid?: string };

export type ExistingResultGroup = {
  label: string;
  results: Record<string, ResultItemMin>;
};

type HypotheticalNodeGenerator = (params: {
  node: string;
  relationType: RelationDetails;
}) => Promise<string>;

type EmbeddingFunc = (text: string) => Promise<EmbeddingVectorType>;

type SearchFunc = (params: {
  queryEmbedding: EmbeddingVectorType;
  indexData: Result[];
}) => Promise<NodeSearchResult[]>;

const API_CONFIG = {
  LLM: {
    URL: `${nextApiRoot()}/llm/openai/chat`,
    MODEL: "gpt-4.1",
    TIMEOUT_MS: 30_000,
    MAX_TOKENS: 104,
    TEMPERATURE: 0.9,
  },
  EMBEDDINGS_URL: `${nextApiRoot()}/embeddings/openai/small`,
} as const;

const handleApiError = async (
  response: Response,
  context: string,
): Promise<never> => {
  const errorText = await response.text();
  let errorData: unknown;
  try {
    errorData = JSON.parse(errorText);
  } catch (e) {
    errorData = { error: `Server responded with ${response.status}` };
  }
  console.error(
    `${context} failed with status ${response.status}. Error:`,
    errorData,
  );
  throw new Error(
    `${context} failed with status ${response.status}. Response: ${errorText}`,
  );
};

const generateHypotheticalNode: HypotheticalNodeGenerator = async ({
  node,
  relationType,
}) => {
  const { relationLabel, relatedNodeText, relatedNodeFormat } = relationType;

  const userPromptContent = `Given the source discourse node \`\`\`${node}\`\`\`,
and considering the relation \`\`\`${relationLabel}\`\`\`
which typically connects to a node of type \`\`\`${relatedNodeText}\`\`\`
(formatted like \`\`\`${relatedNodeFormat}\`\`\`),
generate a hypothetical related discourse node text that would plausibly fit this relationship.
Only return the text of the hypothetical node.`;
  const requestBody = {
    documents: [{ role: "user", content: userPromptContent }],
    passphrase: "",
    settings: {
      model: API_CONFIG.LLM.MODEL,
      maxTokens: API_CONFIG.LLM.MAX_TOKENS,
      temperature: API_CONFIG.LLM.TEMPERATURE,
    },
  };

  let response: Response | null = null;
  try {
    const signal = AbortSignal.timeout(API_CONFIG.LLM.TIMEOUT_MS);
    response = await fetch(API_CONFIG.LLM.URL, {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
      },
      body: JSON.stringify(requestBody),
      signal,
    });

    if (!response.ok) {
      await handleApiError(response, "Hypothetical node generation");
    }

    return await response.text();
  } catch (error: unknown) {
    if (
      error instanceof Error &&
      (error.name === "AbortError" || error.name === "TimeoutError")
    ) {
      console.error("Hypothetical node generation timed out", error);
      return `Error: Failed to generate hypothetical node. Request timed out.`;
    }
    console.error("Hypothetical node generation failed:", error);
    return `Error: Failed to generate hypothetical node. ${
      error instanceof Error ? error.message : String(error)
    }`;
  }
};

const createEmbedding: EmbeddingFunc = async (
  text: string,
): Promise<EmbeddingVectorType> => {
  if (!text.trim()) throw new Error("Input text for embedding is empty.");

  try {
    const response = await fetch(API_CONFIG.EMBEDDINGS_URL, {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({ input: text }),
    });

    if (!response.ok) {
      await handleApiError(response, "Embedding creation");
    }

    const data = (await response.json()) as ApiEmbeddingResponse;
    if (!data?.data?.[0]?.embedding) {
      throw new Error("Invalid API response format from embedding service.");
    }
    return data.data[0].embedding;
  } catch (error: unknown) {
    console.error("Error creating embedding:", error);
    throw error;
  }
};

const searchEmbeddings: SearchFunc = async ({
  queryEmbedding,
  indexData,
}): Promise<NodeSearchResult[]> => {
  if (!indexData?.length) {
    return [];
  }
  const supabaseClient = await getLoggedInClient();
  if (!supabaseClient) return [];

  const subsetRoamUids = indexData.map((node) => node.uid);

  const { data, error } = await supabaseClient.rpc(
    "match_embeddings_for_subset_nodes",
    {
      p_query_embedding: JSON.stringify(queryEmbedding),
      p_subset_roam_uids: subsetRoamUids,
    },
  );

  if (error) {
    console.error("Embedding search failed:", error);
    throw new Error("Embedding search failed");
  }
  const results = data;
  if (!Array.isArray(results)) {
    console.error("Embedding search response was not an array:", results);
    throw new Error("Invalid API response format: Expected an array.");
  }

  const mappedResults = results.map((item: ApiSupabaseResultItem) => ({
    object: { uid: item.roam_uid, text: item.text_content },
    score: item.similarity,
  }));
  return mappedResults;
};

const searchAgainstCandidates = async ({
  hypotheticalTexts,
  indexData,
}: {
  hypotheticalTexts: string[];
  indexData: CandidateNodeWithEmbedding[];
}): Promise<NodeSearchResult[][]> => {
  if (!hypotheticalTexts?.length || !indexData?.length) {
    return [];
  }

  const results = await Promise.all(
    hypotheticalTexts.map(async (hypoText) => {
      try {
        const queryEmbedding = await createEmbedding(hypoText);
        return await searchEmbeddings({
          queryEmbedding,
          indexData,
        });
      } catch (error: unknown) {
        let errorMessage = `Search failed for hypothetical text "${hypoText}".`;
        let errorStack;
        if (error instanceof Error) {
          errorMessage += ` Message: ${error.message}`;
          errorStack = error.stack;
        } else {
          errorMessage += ` Error: ${String(error)}`;
        }
        console.error(
          `Search exception:`,
          errorMessage,
          errorStack && `Stack: ${errorStack}`,
        );
        return [];
      }
    }),
  );
  return results;
};

const combineScores = (
  allSearchResults: NodeSearchResult[][],
): Map<string, number> => {
  const maxScores = new Map<string, number>();
  allSearchResults.forEach((resultSet) => {
    resultSet.forEach((result) => {
      const currentMax = maxScores.get(result.object.uid) ?? -Infinity;
      if (result.score > currentMax) {
        maxScores.set(result.object.uid, result.score);
      }
    });
  });
  return maxScores;
};

const rankNodes = ({
  maxScores,
  candidateNodes,
}: {
  maxScores: Map<string, number>;
  candidateNodes: CandidateNodeWithEmbedding[];
}): SuggestedNode[] => {
  if (!candidateNodes?.length) {
    return [];
  }

  const nodeMap = new Map<string, CandidateNodeWithEmbedding>(
    candidateNodes.map((node) => [node.uid, node]),
  );

  const combinedResults: { node: SuggestedNode; score: number }[] = [];
  maxScores.forEach((score, uid) => {
    const fullNode = nodeMap.get(uid);
    if (fullNode) {
      combinedResults.push({ node: fullNode, score });
    }
  });

  combinedResults.sort((a, b) => b.score - a.score);
  return combinedResults.map((item) => item.node);
};

export const findSimilarNodesUsingHyde = async ({
  candidateNodes,
  currentNodeText,
  relationDetails,
}: {
  candidateNodes: CandidateNodeWithEmbedding[];
  currentNodeText: string;
  relationDetails: RelationDetails[];
}): Promise<SuggestedNode[]> => {
  if (
    !candidateNodes?.length ||
    !currentNodeText?.trim() ||
    !relationDetails?.length
  ) {
    return [];
  }

  try {
    const hypotheticalTexts = (
      await Promise.all(
        relationDetails.map((relationType) =>
          generateHypotheticalNode({ node: currentNodeText, relationType }),
        ),
      )
    ).filter((text) => !text.startsWith("Error:"));

    if (!hypotheticalTexts.length) {
      console.warn("No valid hypothetical nodes were generated. Exiting.");
      return [];
    }

    const allSearchResults = await searchAgainstCandidates({
      hypotheticalTexts,
      indexData: candidateNodes,
    });

    const maxScores = combineScores(allSearchResults);
    return rankNodes({ maxScores, candidateNodes });
  } catch (error: unknown) {
    let errorMessage = "Failed to find similar nodes.";
    let errorStack;
    if (error instanceof Error) {
      errorMessage += ` Message: ${error.message}`;
      errorStack = error.stack;
    } else {
      errorMessage += ` Error: ${String(error)}`;
    }
    console.error(
      "Similar nodes search exception:",
      errorMessage,
      errorStack && `Stack: ${errorStack}`,
    );
    return [];
  }
};

export const getAllPageByUidAsync = async (): Promise<[string, string][]> => {
  return (await window.roamAlphaAPI.data.backend.q(
    "[:find ?pageName ?pageUid :where [?e :node/title ?pageName] [?e :block/uid ?pageUid]]",
  )) as [string, string][];
};

export const extractPagesFromChildBlock = async (
  tag: string,
): Promise<{ uid: string; text: string }[]> => {
  const results = (await window.roamAlphaAPI.data.backend.q(
    `[:find ?uid ?title
      :where [?b :node/title "${normalizePageTitle(tag)}"]
        [?a :block/refs ?b]
        [?p :block/children ?a]
        [?p :block/refs ?rf]
        [?rf :block/uid ?uid]
        [?rf :node/title ?title]]]`,
  )) as Array<[string, string]>;
  return results.map(([uid, title]) => ({ uid, text: title }));
};

export const extractPagesFromParentBlock = async (
  tag: string,
): Promise<{ uid: string; text: string }[]> => {
  const results = (await window.roamAlphaAPI.data.backend.q(
    `[:find ?uid ?title
      :where [?b :node/title "${normalizePageTitle(tag)}"]
        [?a :block/refs ?b]
        [?p :block/parents ?a]
        [?p :block/refs ?rf]
        [?rf :block/uid ?uid]
        [?rf :node/title ?title]]]`,
  )) as Array<[string, string]>;
  return results.map(([uid, title]) => ({ uid, text: title }));
};

export type PerformHydeSearchParams = {
  useAllPagesForSuggestions: boolean;
  selectedPages: string[];
  discourseNode: false | DiscourseNode;
  blockUid: string;
  validTypes: string[];
  existingResults: ExistingResultGroup[];
  uniqueRelationTypeTriplets: RelationDetails[];
  pageTitle: string;
};

export const performHydeSearch = async ({
  useAllPagesForSuggestions,
  selectedPages,
  discourseNode,
  blockUid,
  validTypes,
  existingResults,
  uniqueRelationTypeTriplets,
  pageTitle,
}: PerformHydeSearchParams): Promise<SuggestedNode[]> => {
  if (!useAllPagesForSuggestions && selectedPages.length === 0) {
    return [];
  }

  if (!discourseNode) {
    return [];
  }

  const shouldGrabFromReferencedPages =
    getGlobalSetting<boolean>([
      GLOBAL_KEYS.suggestiveMode,
      SUGGESTIVE_MODE_KEYS.includeCurrentPageRelations,
    ]) ?? true;
  const shouldGrabParentChildContext =
    getGlobalSetting<boolean>([
      GLOBAL_KEYS.suggestiveMode,
      SUGGESTIVE_MODE_KEYS.includeParentAndChildBlocks,
    ]) ?? true;

  let candidateNodesForHyde: SuggestedNode[] = [];

  const existingUids = new Set<string>(
    existingResults
      .flatMap((group) => Object.values(group.results).map((item) => item.uid))
      .filter((uid): uid is string => !!uid),
  );

  if (useAllPagesForSuggestions) {
    const context = await getSupabaseContext();
    if (!context) return [];
    const supabase = await getLoggedInClient();
    const spaceId = context.spaceId;
    if (!supabase) return [];

    candidateNodesForHyde = (
      await getNodesByType({
        supabase,
        spaceId,
        fields: { concepts: ["source_local_id", "name"], content: [] },
        ofTypes: validTypes,
        pagination: { limit: 1000 },
      })
    )
      .map((c) => {
        const node = findDiscourseNode({
          uid: c.source_local_id || "",
        });
        return {
          uid: c.source_local_id || "",
          text: c.name || "",
          type: node ? node.type : "",
        };
      })
      .filter((n) => n.uid && n.text && n.type);
  } else {
    const referenced: { uid: string; text: string }[] = [];
    if (shouldGrabFromReferencedPages) {
      referenced.push(...(await getAllReferencesOnPage(pageTitle)));
      for (const p of selectedPages) {
        referenced.push(...(await getAllReferencesOnPage(p)));
      }
    }
    if (shouldGrabParentChildContext) {
      referenced.push(...(await extractPagesFromChildBlock(pageTitle)));
      referenced.push(...(await extractPagesFromParentBlock(pageTitle)));
      for (const p of selectedPages) {
        referenced.push(...(await extractPagesFromChildBlock(p)));
        referenced.push(...(await extractPagesFromParentBlock(p)));
      }
    }
    const uniqueReferenced = Array.from(
      new Map(referenced.map((x) => [x.uid, x])).values(),
    );
    candidateNodesForHyde = uniqueReferenced
      .map((n) => {
        const node = findDiscourseNode({ uid: n.uid });
        if (
          !node ||
          node.backedBy === "default" ||
          !validTypes.includes(node.type) ||
          existingUids.has(n.uid) ||
          n.uid === blockUid
        ) {
          return null;
        }
        return {
          uid: n.uid,
          text: n.text,
          type: node.type,
        } as SuggestedNode;
      })
      .filter((n): n is SuggestedNode => n !== null);
  }

  if (candidateNodesForHyde.length && uniqueRelationTypeTriplets.length) {
    const found = await findSimilarNodesUsingHyde({
      candidateNodes: candidateNodesForHyde,
      currentNodeText: pageTitle,
      relationDetails: uniqueRelationTypeTriplets,
    });
    return found;
  }
  return [];
};

export type VectorMatch = {
  node: Result;
  score: number;
};

export const findSimilarNodesVectorOnly = async ({
  text,
  threshold = 0.4,
  limit = 15,
}: {
  text: string;
  threshold?: number;
  limit?: number;
}): Promise<VectorMatch[]> => {
  if (!text.trim()) {
    return [];
  }

  try {
    const supabase = await getLoggedInClient();
    if (!supabase) return [];

    const queryEmbedding = await createEmbedding(text);

    const { data, error } = await supabase
      .rpc("match_content_embeddings", {
        query_embedding: JSON.stringify(queryEmbedding),
        match_threshold: threshold,
      })
      .limit(limit);

    if (error) {
      console.error("Vector search failed:", error);
      throw error;
    }

    if (!data || !Array.isArray(data)) return [];

    const results: VectorMatch[] = data.map((item) => ({
      node: {
        uid: item.roam_uid,
        text: item.text_content,
      },
      score: item.similarity,
    }));

    return results;
  } catch (error) {
    console.error("Error in vector-only similar nodes search:", error);
    renderToast({
      content: `Error in vector-only similar nodes search: ${
        error instanceof Error ? error.message : String(error)
      }`,
      intent: "danger",
      id: "vector-search-error",
    });
    return [];
  }
};