import { describe, expect, test } from "bun:test";
import { Tensor } from "@huggingface/transformers";
import {
  PROMPT_PREFIX_KWARG,
  WHISPER_DECODER_POSITIONS,
  buildPromptPrefix,
  installPromptSupport,
  type PromptableModel,
} from "../src/worker/prompt";

const PREV = 50362;

/**
 * Stand-in for a Transformers.js Whisper model: records the kwargs each
 * generate() call received and, like the real one, returns the decoder input
 * followed by the "generated" tokens as an int64 [1, n] tensor.
 */
function fakeModel(initTokens: number[], generated: number[] = [50365, 400, 401, 50400]) {
  const calls: Record<string, unknown>[] = [];
  const model: PromptableModel & { calls: typeof calls } = {
    calls,
    _prepare_generation_config: (_config, kwargs) => ({ fromKwargs: kwargs }),
    _retrieve_init_tokens: () => [...initTokens],
    async generate(kwargs) {
      calls.push(kwargs);
      const input = (kwargs.decoder_input_ids as number[] | undefined) ?? initTokens;
      const sequence = [...input, ...generated];
      return new Tensor("int64", BigInt64Array.from(sequence.map(BigInt)), [1, sequence.length]);
    },
  };
  return model;
}

function tokens(output: unknown): number[] {
  return (output as Tensor).tolist()[0]!.map(Number);
}

describe("installPromptSupport (D-WS2a)", () => {
  test("without a prompt kwarg, generate is called exactly as before", async () => {
    const model = fakeModel([50258, 50259, 50360]);
    installPromptSupport(model);
    const kwargs = { max_new_tokens: 440, return_timestamps: true };
    const output = await model.generate({ ...kwargs });
    expect(model.calls).toEqual([kwargs]);
    expect(tokens(output)).toEqual([50258, 50259, 50360, 50365, 400, 401, 50400]);
  });

  test("an empty prefix is the same as no prompt", async () => {
    const model = fakeModel([50258]);
    installPromptSupport(model);
    await model.generate({ max_new_tokens: 440, [PROMPT_PREFIX_KWARG]: [] });
    expect(model.calls).toEqual([{ max_new_tokens: 440 }]);
  });

  test("a prefix is prepended to the init tokens and stripped from the output", async () => {
    const model = fakeModel([50258, 50259, 50360]);
    installPromptSupport(model);
    const output = await model.generate({
      max_new_tokens: 440,
      language: "en",
      [PROMPT_PREFIX_KWARG]: [PREV, 1, 2],
    });
    const call = model.calls[0]!;
    expect(call.decoder_input_ids).toEqual([PREV, 1, 2, 50258, 50259, 50360]);
    expect(call).not.toHaveProperty(PROMPT_PREFIX_KWARG);
    expect(call.language).toBe("en");
    // Prompt tokens would otherwise reach _decode_asr as transcript text.
    expect(tokens(output)).toEqual([50258, 50259, 50360, 50365, 400, 401, 50400]);
  });

  test("works for English-only models (init tokens are just <|startoftranscript|>)", async () => {
    const model = fakeModel([50257]);
    installPromptSupport(model);
    const output = await model.generate({ max_new_tokens: 440, [PROMPT_PREFIX_KWARG]: [50361, 7] });
    expect(model.calls[0]!.decoder_input_ids).toEqual([50361, 7, 50257]);
    expect(tokens(output)[0]).toBe(50257);
  });

  test("caps max_new_tokens so prompt + init + new tokens fit Whisper's 448 positions", async () => {
    const model = fakeModel([50258, 50259, 50360]);
    installPromptSupport(model);
    const prefix = [PREV, ...Array.from({ length: 223 }, (_, i) => 1000 + i)];
    await model.generate({ max_new_tokens: 440, [PROMPT_PREFIX_KWARG]: prefix });
    expect(WHISPER_DECODER_POSITIONS).toBe(448);
    expect(model.calls[0]!.max_new_tokens).toBe(448 - 227);
    // A short prompt keeps the D-021 value.
    await model.generate({ max_new_tokens: 440, [PROMPT_PREFIX_KWARG]: [PREV, 5] });
    expect(model.calls[1]!.max_new_tokens).toBe(440);
  });

  test("installing twice wraps once", async () => {
    const model = fakeModel([50258]);
    installPromptSupport(model);
    installPromptSupport(model);
    await model.generate({ max_new_tokens: 440, [PROMPT_PREFIX_KWARG]: [PREV, 5] });
    expect(model.calls).toHaveLength(1);
    expect(model.calls[0]!.decoder_input_ids).toEqual([PREV, 5, 50258]);
  });
});

describe("buildPromptPrefix", () => {
  const tokenizer = {
    encode: (text: string, _options: { add_special_tokens: false }) =>
      [...text].map((c) => c.charCodeAt(0)),
    convert_tokens_to_ids: (token: string) => (token === "<|startofprev|>" ? PREV : -1),
  };

  test("no hint → no prefix", () => {
    expect(buildPromptPrefix(tokenizer, "")).toEqual([]);
    expect(buildPromptPrefix(tokenizer, " \n ")).toEqual([]);
  });

  test("a hint → <|startofprev|> then the space-led, full-stopped hint tokens", () => {
    expect(buildPromptPrefix(tokenizer, "Max\nVoid")).toEqual([
      PREV,
      ...[..." Max, Void."].map((c) => c.charCodeAt(0)),
    ]);
  });

  test("a huge hint stays within the 223-token budget", () => {
    const prefix = buildPromptPrefix(tokenizer, "Maxwell, ".repeat(500));
    expect(prefix.length).toBeLessThanOrEqual(224);
  });
});