const fs = require('fs/promises');
const os = require('os');
const path = require('path');
const { mockClient } = require('aws-sdk-client-mock');
const { BedrockRuntimeClient, InvokeModelCommand } = require('@aws-sdk/client-bedrock-runtime');
const { chunkBook } = require('../src/chunk');
const { applyCorrections } = require('../src/corrections');
const { embedMany } = require('../src/embed');
const { LocalFsStorage } = require('../src/storage');
const { buildIndexBody } = require('../src/indexMapping');
const { bulkWrite, deleteStale, verifyDocs, indexBook } = require('../src/indexer');
const { migrateRun, swapAlias, cutover, rollback } = require('../src/migrate');

const bedrock = mockClient(BedrockRuntimeClient);

function paragraph(words) {
  return Array.from({ length: words }, (_, index) => `w${index}`).join(' ');
}

describe('chunking', () => {
  const chapters = [
    { chapterNo: 1, title: 'One', startPage: 1, endPage: 1 },
    { chapterNo: 2, title: 'Two', startPage: 2, endPage: 2 },
  ];
  const pages = [
    { page: 1, text: `${paragraph(80)}\n\n${paragraph(80)}` },
    { page: 2, text: `${paragraph(40)}\n\n1.1 Harbor notes\n\n${paragraph(40)}` },
  ];

  it('stays inside a chapter, covers pages, and links neighbors', () => {
    const chunks = chunkBook({ bookId: 5, pages, chapters, targetTokens: 30, overlapTokens: 8 });
    expect(chunks.length).toBeGreaterThan(2);
    for (const chunk of chunks) {
      expect(chunk.pageStart).toBe(chunk.chapterNo);
      expect(chunk.pageEnd).toBe(chunk.chapterNo);
      expect(chunk.text).not.toMatch(/w0[\s\S]*Chapter/);
    }
    const chapter2 = chunks.filter((chunk) => chunk.chapterNo === 2);
    expect(chapter2.some((chunk) => chunk.sectionTitle.includes('Harbor'))).toBe(true);
    expect(chunks[0].chunkId).toBe('5:chunk:00001');
    expect(chunks[0].prevChunkId).toBeNull();
    expect(chunks[0].nextChunkId).toBe(chunks[1].chunkId);
    expect(chunks[chunks.length - 1].nextChunkId).toBeNull();
    const again = chunkBook({ bookId: 5, pages, chapters, targetTokens: 30, overlapTokens: 8 });
    expect(again.map((chunk) => chunk.chunkId)).toEqual(chunks.map((chunk) => chunk.chunkId));
  });

  it('applies corrections before the chunk text is built', () => {
    const corrected = applyCorrections({
      pages,
      chapters,
      corrections: {
        textReplacements: [{ page: 1, find: 'w0', replace: 'FIXED' }],
        dropPages: [2],
        chapterOverrides: [{ title: 'Only', startPage: 1 }],
        metadata: { title: 'Given Title', author: 'Given Author' },
      },
      pageCount: 2,
    });
    expect(corrected.pages).toHaveLength(1);
    expect(corrected.pages[0].text).toContain('FIXED');
    expect(corrected.chapters[0].title).toBe('Only');
    expect(corrected.metadata.author).toBe('Given Author');
    const chunks = chunkBook({
      bookId: 5,
      pages: corrected.pages,
      chapters: corrected.chapters,
      targetTokens: 30,
      overlapTokens: 0,
    });
    expect(chunks.every((chunk) => chunk.chapterNo !== 2)).toBe(true);
    expect(chunks.some((chunk) => chunk.text.includes('FIXED'))).toBe(true);
  });
});

describe('embedding cache', () => {
  it('does not call the model again for the same text', async () => {
    bedrock.reset();
    const vector = Array.from({ length: 4 }, () => 0.1);
    bedrock.on(InvokeModelCommand).resolves({
      body: new TextEncoder().encode(JSON.stringify({ embedding: vector })),
    });
    const dir = await fs.mkdtemp(path.join(os.tmpdir(), 'embed-'));
    const storage = new LocalFsStorage(dir);
    const client = new BedrockRuntimeClient({ region: 'ap-south-1' });
    const first = await embedMany({
      client,
      modelId: 'amazon.titan-embed-text-v2:0',
      dimensions: 4,
      texts: ['alpha', 'beta'],
      storage,
      bookId: 1,
      concurrency: 2,
    });
    const second = await embedMany({
      client,
      modelId: 'amazon.titan-embed-text-v2:0',
      dimensions: 4,
      texts: ['alpha', 'beta'],
      storage,
      bookId: 1,
      concurrency: 2,
    });
    expect(first.calls).toBe(2);
    expect(second.calls).toBe(0);
    expect(second.vectors).toEqual(first.vectors);
    expect(bedrock.commandCalls(InvokeModelCommand)).toHaveLength(2);
    await fs.rm(dir, { recursive: true, force: true });
  });
});

describe('index mapping', () => {
  it('matches the OpenSearch 2.7 lucene cosine mapping', () => {
    const body = buildIndexBody(1024);
    expect(body.settings).toEqual({ index: { knn: true } });
    expect(body.mappings.properties.embedding).toEqual({
      type: 'knn_vector',
      dimension: 1024,
      method: { name: 'hnsw', space_type: 'cosinesimil', engine: 'lucene' },
    });
    expect(Object.keys(body.mappings.properties).sort()).toEqual(
      [
        'book_id',
        'chapter_no',
        'chapter_title',
        'chunk_id',
        'contextual_text',
        'created_at',
        'embedding',
        'ingest_version',
        'language',
        'level',
        'next_chunk_id',
        'ordinal',
        'page_end',
        'page_start',
        'prev_chunk_id',
        'section_title',
        'text',
        'token_count',
      ].sort(),
    );
    expect(body.mappings.properties.text.fields.english.analyzer).toBe('english');
  });
});

function doc(id, source) {
  return { _id: id, level: source.level, _source: source };
}

describe('indexer', () => {
  it('retries failed bulk items and deletes only stale versions of the same book', async () => {
    const calls = [];
    const client = {
      async bulk(params) {
        calls.push(params);
        if (calls.length === 1) {
          return {
            errors: true,
            items: [{ index: { _id: 'a', status: 201 } }, { index: { _id: 'b', status: 429, error: { type: 'rejected' } } }],
          };
        }
        return { errors: false, items: [{ index: { _id: 'b', status: 201 } }] };
      },
      async deleteByQuery(params) {
        calls.push(params);
      },
    };
    const docs = [doc('a', { level: 'chunk' }), doc('b', { level: 'chunk' })];
    await bulkWrite(client, 'books_current_v2', docs);
    expect(calls[1].body.map((row) => row.index && row.index._id).filter(Boolean)).toEqual(['b']);
    await deleteStale(client, 'books_current_v2', 9, 'v9');
    expect(calls[2].body.query).toEqual({
      bool: {
        must: [{ term: { book_id: '9' } }],
        must_not: [{ term: { ingest_version: 'v9' } }],
      },
    });
  });

  it('fails verification when a page is missing or a vector is the wrong length', () => {
    const vector = [0.1, 0.2];
    const good = {
      level: 'chunk',
      chunk_id: '1:chunk:00001',
      chapter_no: 1,
      page_start: 1,
      page_end: 1,
      embedding: vector,
    };
    expect(() => verifyDocs([good, { level: 'book', chapter_no: 0, embedding: vector }], { dimensions: 2, pages: [1, 2] })).toThrow(/page 2/);
    expect(() =>
      verifyDocs(
        [{ ...good, embedding: [0.1] }, { level: 'book', chapter_no: 0, embedding: vector }],
        { dimensions: 2, pages: [1] },
      ),
    ).toThrow(/bad embedding/);
  });

  it('does not call OpenSearch in local mode', async () => {
    const vector = [0.2, 0.3];
    const docs = [
      doc('1:chunk:00001', {
        level: 'chunk',
        chunk_id: '1:chunk:00001',
        chapter_no: 1,
        page_start: 1,
        page_end: 1,
        embedding: vector,
      }),
      doc('1:book', { level: 'book', chapter_no: 0, embedding: vector }),
    ];
    const result = await indexBook({
      bookId: 1,
      docs,
      client: null,
      config: { embedDimensions: 2, ingestVersion: 'v1', indexAlias: 'books_current', indexVersion: 'v2' },
      pages: [1],
      logger: { info() {} },
    });
    expect(result.chunkCount).toBe(1);
    expect(result.indexName).toBeNull();
  });
});

describe('migration', () => {
  it('dry-run makes no rebuild calls', async () => {
    let calls = 0;
    const results = await migrateRun({
      bookIds: ['3', '4'],
      dryRun: true,
      rebuild: async () => {
        calls += 1;
      },
    });
    expect(calls).toBe(0);
    expect(results.map((row) => row.status)).toEqual(['skipped', 'skipped']);
  });

  it('continues after one book fails', async () => {
    const results = await migrateRun({
      bookIds: ['3', '4'],
      dryRun: false,
      rebuild: async (bookId) => {
        if (bookId === '3') throw new Error('missing pdf');
      },
    });
    expect(results[0]).toMatchObject({ status: 'failed', error: 'missing pdf' });
    expect(results[1].status).toBe('success');
  });

  it('swaps the alias in one request and rollback restores the previous target', async () => {
    const calls = [];
    const client = {
      async indices() {},
    };
    client.indices = {
      async updateAliases(params) {
        calls.push(params);
      },
      async getAlias() {
        return { books_current_v1: { aliases: { books_current: {} } } };
      },
    };
    const actions = await swapAlias(client, { alias: 'books_current', addIndex: 'books_current_v2', removeIndex: 'books_current_v1' });
    expect(calls).toHaveLength(1);
    expect(actions).toEqual([
      { remove: { index: 'books_current_v1', alias: 'books_current' } },
      { add: { index: 'books_current_v2', alias: 'books_current' } },
    ]);

    const state = {
      value: null,
      async get() {
        return this.value;
      },
      async put(value) {
        this.value = value;
      },
    };
    await cutover({ client, alias: 'books_current', toIndex: 'books_current_v2', state });
    expect(calls).toHaveLength(2);
    await rollback({ client, state });
    expect(calls).toHaveLength(3);
    expect(calls[2].body.actions).toEqual([
      { remove: { index: 'books_current_v2', alias: 'books_current' } },
      { add: { index: 'books_current_v1', alias: 'books_current' } },
    ]);
  });
});