import assert from "node:assert/strict";
import fs from "node:fs/promises";
import path from "node:path";
import { fileURLToPath } from "node:url";
import vm from "node:vm";

async function loadRetrieve(stubs) {
  const __dirname = path.dirname(fileURLToPath(import.meta.url));
  const retrieverPath = path.resolve(__dirname, "../retriever.js");
  const source = await fs.readFile(retrieverPath, "utf8");
  const transformed = `${source
    .replace(/^import[\s\S]*?from\s+["'][^"']+["'];\r?\n/gm, "")
    .replace("export async function retrieve", "async function retrieve")}
this.retrieve = retrieve;
`;

  const context = vm.createContext({
    console: { log() {}, error() {}, warn() {} },
    ...stubs,
  });
  new vm.Script(transformed).runInContext(context);
  return context.retrieve;
}

function createGraph() {
  const nodes = [
    {
      id: "rule-1",
      type: "rule",
      importance: 9,
      createdTime: 1,
      archived: false,
      fields: { title: "规则一" },
      seqRange: [1, 1],
    },
    {
      id: "rule-2",
      type: "rule",
      importance: 7,
      createdTime: 2,
      archived: false,
      fields: { title: "规则二" },
      seqRange: [2, 2],
    },
    {
      id: "rule-3",
      type: "rule",
      importance: 3,
      createdTime: 3,
      archived: false,
      fields: { title: "规则三" },
      seqRange: [3, 3],
    },
  ];
  return { nodes, edges: [] };
}

function createGraphHelpers(graph) {
  return {
    getActiveNodes(target, type = null) {
      const source = target?.nodes || graph.nodes;
      return source.filter(
        (node) => !node.archived && (!type || node.type === type),
      );
    },
    getNode(target, id) {
      return (target?.nodes || graph.nodes).find((node) => node.id === id) || null;
    },
    getNodeEdges(target, nodeId) {
      return (target?.edges || graph.edges).filter(
        (edge) => edge.fromId === nodeId || edge.toId === nodeId,
      );
    },
    buildTemporalAdjacencyMap() {
      return new Map();
    },
  };
}

const schema = [{ id: "rule", label: "规则", alwaysInject: false }];

const state = {
  vectorCalls: [],
  diffusionCalls: [],
  llmCalls: [],
  llmCandidateCount: 0,
  llmResponse: { selected_ids: ["rule-2", "rule-1"] },
  llmOptions: [],
};

const graph = createGraph();
const helpers = createGraphHelpers(graph);
const retrieve = await loadRetrieve({
  ...helpers,
  MEMORY_SCOPE_BUCKETS: {
    CHARACTER_POV: "characterPov",
    USER_POV: "userPov",
    OBJECTIVE_CURRENT_REGION: "objectiveCurrentRegion",
    OBJECTIVE_ADJACENT_REGION: "objectiveAdjacentRegion",
    OBJECTIVE_GLOBAL: "objectiveGlobal",
    OTHER_POV: "otherPov",
  },
  normalizeMemoryScope(scope = {}) {
    return {
      layer: scope.layer === "pov" ? "pov" : "objective",
      ownerType: scope.ownerType || "",
      ownerId: scope.ownerId || "",
      ownerName: scope.ownerName || "",
      regionPrimary: scope.regionPrimary || "",
      regionPath: Array.isArray(scope.regionPath) ? scope.regionPath : [],
      regionSecondary: Array.isArray(scope.regionSecondary)
        ? scope.regionSecondary
        : [],
    };
  },
  getScopeRegionKey(scope = {}) {
    return String(scope.regionPrimary || "");
  },
  classifyNodeScopeBucket(node, { activeRegion = "" } = {}) {
    if (node?.scope?.layer === "pov") {
      return node?.scope?.ownerType === "user"
        ? "userPov"
        : "characterPov";
    }
    if (
      activeRegion &&
      String(node?.scope?.regionPrimary || "").trim() === String(activeRegion).trim()
    ) {
      return "objectiveCurrentRegion";
    }
    return "objectiveGlobal";
  },
  resolveScopeBucketWeight(bucket, overrides = {}) {
    return Number(overrides?.[bucket] ?? 1) || 1;
  },
  describeMemoryScope(scope = {}) {
    return `${scope.layer || "objective"}:${scope.ownerType || ""}:${scope.regionPrimary || ""}`;
  },
  describeScopeBucket(bucket = "") {
    return String(bucket || "");
  },
  buildTaskPrompt() {
    return { systemPrompt: "" };
  },
  applyTaskRegex(_settings, _taskType, _stage, text) {
    return text;
  },
  splitIntentSegments(text) {
    if (String(text).includes("和")) {
      return String(text).split("和").map((item) => item.trim());
    }
    return [];
  },
  mergeVectorResults(groups, limit) {
    const merged = new Map();
    let rawHitCount = 0;
    for (const group of groups) {
      for (const item of group) {
        rawHitCount += 1;
        const existing = merged.get(item.nodeId);
        if (!existing || item.score > existing.score) {
          merged.set(item.nodeId, item);
        }
      }
    }
    return {
      rawHitCount,
      results: [...merged.values()].slice(0, limit),
    };
  },
  collectSupplementalAnchorNodeIds() {
    return [];
  },
  isEligibleAnchorNode(node) {
    return Boolean(node?.fields?.title || node?.fields?.name);
  },
  createCooccurrenceIndex() {
    return { map: new Map(), source: "batchJournal", batchCount: 0, pairCount: 0 };
  },
  applyCooccurrenceBoost(baseScores) {
    return { scores: new Map(baseScores), boostedNodes: [] };
  },
  applyDiversitySampling(candidates, { k }) {
    return {
      applied: true,
      reason: "",
      selected: candidates.slice(0, k).reverse(),
      beforeCount: candidates.length,
      afterCount: Math.min(k, candidates.length),
    };
  },
  async runResidualRecall() {
    return { triggered: false, hits: [], skipReason: "residual-disabled-test" };
  },
  hybridScore: ({
    graphScore = 0,
    vectorScore = 0,
    lexicalScore = 0,
    importance = 0,
  }) => graphScore + vectorScore + lexicalScore + importance,
  reinforceAccessBatch() {},
  validateVectorConfig() {
    return { valid: true };
  },
  async findSimilarNodesByText(_graph, message, _embeddingConfig, topK) {
    state.vectorCalls.push({ topK, message });
    return [
      { nodeId: "rule-1", score: 0.9 },
      { nodeId: "rule-2", score: 0.8 },
      { nodeId: "rule-3", score: 0.7 },
    ];
  },
  diffuseAndRank(_adjacencyMap, seeds, options) {
    state.diffusionCalls.push({ seeds, options });
    return [
      { nodeId: "rule-2", energy: 1.2 },
      { nodeId: "rule-3", energy: 0.9 },
    ];
  },
  async callLLMForJSON(params = {}) {
    const { userPrompt = "" } = params;
    state.llmOptions.push({ ...params });
    state.llmCalls.push(userPrompt);
    state.llmCandidateCount = userPrompt
      .split("\n")
      .filter((line) => line.trim().startsWith("[")).length;
    if (params.returnFailureDetails) {
      if (state.llmResponse?.ok === false) {
        return state.llmResponse;
      }
      return {
        ok: true,
        data: state.llmResponse,
        errorType: "",
        failureReason: "",
        attempts: 1,
      };
    }
    return state.llmResponse;
  },
    getSTContextForPrompt() {
      return {};
  },
});

state.vectorCalls.length = 0;
state.diffusionCalls.length = 0;
state.llmCalls.length = 0;
const noStageResult = await retrieve({
  graph,
  userMessage: "只看当前规则",
  recentMessages: [],
  embeddingConfig: {},
  schema,
  options: {
    topK: 2,
    maxRecallNodes: 2,
    enableVectorPrefilter: false,
    enableGraphDiffusion: false,
    enableLLMRecall: false,
  },
});
assert.equal(state.vectorCalls.length, 0);
assert.equal(state.diffusionCalls.length, 0);
assert.equal(state.llmCalls.length, 0);
assert.deepEqual(Array.from(noStageResult.selectedNodeIds), ["rule-2", "rule-1"]);

state.vectorCalls.length = 0;
await retrieve({
  graph,
  userMessage: "他后来怎么做？",
  recentMessages: [
    "[assistant]: 他提到了规则二的限制",
    "[user]: 我们先看规则一",
    "[user]: 他后来怎么做？",
  ],
  embeddingConfig: {},
  schema,
  options: {
    topK: 4,
    maxRecallNodes: 2,
    enableVectorPrefilter: true,
    enableGraphDiffusion: false,
    enableLLMRecall: false,
    enableMultiIntent: false,
    enableContextQueryBlend: true,
  },
});
assert.deepEqual(
  state.vectorCalls.map((item) => item.message),
  ["他后来怎么做？", "他提到了规则二的限制", "我们先看规则一"],
);

state.vectorCalls.length = 0;
state.diffusionCalls.length = 0;
state.llmCalls.length = 0;
state.llmOptions.length = 0;
state.llmCandidateCount = 0;
state.llmResponse = { selected_ids: ["rule-2", "rule-1"] };
const llmPoolResult = await retrieve({
  graph,
  userMessage: "请根据规则给出结论",
  recentMessages: ["用户：现在该怎么做？"],
  embeddingConfig: {},
  schema,
  options: {
    topK: 4,
    maxRecallNodes: 2,
    enableVectorPrefilter: true,
    enableGraphDiffusion: false,
    enableLLMRecall: true,
    llmCandidatePool: 2,
  },
});
assert.deepEqual(state.vectorCalls, [
  { topK: 4, message: "请根据规则给出结论" },
  { topK: 4, message: "现在该怎么做？" },
]);
assert.equal(state.diffusionCalls.length, 0);
assert.equal(state.llmCandidateCount, 2);
assert.deepEqual(Array.from(llmPoolResult.selectedNodeIds), ["rule-2", "rule-1"]);
assert.equal(llmPoolResult.meta.retrieval.llm.status, "llm");
assert.equal(llmPoolResult.meta.retrieval.llm.candidatePool, 2);
assert.equal(llmPoolResult.meta.retrieval.vectorMergedHits, 3);
assert.equal(llmPoolResult.meta.retrieval.diversityApplied, true);
assert.equal(llmPoolResult.meta.retrieval.candidatePoolBeforeDpp, 3);
assert.equal(llmPoolResult.meta.retrieval.candidatePoolAfterDpp, 2);
assert.equal(state.llmOptions[0].returnFailureDetails, true);
assert.equal(state.llmOptions[0].maxRetries, 2);
assert.equal(state.llmOptions[0].maxCompletionTokens, 512);

state.vectorCalls.length = 0;
state.diffusionCalls.length = 0;
state.llmCalls.length = 0;
state.llmOptions.length = 0;
await retrieve({
  graph,
  userMessage: "规则一和规则二有什么关联",
  recentMessages: [],
  embeddingConfig: {},
  schema,
  options: {
    topK: 3,
    maxRecallNodes: 2,
    enableVectorPrefilter: true,
    enableGraphDiffusion: true,
    diffusionTopK: 7,
    enableLLMRecall: false,
    enableMultiIntent: true,
    multiIntentMaxSegments: 4,
    enableTemporalLinks: true,
    temporalLinkStrength: 0.2,
    teleportAlpha: 0.15,
  },
});
assert.equal(state.vectorCalls.length, 3);
assert.deepEqual(
  state.vectorCalls.map((item) => item.topK),
  [3, 3, 3],
);
assert.equal(state.diffusionCalls.length, 1);
assert.equal(state.diffusionCalls[0].options.topK, 7);
assert.equal(state.diffusionCalls[0].options.teleportAlpha, 0.15);
assert.equal(noStageResult.meta.retrieval.llm.status, "disabled");

state.vectorCalls.length = 0;
state.diffusionCalls.length = 0;
state.llmCalls.length = 0;
state.llmOptions.length = 0;
state.llmResponse = {
  ok: false,
  errorType: "invalid-json",
  failureReason: "输出不是有效 JSON，请严格返回紧凑 JSON 对象",
};
const fallbackResult = await retrieve({
  graph,
  userMessage: "LLM 这次会坏掉",
  recentMessages: ["用户：请回忆相关规则"],
  embeddingConfig: {},
  schema,
  options: {
    topK: 4,
    maxRecallNodes: 2,
    enableVectorPrefilter: true,
    enableGraphDiffusion: false,
    enableLLMRecall: true,
    llmCandidatePool: 2,
  },
});
assert.equal(fallbackResult.meta.retrieval.llm.status, "fallback");
assert.match(fallbackResult.meta.retrieval.llm.reason, /有效 JSON|回退到评分排序/);
assert.equal(fallbackResult.meta.retrieval.llm.fallbackType, "invalid-json");

const sceneGraph = {
  nodes: [
    {
      id: "event-1",
      type: "event",
      importance: 10,
      createdTime: 1,
      archived: false,
      fields: { title: "事件一" },
      seqRange: [1, 1],
    },
    {
      id: "character-1",
      type: "character",
      importance: 6,
      createdTime: 2,
      archived: false,
      fields: { name: "Alice" },
      seqRange: [1, 1],
    },
    {
      id: "location-1",
      type: "location",
      importance: 5,
      createdTime: 3,
      archived: false,
      fields: { title: "大厅" },
      seqRange: [1, 1],
    },
  ],
  edges: [
    { fromId: "event-1", toId: "character-1", relation: "mentions" },
    { fromId: "event-1", toId: "location-1", relation: "occurs_at" },
  ],
};
const sceneSchema = [
  { id: "event", label: "事件", alwaysInject: false },
  { id: "character", label: "角色", alwaysInject: false },
  { id: "location", label: "地点", alwaysInject: false },
];
const cappedResult = await retrieve({
  graph: sceneGraph,
  userMessage: "只看这一个场景",
  recentMessages: [],
  embeddingConfig: {},
  schema: sceneSchema,
  options: {
    topK: 3,
    maxRecallNodes: 1,
    enableVectorPrefilter: false,
    enableGraphDiffusion: false,
    enableLLMRecall: false,
    enableProbRecall: false,
  },
});
assert.equal(cappedResult.selectedNodeIds.length, 1);

const lexicalGraph = {
  nodes: [
    {
      id: "char-1",
      type: "character",
      importance: 1,
      createdTime: 1,
      archived: false,
      fields: { name: "Alice", summary: "常驻角色" },
      seqRange: [1, 1],
    },
    {
      id: "char-2",
      type: "character",
      importance: 1,
      createdTime: 1,
      archived: false,
      fields: { name: "Bob", summary: "常驻角色" },
      seqRange: [1, 1],
    },
  ],
  edges: [],
};
const lexicalSchema = [{ id: "character", label: "角色", alwaysInject: false }];
const lexicalResult = await retrieve({
  graph: lexicalGraph,
  userMessage: "Alice 现在怎么样了",
  recentMessages: [],
  embeddingConfig: {},
  schema: lexicalSchema,
  options: {
    topK: 2,
    maxRecallNodes: 1,
    enableVectorPrefilter: false,
    enableGraphDiffusion: false,
    enableLLMRecall: false,
    enableDiversitySampling: false,
    enableLexicalBoost: true,
  },
});
assert.deepEqual(Array.from(lexicalResult.selectedNodeIds), ["char-1"]);
assert.equal(lexicalResult.meta.retrieval.queryBlendActive, false);
assert.equal(lexicalResult.meta.retrieval.lexicalBoostedNodes, 1);
assert.equal(lexicalResult.meta.retrieval.lexicalTopHits[0]?.nodeId, "char-1");

const scopedGraph = {
  nodes: [
    {
      id: "obj-global",
      type: "event",
      importance: 8,
      createdTime: 1,
      archived: false,
      fields: { title: "旧王都事件" },
      seqRange: [1, 1],
      scope: { layer: "objective", regionPrimary: "旧城区" },
    },
    {
      id: "char-pov",
      type: "pov_memory",
      importance: 4,
      createdTime: 2,
      archived: false,
      fields: { summary: "艾琳觉得钟楼入口非常可疑" },
      seqRange: [2, 2],
      scope: {
        layer: "pov",
        ownerType: "character",
        ownerId: "艾琳",
        ownerName: "艾琳",
        regionPrimary: "钟楼",
      },
    },
  ],
  edges: [],
  historyState: {
    activeRegion: "钟楼",
    activeCharacterPovOwner: "艾琳",
    activeUserPovOwner: "玩家",
  },
};
const scopedSchema = [
  { id: "event", label: "事件", alwaysInject: true },
  { id: "pov_memory", label: "主观记忆", alwaysInject: false },
];
const scopedResult = await retrieve({
  graph: scopedGraph,
  userMessage: "钟楼里到底有什么",
  recentMessages: [],
  embeddingConfig: {},
  schema: scopedSchema,
  options: {
    topK: 2,
    maxRecallNodes: 1,
    enableVectorPrefilter: false,
    enableGraphDiffusion: false,
    enableLLMRecall: false,
    enableDiversitySampling: false,
    enableScopedMemory: true,
    activeRegion: "钟楼",
    activeCharacterPovOwner: "艾琳",
  },
});
assert.deepEqual(Array.from(scopedResult.selectedNodeIds), ["char-pov"]);
assert.equal(scopedResult.meta.retrieval.activeRegion, "钟楼");
assert.ok(Array.isArray(scopedResult.scopeBuckets.characterPov));
assert.equal(scopedResult.scopeBuckets.characterPov[0]?.id, "char-pov");

console.log("retrieval-config tests passed");