import {
  // Models
  AutoModelForSeq2SeqLM,
  AutoModelForCausalLM,
  LlamaForCausalLM,
  LlavaForConditionalGeneration,

  // Tokenizers
  AutoTokenizer,
  LlamaTokenizer,

  // Processors
  AutoProcessor,
  Processor,

  // Other
  TextStreamer,
  RawImage,
} from "../../src/transformers.js";

import { init, MAX_TEST_EXECUTION_TIME, MAX_MODEL_LOAD_TIME, MAX_MODEL_DISPOSE_TIME, DEFAULT_MODEL_OPTIONS } from "../init.js";

// Initialise the testing environment
init();

// Helper function to generate text
const generate = async (model, tokenizer, text, options) => {
  const inputs = tokenizer(text);
  return await model.generate({
    ...inputs,
    ...options,
  });
};

describe("Generation parameters", () => {
  // List all models which will be tested
  const models = [
    "hf-internal-testing/tiny-random-T5ForConditionalGeneration", // encoder-decoder
    "hf-internal-testing/tiny-random-LlamaForCausalLM", // decoder-only
  ];
  const DUMMY_TEXT = "hello";

  describe(`encoder-decoder (${models[0]})`, () => {
    const model_id = models[0];

    let model;
    let tokenizer;
    beforeAll(async () => {
      model = await AutoModelForSeq2SeqLM.from_pretrained(model_id, DEFAULT_MODEL_OPTIONS);
      tokenizer = await AutoTokenizer.from_pretrained(model_id);
    }, MAX_MODEL_LOAD_TIME);

    // NOTE: Since `max_length` defaults to 20, this case also tests that.
    it(
      "default",
      async () => {
        const outputs = await generate(model, tokenizer, DUMMY_TEXT, {});
        expect(outputs.dims.at(-1)).toEqual(20);
      },
      MAX_TEST_EXECUTION_TIME,
    );

    it(
      "max_new_tokens",
      async () => {
        const MAX_NEW_TOKENS = 5;
        const outputs = await generate(model, tokenizer, DUMMY_TEXT, {
          max_new_tokens: MAX_NEW_TOKENS,
        });
        expect(outputs.dims.at(-1)).toEqual(MAX_NEW_TOKENS + 1); // + 1 due to forced BOS token
      },
      MAX_TEST_EXECUTION_TIME,
    );

    it(
      "min_length",
      async () => {
        const MIN_LENGTH = 3;
        const MAX_LENGTH = 5;
        const outputs = await generate(model, tokenizer, DUMMY_TEXT, {
          eos_token_id: 0,
          min_length: MIN_LENGTH,
          max_length: MAX_LENGTH,
        });
        expect(outputs.tolist()).toEqual([[0n, 11924n, 11924n, 11924n, 11924n]]);
        expect(outputs.dims.at(-1)).toBeGreaterThanOrEqual(MIN_LENGTH);
      },
      MAX_TEST_EXECUTION_TIME,
    );

    it(
      "min_new_tokens",
      async () => {
        const MIN_NEW_TOKENS = 2;
        const MAX_LENGTH = 5;
        const outputs = await generate(model, tokenizer, DUMMY_TEXT, {
          eos_token_id: 0,
          min_new_tokens: MIN_NEW_TOKENS,
          max_length: MAX_LENGTH,
        });
        expect(outputs.tolist()).toEqual([[0n, 11924n, 11924n, 11924n, 11924n]]);
        expect(outputs.dims.at(-1)).toBeGreaterThanOrEqual(MIN_NEW_TOKENS);
      },
      MAX_TEST_EXECUTION_TIME,
    );

    afterAll(async () => {
      await model?.dispose();
    }, MAX_MODEL_DISPOSE_TIME);
  });

  describe(`decoder-only (${models[1]})`, () => {
    const model_id = models[1];

    let model;
    let tokenizer;
    beforeAll(async () => {
      model = await AutoModelForCausalLM.from_pretrained(model_id, DEFAULT_MODEL_OPTIONS);
      tokenizer = await AutoTokenizer.from_pretrained(model_id);
    }, MAX_MODEL_LOAD_TIME);

    // NOTE: Since `max_length` defaults to 20, this case also tests that.
    it(
      "default",
      async () => {
        const outputs = await generate(model, tokenizer, DUMMY_TEXT, {});
        expect(outputs.dims.at(-1)).toEqual(20);
      },
      MAX_TEST_EXECUTION_TIME,
    );

    it(
      "max_new_tokens",
      async () => {
        const MAX_NEW_TOKENS = 5;
        const PROMPT_LENGTH = 2; // BOS + DUMMY_TEXT
        const outputs = await generate(model, tokenizer, DUMMY_TEXT, {
          max_new_tokens: MAX_NEW_TOKENS,
        });
        const expected_length = PROMPT_LENGTH + MAX_NEW_TOKENS;
        expect(outputs.dims.at(-1)).toEqual(expected_length);
      },
      MAX_TEST_EXECUTION_TIME,
    );

    it(
      "min_length",
      async () => {
        const MIN_LENGTH = 4;
        const outputs = await generate(model, tokenizer, DUMMY_TEXT, {
          eos_token_id: [
            18547, // min_length will suppress this token (generated by default)
            16012, // stop at this token
          ],
          min_length: MIN_LENGTH,
        });
        expect(outputs.tolist()).toEqual([[1n, 22172n, 31583n, 18824n, 16621n, 8136n, 16012n]]);
        expect(outputs.dims.at(-1)).toBeGreaterThanOrEqual(MIN_LENGTH);
      },
      MAX_TEST_EXECUTION_TIME,
    );

    it(
      "min_new_tokens",
      async () => {
        const MIN_NEW_TOKENS = 2;
        const outputs = await generate(model, tokenizer, DUMMY_TEXT, {
          eos_token_id: [
            18547, // min_new_tokens will suppress this token (generated by default)
            16012, // stop at this token
          ],
          min_new_tokens: MIN_NEW_TOKENS,
        });
        expect(outputs.tolist()).toEqual([[1n, 22172n, 31583n, 18824n, 16621n, 8136n, 16012n]]);
        expect(outputs.dims.at(-1)).toBeGreaterThanOrEqual(MIN_NEW_TOKENS);
      },
      MAX_TEST_EXECUTION_TIME,
    );

    afterAll(async () => {
      await model?.dispose();
    }, MAX_MODEL_DISPOSE_TIME);
  });
});

describe("Streamers", () => {
  describe("decoder-only", () => {
    const model_id = "hf-internal-testing/tiny-random-LlamaForCausalLM";
    let model, tokenizer;
    beforeAll(async () => {
      model = await AutoModelForCausalLM.from_pretrained(model_id, DEFAULT_MODEL_OPTIONS);
      tokenizer = await AutoTokenizer.from_pretrained(model_id);
    }, MAX_MODEL_LOAD_TIME);

    it(
      "batch_size=1",
      async () => {
        const target_chunks = ["helloerdingsdelete ", "melytabular ", "Stadiumoba ", "alcune ", "drug"];
        const chunks = [];
        const callback_function = (text) => {
          chunks.push(text);
        };
        const streamer = new TextStreamer(tokenizer, { callback_function, skip_special_tokens: true });

        const inputs = tokenizer("hello");
        const outputs = await model.generate({
          ...inputs,
          max_length: 10,
          streamer,
        });
        expect(outputs.tolist()).toEqual([[1n, 22172n, 18547n, 8143n, 22202n, 9456n, 17213n, 15330n, 26591n, 15721n]]);
        expect(chunks).toEqual(target_chunks);
      },
      MAX_TEST_EXECUTION_TIME,
    );

    afterAll(async () => {
      await model?.dispose();
    }, MAX_MODEL_DISPOSE_TIME);
  });
});

describe("PKV caching", () => {
  describe("LlamaForCausalLM", () => {
    const model_id = "hf-internal-testing/tiny-random-LlamaForCausalLM";
    /** @type {LlamaForCausalLM} */
    let model;
    /** @type {LlamaTokenizer} */
    let tokenizer;
    beforeAll(async () => {
      model = await LlamaForCausalLM.from_pretrained(model_id, DEFAULT_MODEL_OPTIONS);
      tokenizer = await LlamaTokenizer.from_pretrained(model_id);
    }, MAX_MODEL_LOAD_TIME);

    it(
      "batch_size=1",
      async () => {
        const inputs = tokenizer("1");

        // Generate first sequence w/o PKV
        // NOTE: `return_dict_in_generate=true` is required to get PKV
        const { past_key_values, sequences } = await model.generate({
          ...inputs,
          max_new_tokens: 5,
          do_sample: false,
          return_dict_in_generate: true,
        });

        // Update output with new text
        const decoded = tokenizer.batch_decode(sequences, {
          skip_special_tokens: false,
        })[0];
        const new_inputs = tokenizer(decoded + "2", {
          add_special_tokens: false,
        });

        // Run w/o PKV
        const generated_ids = await model.generate({
          ...new_inputs,
          max_new_tokens: 3,
          do_sample: false,
        });

        // Run w/ PKV
        const generated_ids_pkv = await model.generate({
          ...new_inputs,
          past_key_values,
          max_new_tokens: 3,
          do_sample: false,
        });

        const target = [[1n, 259n, 29896n, 24959n, 22063n, 17192n, 12189n, 22468n, 29906n, 3399n, 24823n, 26470n]];

        expect(generated_ids.tolist()).toEqual(target);
        expect(generated_ids_pkv.tolist()).toEqual(target);
      },
      MAX_TEST_EXECUTION_TIME,
    );

    afterAll(async () => {
      await model?.dispose();
    }, MAX_MODEL_DISPOSE_TIME);
  });

  describe("LlavaForConditionalGeneration", () => {
    const model_id = "Xenova/tiny-random-LlavaForConditionalGeneration";
    /** @type {LlavaForConditionalGeneration} */
    let model;
    /** @type {PreTrainedTokenizer} */
    let tokenizer;
    /** @type {Processor} */
    let processor;
    beforeAll(async () => {
      model = await LlavaForConditionalGeneration.from_pretrained(model_id, DEFAULT_MODEL_OPTIONS);
      tokenizer = await AutoTokenizer.from_pretrained(model_id);
      processor = await AutoProcessor.from_pretrained(model_id);
    }, MAX_MODEL_LOAD_TIME);

    it(
      "batch_size=1",
      async () => {
        const text_inputs = tokenizer("<image>hello");

        // Empty white image
        const dims = [224, 224, 3];
        const image = new RawImage(new Uint8ClampedArray(dims[0] * dims[1] * dims[2]).fill(255), ...dims);
        const vision_inputs = await processor(image);

        // Generate first sequence w/o PKV
        // NOTE: `return_dict_in_generate=true` is required to get PKV
        const { past_key_values, sequences } = await model.generate({
          ...text_inputs,
          ...vision_inputs,
          max_new_tokens: 5,
          do_sample: false,
          return_dict_in_generate: true,
        });

        // Update output with new text
        const decoded = tokenizer.batch_decode(sequences).map((x) => x + "new");
        const new_inputs = tokenizer(decoded, {
          add_special_tokens: false,
        });

        // Run w/o PKV
        const generated_ids = await model.generate({
          ...new_inputs,
          ...vision_inputs,
          max_new_tokens: 3,
          do_sample: false,
        });

        // Run w/ PKV
        const generated_ids_pkv = await model.generate({
          ...new_inputs,
          past_key_values,
          max_new_tokens: 3,
          do_sample: false,
        });

        const target = [[1n, 32000n, 29871n, 23927n, 359n, 1519n, 568n, 5769n, 1330n, 21544n, 11568n, 1482n, 7258n, 1250n, 16117n]];
        expect(generated_ids.tolist()).toEqual(target);
        expect(generated_ids_pkv.tolist()).toEqual(target);
      },
      MAX_TEST_EXECUTION_TIME,
    );

    afterAll(async () => {
      await model?.dispose();
    }, MAX_MODEL_DISPOSE_TIME);
  });
});
