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,
  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, importance = 0 }) =>
    graphScore + vectorScore + 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;
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: "请根据规则给出结论" }]);
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);

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);

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