import { describe, expect, it, vi } from 'vitest'
import { CohereRerankClient } from './cohere-rerank-client'

function response(results: unknown, init?: ResponseInit): Response {
  return new Response(JSON.stringify({ results }), init)
}

describe('CohereRerankClient', () => {
  it('posts the exact Cohere/Jina request to the exact configured endpoint', async () => {
    const transport = vi.fn<typeof fetch>(async () =>
      response([
        { index: 1, relevance_score: 0.9 },
        { index: 0, relevance_score: 0.4 }
      ])
    )
    const client = new CohereRerankClient({
      endpoint: 'https://rerank.example/custom/v1/rerank?version=2',
      model: 'vendor/rerank-large',
      apiKey: 'rerank-secret',
      fetch: transport
    })

    await expect(
      client.rerank('find this', ['first', 'second'], 2)
    ).resolves.toEqual([
      { index: 1, relevanceScore: 0.9 },
      { index: 0, relevanceScore: 0.4 }
    ])
    expect(transport).toHaveBeenCalledTimes(1)
    const [endpoint, init] = transport.mock.calls[0] ?? []
    expect(endpoint).toBe(
      'https://rerank.example/custom/v1/rerank?version=2'
    )
    expect(init).toMatchObject({
      method: 'POST',
      redirect: 'error'
    })
    expect(init?.headers).toEqual({
      accept: 'application/json',
      'content-type': 'application/json',
      authorization: 'Bearer rerank-secret'
    })
    expect(JSON.parse(String(init?.body))).toEqual({
      model: 'vendor/rerank-large',
      query: 'find this',
      documents: ['first', 'second'],
      top_n: 2,
      return_documents: false
    })
  })

  it('uses safe defaults and supports endpoints without authentication', async () => {
    const transport = vi.fn<typeof fetch>(async () =>
      response([{ index: 0, relevance_score: 1 }])
    )
    const client = new CohereRerankClient({ fetch: transport })

    await client.rerank('query', ['document'], 1)
    expect(transport.mock.calls[0]?.[0]).toBe(
      'https://api.cohere.com/v1/rerank'
    )
    expect(transport.mock.calls[0]?.[1]?.headers).not.toHaveProperty(
      'authorization'
    )
    expect(JSON.parse(String(transport.mock.calls[0]?.[1]?.body))).toMatchObject(
      { model: 'rerank-v3.5' }
    )
  })

  it('distinguishes timeout from caller cancellation', async () => {
    const waitForAbort = vi.fn<typeof fetch>(
      async (_input, init) =>
        new Promise<Response>((_resolve, reject) => {
          init?.signal?.addEventListener(
            'abort',
            () => reject(init.signal?.reason),
            { once: true }
          )
        })
    )
    const client = new CohereRerankClient({
      timeoutMs: 10,
      fetch: waitForAbort
    })
    await expect(client.rerank('query', ['document'], 1)).rejects.toMatchObject({
      name: 'TimeoutError',
      message: 'Rerank request timed out'
    })

    const caller = new AbortController()
    const cancelled = client.rerank('query', ['document'], 1, caller.signal)
    caller.abort(new Error('secret caller reason'))
    await expect(cancelled).rejects.toMatchObject({
      name: 'AbortError',
      message: 'Rerank request was cancelled'
    })

    const preCancelled = new AbortController()
    preCancelled.abort(new Error('cancel before transport'))
    await expect(
      client.rerank('query', ['document'], 1, preCancelled.signal)
    ).rejects.toMatchObject({
      name: 'AbortError',
      message: 'Rerank request was cancelled'
    })
    expect(waitForAbort).toHaveBeenCalledTimes(2)
  })

  it.each([400, 401, 404, 429, 500, 503])(
    'reports HTTP %i without reading or exposing the response body',
    async (status) => {
      const secretBody = 'secret response body from https://private.example'
      const client = new CohereRerankClient({
        endpoint: 'https://rerank.example/v1/rerank',
        apiKey: 'secret-key',
        fetch: async () => new Response(secretBody, { status })
      })
      const error = await client
        .rerank('query', ['document'], 1)
        .catch((caught: unknown) => caught)
      expect(error).toMatchObject({
        message: `Rerank request failed with HTTP ${status}`
      })
      expect(String(error)).not.toContain(secretBody)
      expect(String(error)).not.toContain('secret-key')
      expect(String(error)).not.toContain('rerank.example')
    }
  )

  it.each([
    ['invalid JSON', () => new Response('{')],
    ['missing results', () => new Response('{}')],
    [
      'extra root fields',
      () => new Response('{"results":[],"meta":{"secret":true}}')
    ],
    [
      'provider documents',
      () =>
        response([
          {
            index: 0,
            relevance_score: 0.8,
            document: { text: 'must not be consumed' }
          }
        ])
    ],
    [
      'duplicate indexes',
      () =>
        response([
          { index: 0, relevance_score: 0.8 },
          { index: 0, relevance_score: 0.7 }
        ])
    ],
    [
      'out-of-range indexes',
      () => response([{ index: 2, relevance_score: 0.8 }])
    ],
    [
      'scores above one',
      () => response([{ index: 0, relevance_score: 1.1 }])
    ],
    [
      'non-numeric scores',
      () => response([{ index: 0, relevance_score: 'NaN' }])
    ]
  ])('rejects malformed response: %s', async (_name, makeResponse) => {
    const client = new CohereRerankClient({
      fetch: async () => makeResponse()
    })
    const documents =
      _name === 'duplicate indexes' ? ['one', 'two'] : ['one']
    await expect(
      client.rerank('query', documents, documents.length)
    ).rejects.toThrow()
  })

  it('rejects non-finite scores encoded with overflowing JSON numbers', async () => {
    for (const relevanceScore of ['1e400', '-1e400']) {
      const client = new CohereRerankClient({
        fetch: async () =>
          new Response(
            `{"results":[{"index":0,"relevance_score":${relevanceScore}}]}`
          )
      })
      await expect(client.rerank('query', ['one'], 1)).rejects.toThrow(
        'invalid score'
      )
    }
  })

  it('requires exactly topN unique results and allows that to be fewer than candidates', async () => {
    const accepted = new CohereRerankClient({
      fetch: async () =>
        response([
          { index: 3, relevance_score: 0.9 },
          { index: 1, relevance_score: 0.8 }
        ])
    })
    await expect(
      accepted.rerank('query', ['zero', 'one', 'two', 'three'], 2)
    ).resolves.toHaveLength(2)

    for (const count of [1, 3, 4]) {
      const rejected = new CohereRerankClient({
        fetch: async () =>
          response(
            Array.from({ length: count }, (_, index) => ({
              index,
              relevance_score: 1 - index / 10
            }))
          )
      })
      await expect(
        rejected.rerank('query', ['zero', 'one', 'two', 'three'], 2)
      ).rejects.toThrow('exactly 2 results')
    }
  })

  it('sorts scores descending and ties by original document index', async () => {
    const client = new CohereRerankClient({
      fetch: async () =>
        response([
          { index: 3, relevance_score: 0.5 },
          { index: 2, relevance_score: 0.9 },
          { index: 0, relevance_score: 0.5 },
          { index: 1, relevance_score: 0.9 }
        ])
    })
    await expect(
      client.rerank('query', ['zero', 'one', 'two', 'three'], 4)
    ).resolves.toEqual([
      { index: 1, relevanceScore: 0.9 },
      { index: 2, relevanceScore: 0.9 },
      { index: 0, relevanceScore: 0.5 },
      { index: 3, relevanceScore: 0.5 }
    ])
  })

  it('enforces query, candidate, document and encoded body bounds', async () => {
    const transport = vi.fn<typeof fetch>()
    const client = new CohereRerankClient({ fetch: transport })
    await expect(client.rerank('x'.repeat(4_001), ['one'], 1)).rejects.toThrow(
      'query must be at most 4000'
    )
    await expect(client.rerank('query', [], 1)).rejects.toThrow(
      'documents must contain'
    )
    await expect(
      client.rerank('query', Array.from({ length: 101 }, () => 'x'), 1)
    ).rejects.toThrow('documents must contain')
    await expect(client.rerank('query', ['x'.repeat(8_001)], 1)).rejects.toThrow(
      'documents[0] must be at most 8000'
    )
    // UTF-8 can exceed the body bound while every string remains under its
    // character limit.
    await expect(
      client.rerank(
        'query',
        Array.from({ length: 100 }, () => '汉'.repeat(8_000)),
        100
      )
    ).rejects.toThrow('request body is too large')
    expect(transport).not.toHaveBeenCalled()
  })

  it('bounds declared and streamed response bodies to one MiB', async () => {
    const declared = new CohereRerankClient({
      fetch: async () =>
        new Response('{}', {
          headers: { 'content-length': String(1024 * 1024 + 1) }
        })
    })
    await expect(declared.rerank('query', ['one'], 1)).rejects.toThrow(
      'response is too large'
    )

    const streamed = new CohereRerankClient({
      fetch: async () =>
        new Response(new Uint8Array(1024 * 1024 + 1))
    })
    await expect(streamed.rerank('query', ['one'], 1)).rejects.toThrow(
      'response is too large'
    )
  })

  it('rejects unsafe endpoints without echoing their value', () => {
    const endpoint = 'file:///private/secret'
    expect(() => new CohereRerankClient({ endpoint })).toThrow(
      'endpoint must use HTTP or HTTPS'
    )
    try {
      new CohereRerankClient({ endpoint: 'not-a-url secret-token' })
    } catch (error) {
      expect(String(error)).not.toContain('secret-token')
    }
  })
})