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

describe('runPipelineStep retrieval modes', () => {
  beforeEach(() => {
    vi.resetModules();
  });

  afterEach(() => {
    vi.restoreAllMocks();
    vi.unmock('../src/retrieval/retrieval.mjs');
    vi.unmock('../src/providers/provider.mjs');
    vi.unmock('../src/generator/generator_core.mjs');
    vi.unmock('../src/verifier/verifier_core.mjs');
    vi.unmock('../src/reward/reward_core.mjs');
  });

  async function setupRetrievalMocks() {
    const hybridSearch = vi.fn(async () => [{ content: 'hybrid' }]);
    const bm25Search = vi.fn(async () => [{ content: 'bm25' }]);
    const vectorSearch = vi.fn(async () => [{ content: 'vector' }]);
    const hydeHybrid = vi.fn(async () => [{ content: 'hyde' }]);

    vi.doMock('../src/retrieval/retrieval.mjs', () => ({
      hybridSearch,
      bm25Search,
      vectorSearch,
      hydeHybrid,
    }));

    vi.doMock('../src/providers/provider.mjs', () => ({
      loadProviderFor: () => ({ generate: async () => 'ok' }),
    }));

    vi.doMock('../src/generator/generator_core.mjs', () => ({
      runGenerator: vi.fn(async () => ({ answer: 'a', raw: 'a' })),
    }));
    vi.doMock('../src/verifier/verifier_core.mjs', () => ({
      runVerifier: vi.fn(async () => ({ ok: true, raw: 'yes' })),
    }));
    vi.doMock('../src/reward/reward_core.mjs', () => ({
      runReward: vi.fn(async () => ({ ok: true, score: 1, raw: '1' })),
    }));

    const mod = await import('../src/pipeline/step.mjs');
    return { hybridSearch, bm25Search, vectorSearch, hydeHybrid, runPipelineStep: mod.runPipelineStep };
  }

  it('uses bm25 when retrievalMode=bm25', async () => {
    const mocks = await setupRetrievalMocks();

    await mocks.runPipelineStep({ question: 'q', retrievalMode: 'bm25' });

    expect(mocks.bm25Search).toHaveBeenCalledTimes(1);
    expect(mocks.vectorSearch).not.toHaveBeenCalled();
    expect(mocks.hybridSearch).not.toHaveBeenCalled();
    expect(mocks.hydeHybrid).not.toHaveBeenCalled();
  });

  it('uses hyde when retrievalMode=hyde', async () => {
    const mocks = await setupRetrievalMocks();

    await mocks.runPipelineStep({ question: 'q', retrievalMode: 'hyde' });

    expect(mocks.hydeHybrid).toHaveBeenCalledTimes(1);
    expect(mocks.bm25Search).not.toHaveBeenCalled();
    expect(mocks.vectorSearch).not.toHaveBeenCalled();
    expect(mocks.hybridSearch).not.toHaveBeenCalled();
  });
});

describe('runPipelineBatch question cap', () => {
  beforeEach(() => {
    vi.resetModules();
    process.env.PIPELINE_CHUNK_SOURCE = 'jsonl';
    process.env.PIPELINE_CACHE_DIR = path.join(
      os.tmpdir(),
      `cache-${Date.now()}`,
    );
  });

  afterEach(() => {
    delete process.env.PIPELINE_CHUNK_SOURCE;
    delete process.env.PIPELINE_CACHE_DIR;
    vi.restoreAllMocks();
    vi.unmock('../src/providers/provider.mjs');
    vi.unmock('../src/retrieval/jsonl_chunks.mjs');
    vi.unmock('../src/pipeline/step.mjs');
  });

  it('stops once the question limit is reached in question-first mode', async () => {
    const questionsPerChunk = 'Q1?\nQ2?\nQ3?';

    // Mock question provider + retrieval
    vi.doMock('../src/providers/provider.mjs', () => ({
      loadProviderFor: (stage) => {
        if (stage === 'question') return { generate: async () => questionsPerChunk };
        if (stage === 'verifier')
          return {
            generate: async () =>
              JSON.stringify({ REASONING: ['ok'], SCORE: 1 }),
          };
        if (stage === 'reward') return { generate: async () => '0.9 good' };
        return { generate: async () => '' };
      },
    }));

    // Force question generator to return a fixed list to avoid parser variability
    const runQuestionGenerator = vi.fn(async () => ({
      questions: ['Q1?', 'Q2?', 'Q3?'],
      raw: questionsPerChunk,
      parsed: null,
    }));
    vi.doMock('../src/question/question_core.mjs', () => ({ runQuestionGenerator }));

    vi.doMock('../src/retrieval/jsonl_chunks.mjs', () => ({
      loadRagChunks: vi.fn(async (limit) =>
        Array.from({ length: limit ?? 3 }, (_, i) => ({
          id: `c-${i}`,
          content: `chunk ${i}`,
        })),
      ),
    }));

    // Stub runPipelineStep to avoid model calls; count invocations
    const runPipelineStep = vi.fn(async () => ({
      status: 'accepted',
      context: [],
      gen: { answer: 'a' },
      ver: { ok: true },
      rew: { ok: true },
    }));
    vi.doMock('../src/pipeline/step.mjs', () => ({ runPipelineStep }));

    const outPath = path.join(os.tmpdir(), `test-gold-${Date.now()}.jsonl`);
    const { runPipelineBatch } = await import('../src/pipeline/batch.mjs');
    const result = await runPipelineBatch({
      limit: 2, // question cap
      chunkLimit: 3,
      seedMode: 'question-first',
      outPath,
      verbose: false,
      logger: { log() {}, error() {} },
    });

    expect(result.processed).toBe(2);
    expect(runPipelineStep).toHaveBeenCalledTimes(2);
    expect(result.accepted).toBe(2);
  });
});

describe('fetchChunksFromIndex default size', () => {
  beforeEach(() => {
    vi.resetModules();
    vi.unmock('../src/retrieval/retrieval.mjs');
  });

  afterEach(() => {
    vi.restoreAllMocks();
  });

  it('uses a larger default than 10 when no limit is provided', async () => {
    let capturedOpts = null;

    vi.unmock('../src/retrieval/retrieval.mjs');

    vi.doMock('@elastic/elasticsearch', () => ({
      Client: vi.fn().mockImplementation(() => ({
        search: vi.fn(async (opts) => {
          capturedOpts = opts;
          return { hits: { hits: [] } };
        }),
      })),
    }));

    // avoid accidental network fetch calls in this test
    vi.doMock('node-fetch', () => ({
      default: vi.fn(async () => ({ ok: true, json: async () => ({ embedding: [] }) })),
    }));

    // Re-export the real module so the fetch function exists, but with mocked deps above.
    vi.doMock('../src/retrieval/retrieval.mjs', async (importOriginal) => {
      const actual = await importOriginal();
      return { ...actual };
    });

    const { fetchChunksFromIndex } = await import('../src/retrieval/retrieval.mjs');
    await fetchChunksFromIndex();

    expect(capturedOpts?.size).toBe(100);
  });
});

describe('question-first default chunk source', () => {
  beforeEach(() => {
    vi.resetModules();
  });

  afterEach(() => {
    vi.restoreAllMocks();
    vi.unmock('../src/retrieval/jsonl_chunks.mjs');
    vi.unmock('../src/pipeline/step.mjs');
  });

  it('prefers jsonl chunks by default when no env override is set', async () => {
    const loadRagChunks = vi.fn(async () => [{ id: 'c1', content: 'chunk' }]);

    vi.doMock('../src/retrieval/jsonl_chunks.mjs', () => ({
      loadRagChunks,
    }));

    // Stub pipeline step to avoid model calls
    vi.doMock('../src/pipeline/step.mjs', () => ({
      runPipelineStep: vi.fn(async () => ({
        status: 'accepted',
        context: [],
        gen: { answer: 'a' },
        ver: { ok: true },
        rew: { ok: true },
      })),
    }));

    const { runPipelineBatch } = await import('../src/pipeline/batch.mjs');
    await runPipelineBatch({
      limit: 1,
      seedMode: 'question-first',
      logger: { log() {}, error() {} },
    });

    expect(loadRagChunks).toHaveBeenCalled();
  });
});