import { describe, it, expect, vi, beforeEach } from 'vitest';

// --- We mock both ES client and node-fetch ----
vi.mock('@elastic/elasticsearch', () => {
  return {
    Client: vi.fn().mockImplementation(() => {
      return {
        search: vi.fn(async (opts) => {
          // Interpret the request type:
          if (opts.query?.match) {
            // --- BM25 mock output ----
            return {
              hits: {
                hits: [
                  { _id: 'bm25-1', _score: 2.0, _source: { content: 'A' } },
                  { _id: 'bm25-2', _score: 1.0, _source: { content: 'B' } },
                ]
              }
            };
          }

          if (opts.knn) {
            // --- Dense vector mock output ----
            return {
              hits: {
                hits: [
                  { _id: 'vec-1', _score: 0.9, _source: { content: 'X' } },
                  { _id: 'vec-2', _score: 0.8, _source: { content: 'Y' } },
                ]
              }
            };
          }

          return { hits: { hits: [] } };
        })
      };
    })
  };
});

vi.mock('node-fetch', () => {
  return {
    default: vi.fn(async () => {
      return {
        ok: true,
        json: async () => ({
          embedding: [0.1, 0.2, 0.3]   // Dummy vector
        })
      };
    })
  };
});

// After mocks, import retrieval module:
import {
  embed,
  bm25Search,
  vectorSearch,
  hybridSearch
} from '../src/retrieval/retrieval.mjs';

describe('retrieval.mjs', () => {
  beforeEach(() => {
    vi.clearAllMocks();
  });

  // ------------------------------
  // EMBEDDING
  // ------------------------------
  it('embed() returns an embedding array', async () => {
    const emb = await embed('test');
    expect(Array.isArray(emb)).toBe(true);
    expect(emb).toEqual([0.1, 0.2, 0.3]);
  });

  // ------------------------------
  // BM25
  // ------------------------------
  it('bm25Search() returns valid results', async () => {
    const res = await bm25Search('love', 2);
    expect(res.length).toBe(2);

    expect(res[0]).toMatchObject({
      id: 'bm25-1',
      score: 2.0,
      content: 'A'
    });
  });

  // ------------------------------
  // Dense Vector
  // ------------------------------
  it('vectorSearch() returns vector-based results', async () => {
    const res = await vectorSearch('unity', 2);
    expect(res.length).toBe(2);

    expect(res[0]).toMatchObject({
      id: 'vec-1',
      score: 0.9,
      content: 'X'
    });
  });

  // ------------------------------
  // Hybrid (RRF)
  // ------------------------------
  it('hybridSearch() fuses BM25 and vector results via RRF', async () => {
    const res = await hybridSearch('seeking', 3);

    // Should merge vec + bm25 in some order.
    // Score fusion means 'vec-1' or 'bm25-1' will come first, but we just check the structure.
    expect(res.length).toBeGreaterThan(0);

    const ids = res.map(r => r.id);

    // Confirm that hybrid includes items from BOTH sources
    expect(ids).toContain('vec-1');
    expect(ids).toContain('bm25-1');
  });
});