import { describe, expect, it } from 'vitest';
import {
  buildPaperLexicalIndex,
  combinePaperRanks,
  diversifyPaperRanks,
  isPaperComparisonQuery,
  rankPaperLexical,
  rankPaperSemantic,
  tokenizeScientificText,
} from './search';
import type { PaperIndexGeneration } from './types';

function generation(texts: string[], vectors?: number[][]): PaperIndexGeneration {
  const chunks = texts.map((text, index) => ({
    id: `c${index + 1}`,
    paperId: 'p1',
    kind: 'paragraph' as const,
    text,
    pageStart: index + 1,
    pageEnd: index + 1,
    contentHash: String(index).repeat(64).slice(0, 64),
  }));
  return {
    schemaVersion: 1,
    extractorVersion: 'pdfjs-scientific-v1',
    chunkingVersion: 'scientific-400-900-v1',
    generationId: 'g1',
    paperId: 'p1',
    createdAt: 1,
    pages: [],
    blocks: [],
    chunks,
    lexical: buildPaperLexicalIndex(chunks),
    ...(vectors ? {
      embedding: {
        engine: 'test',
        model: 'test',
        modelVersion: 'v1',
        dimensions: vectors[0]?.length ?? 0,
        vectors: new Float32Array(vectors.flat()).buffer,
      },
    } : {}),
  };
}

describe('Paper Memory hybrid search', () => {
  it('preserves exact scientific tokens for lexical search', () => {
    const tokens = tokenizeScientificText('AdamW beta_2 β_2 10^-4 CLIP ResNet-50 ImageNet-1K');
    expect(tokens).toContain('adamw');
    expect(tokens).toContain('beta_2');
    expect(tokens).toContain('β_2');
    expect(tokens).toContain('10^-4');
    expect(tokens).toContain('resnet-50');
    expect(tokens).toContain('imagenet-1k');
  });

  it('ranks exact lexical evidence even without embeddings', () => {
    const index = generation([
      'We optimize with SGD and cosine decay.',
      'AdamW uses beta_2 = 0.999 with learning rate 10^-4.',
    ]);
    expect(rankPaperLexical(index, 'AdamW beta_2 10^-4')[0]?.chunkIndex).toBe(1);
  });

  it('combines semantic and lexical scores while keeping exact evidence competitive', () => {
    const index = generation([
      'General optimization discussion.',
      'The exact ResNet-50 ImageNet-1K baseline is reported here.',
    ], [[1, 0], [0.4, 0.6]]);
    const semantic = rankPaperSemantic(index, new Float32Array([1, 0]));
    const combined = combinePaperRanks([{ paperId: 'p1', lexical: rankPaperLexical(index, 'ResNet-50'), semantic }]);
    expect(combined[0]?.chunkIndex).toBe(1);
  });


  it('keeps 30-paper comparison retrieval bounded without an ANN dependency', () => {
    const ranked = Array.from({ length: 30 }, (_, paper) =>
      Array.from({ length: 20 }, (_, chunk) => ({
        paperId: `p${paper + 1}`,
        chunkIndex: chunk,
        lexicalScore: 1 / (chunk + 1),
        score: 1 - (paper * 20 + chunk) / 1000,
      })),
    ).flat().sort((a, b) => b.score - a.score);
    const output = diversifyPaperRanks(ranked, 12, 3);
    expect(output).toHaveLength(12);
    expect(Math.max(...[...new Set(output.map((item) => item.paperId))].map((paperId) => output.filter((item) => item.paperId === paperId).length))).toBeLessThanOrEqual(3);
  });

  it('keeps global leaders then caps one paper for comparison diversity', () => {
    const ranked = [
      { paperId: 'a', chunkIndex: 0, lexicalScore: 1, score: 1 },
      { paperId: 'a', chunkIndex: 1, lexicalScore: 0.9, score: 0.9 },
      { paperId: 'a', chunkIndex: 2, lexicalScore: 0.8, score: 0.8 },
      { paperId: 'b', chunkIndex: 0, lexicalScore: 0.7, score: 0.7 },
      { paperId: 'c', chunkIndex: 0, lexicalScore: 0.6, score: 0.6 },
    ];
    const diverse = diversifyPaperRanks(ranked, 5, 2);
    expect(diverse.slice(0, 2).map((item) => item.paperId)).toEqual(['a', 'a']);
    expect(diverse.map((item) => item.paperId)).toContain('b');
    expect(diverse.map((item) => item.paperId)).toContain('c');
    expect(isPaperComparisonQuery('Which papers disagree about methodology?')).toBe(true);
  });
});