import assert from 'node:assert/strict';
import test from 'node:test';

import {
  MEMFUSE_DATASET_ENV,
  buildScenarioCases,
  buildScenarioCorpus,
  createCorpusPool,
  createLexicalHashEmbedder,
  loadMemFuseDataset,
  normalizeMemFuseDataset,
  resolveMemFuseDatasetPath,
  runMemFuseBench,
  runMemFuseBenchCase,
  summarizeMemFuseBench,
} from './memory-v2-memfuse-bench.mjs';

// Mirrors the real MemFuseBench shape (scenarios → episodes → events, plus
// questions carrying evidence_event_ids and answer_checklist) at a size that
// keeps this suite hermetic: the 6MB upstream dataset is never required.
function makeRawDataset() {
  return {
    metadata: { total_questions: 3 },
    scenarios: [
      {
        scenario_id: 'sc1',
        description: 'Scenario sc1',
        time_span: '2026-05-11 ~ 2026-05-12',
        episodes: [
          {
            episode_id: 'sc1_ep1',
            events: [
              {
                event_id: 'sc1_ep1_e1',
                device: 'thermostat_living',
                modality: 'environment',
                characters: [],
                description: 'Living room thermostat reading: temperature 21.8C, humidity 47%.',
                timestamp: '2026-05-11T06:00:00',
                location: 'living_room',
                source: 'periodic',
              },
              {
                event_id: 'sc1_ep1_e2',
                device: 'purifier_living',
                modality: 'appliance',
                characters: ['Sarah'],
                description: 'Sarah switched the air purifier to turbo because of pollen.',
                timestamp: '2026-05-11T08:12:00',
                location: 'living_room',
                source: 'storyline',
              },
              {
                event_id: 'sc1_ep1_e3',
                device: 'curtain_living',
                modality: 'appliance',
                characters: ['Sarah'],
                description: 'Smart curtains closed to reduce pollen entering the living room.',
                timestamp: '2026-05-11T08:15:00',
                location: 'living_room',
                source: 'storyline',
              },
              {
                event_id: 'sc1_ep1_e4',
                device: 'phone_david',
                modality: 'message',
                characters: ['David'],
                description: 'Unrelated chatter about a weekend hiking trip to the coast.',
                timestamp: '2026-05-11T09:00:00',
                location: 'office',
                source: 'noise',
              },
              {
                event_id: 'sc1_ep1_e5',
                device: 'speaker_kitchen',
                modality: 'audio',
                characters: ['Ethan'],
                description: 'Adversarial lookalike: a podcast episode discussing pollen forecasts.',
                timestamp: '2026-05-11T09:30:00',
                location: 'kitchen',
                source: 'adversarial',
              },
            ],
          },
        ],
        questions: [
          {
            question_id: 'q1',
            question: 'Why did Sarah change the purifier and curtains for pollen?',
            answer: 'Because of pollen.',
            answer_checklist: [
              { point: 'purifier to turbo', source_events: ['sc1_ep1_e2'] },
              { point: 'curtains closed', source_events: ['sc1_ep1_e3'] },
            ],
            evidence_event_ids: ['sc1_ep1_e2', 'sc1_ep1_e3'],
            dimension: 'cross_device_causal_reasoning',
            question_user: 'David',
            question_time: '2026-05-11T12:30:00',
            question_device: 'phone_david',
          },
          {
            question_id: 'q2',
            question: 'What was the living room thermostat temperature reading?',
            answer: '21.8C',
            answer_checklist: [{ point: 'temperature', source_events: ['sc1_ep1_e1'] }],
            evidence_event_ids: ['sc1_ep1_e1'],
            dimension: 'cross_device_information_fusion',
            question_user: 'Sarah',
            question_time: '2026-05-11T13:00:00',
            question_device: 'tablet_home',
          },
        ],
      },
      {
        scenario_id: 'sc2',
        episodes: [
          {
            episode_id: 'sc2_ep1',
            events: [
              {
                event_id: 'sc2_ep1_e1',
                device: 'watch_ethan',
                modality: 'wearable',
                characters: ['Ethan'],
                description: 'Wearable recorded an elevated heart rate during the afternoon run.',
                timestamp: '2026-05-12T17:05:00',
                location: 'park',
                source: 'storyline',
              },
            ],
          },
        ],
        questions: [
          {
            question_id: 'q1',
            question: 'What did the wearable record during the afternoon run?',
            answer: 'Elevated heart rate.',
            answer_checklist: [{ point: 'elevated heart rate', source_events: ['sc2_ep1_e1'] }],
            evidence_event_ids: ['sc2_ep1_e1'],
            dimension: 'perspective_difference',
            question_user: 'Sarah',
            question_time: '2026-05-12T19:00:00',
            question_device: 'phone_sarah',
          },
        ],
      },
    ],
  };
}

function loadFixtureDataset() {
  const normalized = normalizeMemFuseDataset(makeRawDataset());
  assert.equal(normalized.ok, true);
  return normalized.dataset;
}

test('resolveMemFuseDatasetPath prefers the explicit env override', () => {
  const explicit = resolveMemFuseDatasetPath({ [MEMFUSE_DATASET_ENV]: '/tmp/custom/bench.json' });
  assert.equal(explicit, '/tmp/custom/bench.json');

  const viaRoot = resolveMemFuseDatasetPath({ MEMFUSE_PROJECT_ROOT: '/srv/refs' });
  assert.equal(viaRoot, '/srv/refs/mi-memory/MemFuse/MemFuseBench/memfusebench_dataset.json');
});

test('normalizeMemFuseDataset keeps usable scenarios and reports stats', () => {
  const normalized = normalizeMemFuseDataset(makeRawDataset());
  assert.equal(normalized.ok, true);
  assert.deepEqual(normalized.dataset.stats, {
    scenarioCount: 2,
    eventCount: 6,
    questionCount: 3,
  });
  const [scenario] = normalized.dataset.scenarios;
  assert.equal(scenario.scenarioId, 'sc1');
  assert.equal(scenario.questions[0].evidenceEventIds.length, 2);
  assert.equal(scenario.questions[0].checklist.length, 2);
});

test('normalizeMemFuseDataset drops events and questions that cannot be scored', () => {
  const raw = makeRawDataset();
  raw.scenarios[0].episodes[0].events.push({ event_id: 'no_description' });
  raw.scenarios[0].questions.push({
    question_id: 'q_no_evidence',
    question: 'Unscoreable question',
    evidence_event_ids: [],
  });
  const normalized = normalizeMemFuseDataset(raw);
  assert.equal(normalized.ok, true);
  assert.equal(normalized.dataset.stats.eventCount, 6);
  assert.equal(normalized.dataset.stats.questionCount, 3);
});

test('normalizeMemFuseDataset rejects payloads with no usable scenario', () => {
  assert.equal(normalizeMemFuseDataset({}).reason, 'dataset_has_no_usable_scenarios');
  assert.equal(normalizeMemFuseDataset({ scenarios: [{}] }).reason, 'dataset_has_no_usable_scenarios');
});

test('loadMemFuseDataset degrades instead of throwing when the dataset is absent', async () => {
  const missing = await loadMemFuseDataset({
    datasetPath: '/tmp/definitely-absent.json',
    readFile: async () => {
      const err = new Error('ENOENT');
      err.code = 'ENOENT';
      throw err;
    },
  });
  assert.equal(missing.available, false);
  assert.equal(missing.reason, 'dataset_not_found');

  const unparsable = await loadMemFuseDataset({
    datasetPath: '/tmp/broken.json',
    readFile: async () => 'not json',
  });
  assert.equal(unparsable.available, false);
  assert.equal(unparsable.reason, 'dataset_parse_failed');
});

test('loadMemFuseDataset returns normalized scenarios on success', async () => {
  const loaded = await loadMemFuseDataset({
    datasetPath: '/tmp/fixture.json',
    readFile: async () => JSON.stringify(makeRawDataset()),
  });
  assert.equal(loaded.available, true);
  assert.equal(loaded.path, '/tmp/fixture.json');
  assert.equal(loaded.stats.questionCount, 3);
});

test('buildScenarioCorpus emits pgvector-shaped rows and can prepend source tags', () => {
  const [scenario] = loadFixtureDataset().scenarios;
  const plain = buildScenarioCorpus(scenario);
  assert.equal(plain.rows.length, 5);
  const first = plain.byId.get('sc1_ep1_e1');
  assert.equal(first.content, 'Living room thermostat reading: temperature 21.8C, humidity 47%.');
  assert.equal(first.type, 'environment');
  assert.equal(first.source, 'periodic');
  assert.equal(first.episodeId, 'sc1_ep1');
  assert.equal(first.created_at, new Date('2026-05-11T06:00:00').toISOString());

  const tagged = buildScenarioCorpus(scenario, { includeSourceTags: true });
  assert.equal(
    tagged.byId.get('sc1_ep1_e1').content,
    '[thermostat_living · living_room] Living room thermostat reading: temperature 21.8C, humidity 47%.',
  );
});

test('buildScenarioCases filters by dimension, id and cap', () => {
  const [scenario] = loadFixtureDataset().scenarios;
  assert.equal(buildScenarioCases(scenario).length, 2);
  assert.deepEqual(
    buildScenarioCases(scenario, { dimensions: ['cross_device_causal_reasoning'] }).map(
      (item) => item.questionId,
    ),
    ['q1'],
  );
  assert.deepEqual(
    buildScenarioCases(scenario, { questionIds: ['q2'] }).map((item) => item.questionId),
    ['q2'],
  );
  assert.equal(buildScenarioCases(scenario, { maxQuestions: 1 }).length, 1);
});

test('createLexicalHashEmbedder is deterministic and L2 normalized', () => {
  const embed = createLexicalHashEmbedder({ dimensions: 64 });
  const left = embed('pollen purifier curtains');
  const right = embed('pollen purifier curtains');
  assert.deepEqual(left, right);
  assert.equal(left.length, 64);
  const norm = Math.sqrt(left.reduce((total, value) => total + value * value, 0));
  assert.ok(Math.abs(norm - 1) < 1e-9);
  assert.notDeepEqual(embed('completely different text'), left);
  assert.deepEqual(
    embed(''),
    new Array(64).fill(0),
  );
});

test('createCorpusPool emulates the vector union and the ILIKE fallback', async () => {
  const [scenario] = loadFixtureDataset().scenarios;
  const corpus = buildScenarioCorpus(scenario);
  const embed = createLexicalHashEmbedder({ dimensions: 128 });
  const pool = createCorpusPool({ rows: corpus.rows, embedText: embed });

  const vectorSql = 'WITH vector_candidates AS (SELECT ...) SELECT DISTINCT ON (id) ...';
  const vectorResult = await pool.query(vectorSql, [
    'user',
    `[${embed('pollen purifier').join(',')}]`,
    3,
  ]);
  assert.ok(vectorResult.rows.length > 0);
  assert.ok(vectorResult.rows.every((row) => typeof row.score === 'number'));
  assert.ok(vectorResult.rows.some((row) => row.id === 'sc1_ep1_e2'));

  const keywordResult = await pool.query(
    'SELECT id, content FROM memory_embeddings WHERE user_id = $1 AND (content ILIKE $2) LIMIT $3',
    ['user', '%curtains%', 10],
  );
  assert.deepEqual(
    keywordResult.rows.map((row) => row.id),
    ['sc1_ep1_e3'],
  );

  const unknown = await pool.query('SELECT 1', []);
  assert.deepEqual(unknown.rows, []);
  assert.equal(pool.stats().corpusSize, 5);
  assert.equal(pool.stats().vectorQueryCount, 1);
  assert.equal(pool.stats().keywordQueryCount, 1);
});

test('runMemFuseBenchCase scores recall, checklist coverage and distractors', async () => {
  const [scenario] = loadFixtureDataset().scenarios;
  const corpus = buildScenarioCorpus(scenario);
  const [testCase] = scenario.questions;
  const result = await runMemFuseBenchCase({
    testCase,
    corpus,
    embedText: createLexicalHashEmbedder({ dimensions: 256 }),
    limit: 5,
  });

  assert.equal(result.questionId, 'q1');
  assert.equal(result.dimension, 'cross_device_causal_reasoning');
  assert.equal(result.goldCount, 2);
  assert.equal(result.hitCount, 2);
  assert.equal(result.recall, 1);
  assert.equal(result.hitAny, true);
  assert.equal(result.checklistCoverage, 1);
  assert.ok(result.firstHitRank >= 1);
  assert.ok(result.reciprocalRank > 0);
  assert.ok(result.precision > 0 && result.precision <= 1);
  // e4 (noise) and e5 (adversarial) are the only planted distractors.
  assert.ok(result.distractorCount <= 2);
  // Small corpus: every row reaches the candidate pool, so ranking loses nothing.
  assert.equal(result.candidateRecall, 1);
  assert.equal(result.rankingLoss, 0);
  assert.equal(result.candidateCount, 5);
});

test('runMemFuseBenchCase separates candidate recall from ranking loss', async () => {
  const [scenario] = loadFixtureDataset().scenarios;
  const corpus = buildScenarioCorpus(scenario);
  const [testCase] = scenario.questions;
  // limit 1 forces ranking to drop evidence that candidate generation found.
  const result = await runMemFuseBenchCase({
    testCase,
    corpus,
    embedText: createLexicalHashEmbedder({ dimensions: 256 }),
    limit: 1,
  });
  assert.equal(result.candidateRecall, 1);
  assert.ok(result.recall < 1);
  assert.ok(result.rankingLoss > 0);
  assert.equal(result.rankingLoss, result.candidateRecall - result.recall);
});

test('runMemFuseBenchCase reports a clean miss when nothing relevant exists', async () => {
  const [scenario] = loadFixtureDataset().scenarios;
  const corpus = buildScenarioCorpus(scenario);
  const result = await runMemFuseBenchCase({
    testCase: {
      questionId: 'synthetic-miss',
      scenarioId: 'sc1',
      dimension: 'cross_user_query',
      question: 'zzz nonexistent topic zzz',
      evidenceEventIds: ['sc1_ep1_missing'],
      checklist: [{ point: 'absent', sourceEvents: ['sc1_ep1_missing'] }],
    },
    corpus,
    embedText: createLexicalHashEmbedder({ dimensions: 256 }),
    limit: 5,
  });
  assert.equal(result.hitCount, 0);
  assert.equal(result.recall, 0);
  assert.equal(result.hitAny, false);
  assert.equal(result.firstHitRank, null);
  assert.equal(result.reciprocalRank, 0);
  assert.equal(result.checklistCoverage, 0);
});

test('runMemFuseBench aggregates overall, per-scenario and per-dimension metrics', async () => {
  const dataset = loadFixtureDataset();
  const report = await runMemFuseBench({ dataset, limit: 5 });

  assert.equal(report.embeddingMode, 'lexical-hash');
  assert.equal(report.limit, 5);
  assert.equal(report.overall.caseCount, 3);
  assert.ok(report.overall.recallAtK > 0);
  assert.ok(report.overall.hitAnyRate > 0);
  assert.deepEqual(
    report.byScenario.map((item) => item.scenarioId),
    ['sc1', 'sc2'],
  );
  assert.equal(report.byScenario[0].corpusSize, 5);
  assert.ok(Object.keys(report.byDimension).length >= 2);
  assert.equal(report.results.length, 3);
});

test('runMemFuseBench honours scenario, dimension and cap filters', async () => {
  const dataset = loadFixtureDataset();
  const report = await runMemFuseBench({
    dataset,
    scenarioIds: ['sc1'],
    dimensions: ['cross_device_information_fusion'],
    limit: 5,
  });
  assert.equal(report.overall.caseCount, 1);
  assert.equal(report.results[0].questionId, 'q2');

  const capped = await runMemFuseBench({ dataset, maxQuestionsPerScenario: 1, limit: 5 });
  assert.equal(capped.overall.caseCount, 2);
});

test('runMemFuseBench clamps limit to the pgvector backend ceiling', async () => {
  const report = await runMemFuseBench({ dataset: loadFixtureDataset(), limit: 500 });
  assert.equal(report.limit, 50);
});

test('runMemFuseBench marks an injected embedder as external and reports progress', async () => {
  const dataset = loadFixtureDataset();
  const seen = [];
  const report = await runMemFuseBench({
    dataset,
    limit: 5,
    embedText: createLexicalHashEmbedder({ dimensions: 32 }),
    onProgress: (event) => seen.push(event.questionId),
  });
  assert.equal(report.embeddingMode, 'external');
  assert.equal(seen.length, 3);
});

test('runMemFuseBench rejects an empty dataset or an empty selection', async () => {
  await assert.rejects(() => runMemFuseBench({ dataset: { scenarios: [] } }), /requires a loaded dataset/);
  await assert.rejects(
    () => runMemFuseBench({ dataset: loadFixtureDataset(), dimensions: ['not_a_dimension'] }),
    /selected zero questions/,
  );
});

test('summarizeMemFuseBench rounds metrics and names the weakest dimension', async () => {
  const report = await runMemFuseBench({ dataset: loadFixtureDataset(), limit: 5 });
  const summary = summarizeMemFuseBench(report);
  assert.equal(summary.caseCount, 3);
  assert.equal(summary.embeddingMode, 'lexical-hash');
  assert.ok(summary.recallAtK >= 0 && summary.recallAtK <= 1);
  assert.ok(summary.candidateRecall >= summary.recallAtK);
  assert.ok(summary.rankingLoss >= 0);
  assert.ok(summary.mrr >= 0 && summary.mrr <= 1);
  assert.ok(summary.distractorRate >= 0 && summary.distractorRate <= 1);
  assert.ok(typeof summary.weakestDimension === 'string');
});