// The extractor over a scripted ModelGateway: the mapping of the gateway's failures to the codes a
// document records, and the retry each one gets. The whole path through the real gateway client is
// proved by test/documents-extraction.test.ts.

import { AiError, type GenerateObjectOptions, type ModelGateway, ModelGatewayError } from "@conusai-kit/platform"
import { describe, expect, it } from "vitest"
import { modelLineExtractor } from "../adapters/model-extractor.js"
import { ExtractionFailed, type ExtractionInput, type ReadDocument } from "../core/ports.js"
import { noCache, type StageCache } from "../core/stage-cache.js"

const document: ReadDocument = {
  engine: { name: "fixture", version: "1" },
  page_count: 1,
  pages: [
    {
      index: 0,
      width: 100,
      height: 100,
      source: "text_layer",
      text: "Reference: QR-1",
      class: "text",
      lines: [{ text: "Reference: QR-1", bbox: [10, 10, 60, 20], confidence: 0.9, order: 0, repeated: false }],
      blocks: [{ id: "p1-b1", bbox: [10, 10, 60, 20], order: 0, lines: [0] }],
    },
  ],
}

const input = (cache: StageCache = noCache): ExtractionInput => ({ document, pageImage: async () => undefined, cache })

/** A gateway that answers from a script, one entry per call: a value is the reply, an Error is thrown. */
function scripted(...script: Array<unknown>): { gateway: ModelGateway; options: GenerateObjectOptions[] } {
  const options: GenerateObjectOptions[] = []
  const gateway: ModelGateway = {
    async generateObject(_schema, _messages, call) {
      options.push(call ?? {})
      const next = script.shift()
      if (next instanceof Error) throw next
      return next as never
    },
  }
  return { gateway, options }
}

const good = { lines: [{ label: "Reference", value: "QR-1", quote: "Reference: QR-1" }] }
const extractorOf = (gateway: ModelGateway) => modelLineExtractor({ gateway, model: "m" })

describe("the model extractor", () => {
  it("names the model and the prompt version, which every proposal records", () => {
    const extractor = extractorOf(scripted().gateway)
    expect(extractor.name).toBe("m")
    expect(extractor.version).toMatch(/^[a-z-]+\.v\d+$/)
  })

  it.each([
    ["refusal", "extraction_refused", 2],
    ["truncated", "extraction_truncated", 2],
    ["empty", "extraction_invalid", 2],
    // a schema failure was already asked again by the gateway: the extractor does not ask a third time
    ["schema", "extraction_invalid", 1],
  ] as const)("a %s is asked again as the rules say, then ends as %s", async (kind, code, calls) => {
    const failure = () => new ModelGatewayError(kind, "x")
    const { gateway, options } = scripted(failure(), failure())
    await expect(extractorOf(gateway).extract(input())).rejects.toMatchObject({ code })
    expect(options).toHaveLength(calls)
  })

  it("answers when the second attempt is good, giving a truncated reply twice the room", async () => {
    const { gateway, options } = scripted(new ModelGatewayError("truncated", "x"), good)
    const { lines } = await extractorOf(gateway).extract(input())
    expect(lines.map((line) => line.value)).toEqual(["QR-1"])
    expect(options.map((call) => call.maxTokens)).toEqual([4096, 8192])
  })

  it("treats a request the gateway will never accept as the document's failure", async () => {
    const bad = () => new AiError("bad image", { status: 400 })
    const { gateway } = scripted(bad(), bad())
    await expect(extractorOf(gateway).extract(input())).rejects.toMatchObject({
      code: "extraction_invalid",
      name: ExtractionFailed.name,
    })
  })

  it.each([
    ["a network error", new AiError("down", { code: "network_error" })],
    ["a 503", new AiError("busy", { status: 503 })],
    ["a missing token scope", new AiError("no", { status: 403 })],
    ["a missing model", new ModelGatewayError("no_model", "x")],
  ])("leaves %s to the queue, which retries the job", async (_name, error) => {
    const { gateway, options } = scripted(error)
    await expect(extractorOf(gateway).extract(input())).rejects.toBe(error)
    expect(options).toHaveLength(1)
  })

  it("reads a cached reply without asking the model", async () => {
    const stored = new Map<string, unknown>()
    const cache: StageCache = {
      async get(stage, key, schema) {
        const value = stored.get(`${stage}/${key}`)
        return value === undefined ? undefined : schema.parse(value)
      },
      async put(stage, key, value) {
        stored.set(`${stage}/${key}`, value)
      },
    }
    const { gateway, options } = scripted(good)
    await extractorOf(gateway).extract(input(cache))
    await extractorOf(gateway).extract(input(cache))
    expect(options).toHaveLength(1)
    // another model is another key
    const other = scripted(good)
    await modelLineExtractor({ gateway: other.gateway, model: "other" }).extract(input(cache))
    expect(other.options).toHaveLength(1)
  })
})