/**
 * @file Tokenization utilities
 *
 * @module tokenizers
 */

import { Tokenizer } from '@huggingface/tokenizers';
import { Template } from '@huggingface/jinja';
import { Callable } from './utils/generic.js';

import { isIntegralNumber, mergeArrays } from './utils/core.js';
import { getModelJSON } from './utils/hub.js';
import { max } from './utils/maths.js';
import { Tensor } from './utils/tensor.js';
import { logger } from './utils/logger.js';
import { get_tokenizer_files } from './utils/model_registry/get_tokenizer_files.js';

/**
 * @typedef {import('./utils/hub.js').PretrainedOptions} PretrainedTokenizerOptions
 */

/**
 * Loads a tokenizer from the specified path.
 * @param {string} pretrained_model_name_or_path The path to the tokenizer directory.
 * @param {PretrainedTokenizerOptions} options Additional options for loading the tokenizer.
 * @returns {Promise<any[]>} A promise that resolves with information about the loaded tokenizer.
 */
export async function loadTokenizer(pretrained_model_name_or_path, options) {
    const tokenizerFiles = await get_tokenizer_files(pretrained_model_name_or_path);
    return await Promise.all(
        tokenizerFiles.map((file) => getModelJSON(pretrained_model_name_or_path, file, true, options)),
    );
}

/**
 * Helper function to convert a tensor to a list before decoding.
 * @param {Tensor} tensor The tensor to convert.
 * @returns {number[]} The tensor as a list.
 */
export function prepareTensorForDecode(tensor) {
    const dims = tensor.dims;
    switch (dims.length) {
        case 1:
            return tensor.tolist();
        case 2:
            if (dims[0] !== 1) {
                throw new Error(
                    'Unable to decode tensor with `batch size !== 1`. Use `tokenizer.batch_decode(...)` for batched inputs.',
                );
            }
            return tensor.tolist()[0];
        default:
            throw new Error(`Expected tensor to have 1-2 dimensions, got ${dims.length}.`);
    }
}

const SPECIAL_TOKEN_ATTRIBUTES = [
    'bos_token',
    'eos_token',
    'unk_token',
    'sep_token',
    'pad_token',
    'cls_token',
    'mask_token',
    // additional_special_tokens (TODO)
];

/**
 * @typedef {{ type: 'text', text: string, [key: string]: any }} TextContent
 * @property {'text'} type The type of content (must be 'text').
 * @property {string} text The text content.
 */

/**
 * @typedef {{ type: 'image', image?: string | import('./utils/image.js').RawImage, [key: string]: any }} ImageContent
 * @property {'image'} type The type of content (must be 'image').
 * @property {string | import('./utils/image.js').RawImage} [image] Optional URL or instance of the image.
 *
 * Note: This works for SmolVLM. Qwen2VL and Idefics3 have different implementations.
 */

/**
 * @typedef {TextContent | ImageContent | { type: string & {}, [key: string]: any }} MessageContent
 * Base type for message content. This is a discriminated union that can be extended with additional content types.
 * Example: `@typedef {TextContent | ImageContent | AudioContent} MessageContent`
 */

/**
 * @typedef {Object} Message
 * @property {'user' | 'assistant' | 'system' | (string & {})} role The role of the message.
 * @property {string | MessageContent[]} content The content of the message. Can be a simple string or an array of content objects.
 */

/**
 *
 * Helper function for padding values of an object, which are each arrays.
 * NOTE: No additional checks are made here for validity of arguments.
 * @param {Record<string, any[]>} item The input object.
 * @param {number} length The length to pad to.
 * @param {(key: string) => any} value_fn Determine the value to fill the array, based on its key.
 * @param {string} side Which side to pad the array.
 * @private
 */
function padHelper(item, length, value_fn, side) {
    for (const key of Object.keys(item)) {
        const diff = length - item[key].length;
        const value = value_fn(key);

        const padData = new Array(diff).fill(value);
        item[key] = side === 'right' ? mergeArrays(item[key], padData) : mergeArrays(padData, item[key]);
    }
}

/**
 * Helper function for truncating values of an object, which are each arrays.
 * NOTE: No additional checks are made here for validity of arguments.
 * @param {Record<string, any[]>} item The input object.
 * @param {number} length The length to truncate to.
 * @private
 */
function truncateHelper(item, length) {
    // Setting .length to a lower value truncates the array in-place:
    // https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/Array/length
    for (const key of Object.keys(item)) {
        item[key].length = length;
    }
}

/**
 * Returns the value of the first matching key in the tokenizer config object.
 * @param {Object} config The tokenizer config object.
 * @param {...string} keys One or more keys to search for in the tokenizer config object.
 * @returns {string|null} The value associated with the first matching key, or null if no match is found.
 * @throws {Error} If an object is found for a matching key and its __type property is not "AddedToken".
 * @private
 */
function getTokenFromConfig(config, ...keys) {
    for (const key of keys) {
        if (!Object.hasOwn(config, key)) continue;
        const item = config[key];
        if (!item) continue;

        if (typeof item === 'object') {
            if (item.__type === 'AddedToken') {
                return item.content;
            } else {
                throw Error(`Unknown token: ${item}`);
            }
        } else {
            return item;
        }
    }
    return null;
}

/**
 *
 * @param {import('@huggingface/tokenizers').Tokenizer} tokenizer
 * @returns {import('@huggingface/tokenizers').AddedToken[]}
 * @private
 */
function getSpecialTokens(tokenizer) {
    const special = [];
    for (const value of tokenizer.get_added_tokens_decoder().values()) {
        if (value.special) special.push(value);
    }
    return special;
}

/**
 * @template {string|string[]} TText
 * @typedef {TText extends string ? number[] : number[][]} BatchEncodingArrayItem
 */

/**
 * @template {string|string[]} TText
 * @template {boolean} [TReturnTensor=true]
 * @typedef {TReturnTensor extends true ? Tensor : BatchEncodingArrayItem<TText>} BatchEncodingItem
 */

/**
 * @template TItem
 * @typedef {Object} BatchEncoding
 * @property {TItem} input_ids List of token ids to be fed to a model.
 * @property {TItem} attention_mask List of indices specifying which tokens should be attended to by the model.
 * @property {TItem} [token_type_ids] List of token type ids to be fed to a model.
 */

/**
 * @template {string|string[]} TText
 * @template {boolean} [TReturnTensor=true]
 * @typedef {Object} TokenizerCallOptions
 * @property {TText extends string ? string|null : string[]|null} [text_pair=null] Optional second sequence to be encoded. If set, must be the same type as text.
 * @property {boolean|'max_length'} [padding=false] Whether to pad the input sequences.
 * @property {boolean} [add_special_tokens=true] Whether or not to add the special tokens associated with the corresponding model.
 * @property {boolean|null} [truncation=null] Whether to truncate the input sequences.
 * @property {number|null} [max_length=null] Maximum length of the returned list and optionally padding length.
 * @property {TReturnTensor} [return_tensor=true] Whether to return the results as Tensors or arrays.
 * @property {boolean|null} [return_token_type_ids=null] Whether to return the token type ids.
 */

/**
 * @typedef {<TText extends string | string[], TReturnTensor extends boolean = true>(text: TText, options?: TokenizerCallOptions<TText, TReturnTensor>) => BatchEncoding<BatchEncodingItem<TText, TReturnTensor>>} PreTrainedTokenizerCallback
 */

/**
 * @template {boolean} [TTokenize=true]
 * @template {boolean} [TReturnTensor=true]
 * @template {boolean} [TReturnDict=true]
 * @typedef {Object} ApplyChatTemplateOptions
 * @property {string|null} [chat_template=null] A Jinja template to use for this conversion.
 * @property {Object[]|null} [tools=null] A list of tools (callable functions) that will be accessible to the model.
 * @property {Record<string, string>[]|null} [documents=null] Documents that will be accessible to the model.
 * @property {boolean} [add_generation_prompt=false] Whether to end the prompt with the token(s) that indicate the start of an assistant message.
 * @property {TTokenize} [tokenize=true] Whether to tokenize the output. If false, the output will be a string.
 * @property {boolean} [padding=false] Whether to pad sequences to the maximum length. Has no effect if tokenize is false.
 * @property {boolean} [truncation=false] Whether to truncate sequences to the maximum length. Has no effect if tokenize is false.
 * @property {number|null} [max_length=null] Maximum length (in tokens) to use for padding or truncation. Has no effect if tokenize is false.
 * @property {TReturnTensor} [return_tensor=true] Whether to return the output as a Tensor or an Array. Has no effect if tokenize is false.
 * @property {TReturnDict} [return_dict=true] Whether to return a dictionary with named outputs. Has no effect if tokenize is false.
 * @property {Object} [tokenizer_kwargs={}] Additional options to pass to the tokenizer.
 */

/**
 * @template {boolean} [TTokenize=true]
 * @template {boolean} [TReturnTensor=true]
 * @template {boolean} [TReturnDict=true]
 * @typedef {TTokenize extends false ? string : TReturnDict extends false ? BatchEncodingItem<string, TReturnTensor> : BatchEncoding<BatchEncodingItem<string, TReturnTensor>>} ApplyChatTemplateReturn
 */

export class PreTrainedTokenizer
    extends /** @type {new (tokenizerJSON: Object, tokenizerConfig: Object) => PreTrainedTokenizerCallback} */ (
        Callable
    )
{
    return_token_type_ids = false;

    padding_side = 'right';
    /**
     * Create a new PreTrainedTokenizer instance.
     * @param {Object} tokenizerJSON The JSON of the tokenizer.
     * @param {Object} tokenizerConfig The config of the tokenizer.
     */
    constructor(tokenizerJSON, tokenizerConfig) {
        super();

        this._tokenizerJSON = tokenizerJSON;
        this._tokenizerConfig = tokenizerConfig;
        this._tokenizer = new Tokenizer(tokenizerJSON, tokenizerConfig);

        this.config = tokenizerConfig;

        this.padding_side = tokenizerConfig.padding_side ?? this.padding_side;

        // Set mask token if present (otherwise will be undefined, which is fine)
        this.mask_token = getTokenFromConfig(tokenizerConfig, 'mask_token');
        this.mask_token_id = this._tokenizer.token_to_id(this.mask_token);

        this.pad_token = getTokenFromConfig(tokenizerConfig, 'pad_token', 'eos_token');
        this.pad_token_id = this._tokenizer.token_to_id(this.pad_token);

        this.sep_token = getTokenFromConfig(tokenizerConfig, 'sep_token');
        this.sep_token_id = this._tokenizer.token_to_id(this.sep_token);

        this.unk_token = getTokenFromConfig(tokenizerConfig, 'unk_token');
        this.unk_token_id = this._tokenizer.token_to_id(this.unk_token);

        this.bos_token = getTokenFromConfig(tokenizerConfig, 'bos_token');
        this.bos_token_id = this._tokenizer.token_to_id(this.bos_token);

        this.eos_token = getTokenFromConfig(tokenizerConfig, 'eos_token');
        this.eos_token_id = this._tokenizer.token_to_id(this.eos_token);

        this.chat_template = tokenizerConfig.chat_template ?? null;
        if (Array.isArray(this.chat_template)) {
            // Chat templates are stored as lists of dicts with fixed key names,
            // we reconstruct that into a single dict while loading them.
            const chat_template = Object.create(null);
            for (const { name, template } of this.chat_template) {
                if (typeof name !== 'string' || typeof template !== 'string') {
                    throw new Error('Chat template must be a list of objects with "name" and "template" properties');
                }
                chat_template[name] = template;
            }
            this.chat_template = chat_template;
        }
        this._compiled_template_cache = new Map();

        const special_tokens = getSpecialTokens(this._tokenizer);
        this.all_special_ids = special_tokens.map((t) => t.id);
        this.all_special_tokens = special_tokens.map((t) => t.content);
    }

    /**
     * Loads a pre-trained tokenizer from the given `pretrained_model_name_or_path`.
     *
     * @param {string} pretrained_model_name_or_path The path to the pre-trained tokenizer.
     * @param {PretrainedTokenizerOptions} options Additional options for loading the tokenizer.
     *
     * @throws {Error} Throws an error if the tokenizer.json or tokenizer_config.json files are not found in the `pretrained_model_name_or_path`.
     * @returns {Promise<PreTrainedTokenizer>} A new instance of the `PreTrainedTokenizer` class.
     */
    static async from_pretrained(
        pretrained_model_name_or_path,
        { progress_callback = null, config = null, cache_dir = null, local_files_only = false, revision = 'main' } = {},
    ) {
        const info = await loadTokenizer(pretrained_model_name_or_path, {
            progress_callback,
            config,
            cache_dir,
            local_files_only,
            revision,
        });

        // @ts-ignore
        return new this(...info);
    }

    get_vocab() {
        return this._tokenizer.get_vocab();
    }

    get model_max_length() {
        return this._tokenizerConfig.model_max_length ?? Infinity;
    }

    get add_eos_token() {
        return this._tokenizerConfig.add_eos_token;
    }
    get add_bos_token() {
        return this._tokenizerConfig.add_bos_token;
    }

    /**
     * Converts a token string (or a sequence of tokens) into a single integer id (or a sequence of ids), using the vocabulary.
     *
     * @template {string|string[]} T
     * @param {T} tokens One or several token(s) to convert to token id(s).
     * @returns {T extends string ? number : number[]} The token id or list of token ids.
     */
    convert_tokens_to_ids(tokens) {
        if (typeof tokens === 'string') {
            return /** @type {any} */ (this._tokenizer.token_to_id(tokens));
        } else {
            return /** @type {any} */ (tokens.map((token) => this._tokenizer.token_to_id(token)));
        }
    }

    /**
     * Encode/tokenize the given text(s).
     * @template {string|string[]} TText
     * @template {boolean} [TReturnTensor=true]
     * @param {TText} text The text to tokenize.
     * @param {TokenizerCallOptions<TText, TReturnTensor>} [options] Additional tokenization options.
     * @returns {BatchEncoding<BatchEncodingItem<TText, TReturnTensor>>} Object to be passed to the model.
     */
    _call(
        // Required positional arguments
        text,
        options = {},
    ) {
        const { text_pair = null, add_special_tokens = true, padding = false, return_token_type_ids = null } = options;
        let { truncation = null, max_length = null } = options;
        const return_tensor = /** @type {TReturnTensor} */ (options.return_tensor ?? true); // Different to HF

        const isBatched = Array.isArray(text);

        let encodedTokens;

        if (isBatched) {
            if (text.length === 0) {
                throw Error('text array must be non-empty');
            }

            if (text_pair !== null) {
                if (!Array.isArray(text_pair)) {
                    throw Error('text_pair must also be an array');
                } else if (text.length !== text_pair.length) {
                    throw Error('text and text_pair must have the same length');
                }

                encodedTokens = text.map((t, i) =>
                    this._encode_plus(t, { text_pair: text_pair[i], add_special_tokens, return_token_type_ids }),
                );
            } else {
                encodedTokens = text.map((x) => this._encode_plus(x, { add_special_tokens, return_token_type_ids }));
            }
        } else {
            if (text === null || text === undefined) {
                throw Error('text may not be null or undefined');
            }

            if (Array.isArray(text_pair)) {
                throw Error(
                    'When specifying `text_pair`, since `text` is a string, `text_pair` must also be a string (i.e., not an array).',
                );
            }

            // For single input, we just wrap in an array, and then unwrap later.
            encodedTokens = [this._encode_plus(text, { text_pair, add_special_tokens, return_token_type_ids })];
        }
        // At this point, `encodedTokens` is batched, of shape [batch_size, tokens].
        // However, array may be jagged. So, we may need pad to max_length.
        if (max_length === null) {
            max_length = this.model_max_length;
        } else if (truncation === null) {
            if (padding === true) {
                logger.warn(
                    '`max_length` is ignored when `padding: true` and there is no truncation strategy. ' +
                        "To pad to max length, use `padding: 'max_length'`.",
                );
                max_length = this.model_max_length;
            } else if (padding === false) {
                logger.warn(
                    'Truncation was not explicitly activated but `max_length` is provided a specific value, please use `truncation: true` to explicitly truncate examples to max length.',
                );
                truncation = true;
            }
        }

        // padding: 'max_length' doesn't require any additional calculation
        // but padding: true has to calculate max_length from the sequences
        if (padding === true) {
            max_length = Math.min(max(encodedTokens.map((x) => x.input_ids.length))[0], max_length ?? Infinity);
        }

        // Ensure it is less than model max length
        max_length = Math.min(max_length, this.model_max_length ?? Infinity);

        if (padding || truncation) {
            // Perform padding and/or truncation
            for (let i = 0; i < encodedTokens.length; ++i) {
                if (encodedTokens[i].input_ids.length === max_length) {
                    continue;
                } else if (encodedTokens[i].input_ids.length > max_length) {
                    // possibly truncate
                    if (truncation) {
                        truncateHelper(encodedTokens[i], max_length);
                    }
                } else {
                    // t.length < max_length
                    // possibly pad
                    if (padding) {
                        padHelper(
                            encodedTokens[i],
                            max_length,
                            (key) => (key === 'input_ids' ? this.pad_token_id : 0),
                            this.padding_side,
                        );
                    }
                }
            }
        }

        const result = {};

        if (return_tensor) {
            if (!(padding && truncation)) {
                // Not, guaranteed that all items have same length, so
                // we perform additional check

                if (
                    encodedTokens.some((x) => {
                        for (const key of Object.keys(x)) {
                            if (x[key].length !== encodedTokens[0][key]?.length) {
                                return true;
                            }
                        }
                        return false;
                    })
                ) {
                    throw Error(
                        'Unable to create tensor, you should probably activate truncation and/or padding ' +
                            "with 'padding=true' and 'truncation=true' to have batched tensors with the same length.",
                    );
                }
            }

            // Now we actually convert to tensor
            // NOTE: In the same way as the python library, we return a batched tensor, regardless of
            // whether we have a single input or multiple inputs.
            const dims = [encodedTokens.length, encodedTokens[0].input_ids.length];

            for (const key of Object.keys(encodedTokens[0])) {
                result[key] = new Tensor(
                    'int64',
                    BigInt64Array.from(encodedTokens.flatMap((x) => x[key]).map(BigInt)),
                    dims,
                );
            }
        } else {
            for (const key of Object.keys(encodedTokens[0])) {
                result[key] = encodedTokens.map((x) => x[key]);
            }

            // If not returning a tensor, we match the input type
            if (!isBatched) {
                // Input was not batched, so we unwrap
                for (const key of Object.keys(result)) {
                    result[key] = result[key][0];
                }
            }
        }

        return /** @type {BatchEncoding<BatchEncodingItem<TText, TReturnTensor>>} */ (result);
    }

    /**
     * Encodes a single text using the preprocessor pipeline of the tokenizer.
     *
     * @param {string|null} text The text to encode.
     * @returns {string[]|null} The encoded tokens.
     */
    _encode_text(text) {
        if (text === null) return null;
        return this._tokenizer.encode(text).tokens;
    }

    /**
     * Encodes a single text or a pair of texts using the model's tokenizer.
     *
     * @param {string} text The text to encode.
     * @param {Object} options An optional object containing the following properties:
     * @param {string|null} [options.text_pair=null] The optional second text to encode.
     * @param {boolean} [options.add_special_tokens=true] Whether or not to add the special tokens associated with the corresponding model.
     * @param {boolean|null} [options.return_token_type_ids=null] Whether to return token_type_ids.
     * @returns {{input_ids: number[], attention_mask: number[], token_type_ids?: number[]}} An object containing the encoded text.
     * @private
     */
    _encode_plus(text, { text_pair = null, add_special_tokens = true, return_token_type_ids = null } = {}) {
        const { ids, attention_mask, token_type_ids } = this._tokenizer.encode(text, {
            text_pair,
            add_special_tokens,
            return_token_type_ids: return_token_type_ids ?? this.return_token_type_ids,
        });
        return {
            input_ids: ids,
            attention_mask,
            ...(token_type_ids ? { token_type_ids } : {}),
        };
    }

    /**
     * Converts a string into a sequence of tokens.
     * @param {string} text The sequence to be encoded.
     * @param {Object} options An optional object containing the following properties:
     * @param {string|null} [options.pair] A second sequence to be encoded with the first.
     * @param {boolean} [options.add_special_tokens=false] Whether or not to add the special tokens associated with the corresponding model.
     * @returns {string[]} The list of tokens.
     */
    tokenize(text, { pair = null, add_special_tokens = false } = {}) {
        return this._tokenizer.tokenize(text, { text_pair: pair, add_special_tokens });
    }

    /**
     * Encodes a single text or a pair of texts using the model's tokenizer.
     *
     * @param {string} text The text to encode.
     * @param {Object} options An optional object containing the following properties:
     * @param {string|null} [options.text_pair=null] The optional second text to encode.
     * @param {boolean} [options.add_special_tokens=true] Whether or not to add the special tokens associated with the corresponding model.
     * @param {boolean|null} [options.return_token_type_ids=null] Whether to return token_type_ids.
     * @returns {number[]} An array of token IDs representing the encoded text(s).
     */
    encode(text, { text_pair = null, add_special_tokens = true, return_token_type_ids = null } = {}) {
        return this._tokenizer.encode(text, {
            text_pair,
            add_special_tokens,
            return_token_type_ids,
        }).ids;
    }

    /**
     * Decode a batch of tokenized sequences.
     * @param {number[][]|Tensor} batch List/Tensor of tokenized input sequences.
     * @param {Object} decode_args (Optional) Object with decoding arguments.
     * @returns {string[]} List of decoded sequences.
     */
    batch_decode(batch, decode_args = {}) {
        if (batch instanceof Tensor) {
            batch = batch.tolist();
        }
        return batch.map((x) => this.decode(x, decode_args));
    }

    /**
     * Decodes a sequence of token IDs back to a string.
     *
     * @param {number[]|bigint[]|Tensor} token_ids List/Tensor of token IDs to decode.
     * @param {Object} [decode_args={}]
     * @param {boolean} [decode_args.skip_special_tokens=false] If true, special tokens are removed from the output string.
     * @param {boolean} [decode_args.clean_up_tokenization_spaces=true] If true, spaces before punctuations and abbreviated forms are removed.
     *
     * @returns {string} The decoded string.
     * @throws {Error} If `token_ids` is not a non-empty array of integers.
     */
    decode(token_ids, decode_args = {}) {
        if (token_ids instanceof Tensor) {
            token_ids = prepareTensorForDecode(token_ids);
        }

        if (!Array.isArray(token_ids) || token_ids.length === 0 || !isIntegralNumber(token_ids[0])) {
            throw Error('token_ids must be a non-empty array of integers.');
        }

        return this.decode_single(token_ids, decode_args);
    }

    /**
     * Decode a single list of token ids to a string.
     * @param {number[]|bigint[]} token_ids List of token ids to decode
     * @param {Object} decode_args Optional arguments for decoding
     * @param {boolean} [decode_args.skip_special_tokens=false] Whether to skip special tokens during decoding
     * @param {boolean|null} [decode_args.clean_up_tokenization_spaces=null] Whether to clean up tokenization spaces during decoding.
     * If null, the value is set to `this.decoder.cleanup` if it exists, falling back to `this.clean_up_tokenization_spaces` if it exists, falling back to `true`.
     * @returns {string} The decoded string
     */
    decode_single(token_ids, { skip_special_tokens = false, clean_up_tokenization_spaces = null }) {
        return this._tokenizer.decode(token_ids, {
            skip_special_tokens,
            clean_up_tokenization_spaces,
        });
    }

    /**
     * Retrieve the chat template string used for tokenizing chat messages. This template is used
     * internally by the `apply_chat_template` method and can also be used externally to retrieve the model's chat
     * template for better generation tracking.
     *
     * @param {Object} options An optional object containing the following properties:
     * @param {string|null} [options.chat_template=null]
     * A Jinja template or the name of a template to use for this conversion.
     * It is usually not necessary to pass anything to this argument,
     * as the model's template will be used by default.
     * @param {Object[]} [options.tools=null]
     * A list of tools (callable functions) that will be accessible to the model. If the template does not
     * support function calling, this argument will have no effect. Each tool should be passed as a JSON Schema,
     * giving the name, description and argument types for the tool. See our
     * [chat templating guide](https://huggingface.co/docs/transformers/main/en/chat_templating#automated-function-conversion-for-tool-use)
     * for more information.
     * @returns {string} The chat template string.
     */
    get_chat_template({ chat_template = null, tools = null } = {}) {
        // First, handle the cases when the model has a dict of multiple templates
        if (this.chat_template && typeof this.chat_template === 'object') {
            const template_dict = this.chat_template;

            if (chat_template !== null && Object.hasOwn(template_dict, chat_template)) {
                // The user can pass the name of a template to the chat template argument instead of an entire template
                chat_template = template_dict[chat_template];
            } else if (chat_template === null) {
                if (tools !== null && 'tool_use' in template_dict) {
                    chat_template = template_dict['tool_use'];
                } else if ('default' in template_dict) {
                    chat_template = template_dict['default'];
                } else {
                    throw Error(
                        `This model has multiple chat templates with no default specified! Please either pass a chat ` +
                            `template or the name of the template you wish to use to the 'chat_template' argument. Available ` +
                            `template names are ${Object.keys(template_dict).sort()}.`,
                    );
                }
            }
        } else if (chat_template === null) {
            // These are the cases when the model has a single template
            // priority: `chat_template` argument > `tokenizer.chat_template`
            if (this.chat_template) {
                chat_template = this.chat_template;
            } else {
                throw Error(
                    'Cannot use apply_chat_template() because tokenizer.chat_template is not set and no template ' +
                        'argument was passed! For information about writing templates and setting the ' +
                        'tokenizer.chat_template attribute, please see the documentation at ' +
                        'https://huggingface.co/docs/transformers/main/en/chat_templating',
                );
            }
        }
        return chat_template;
    }

    /**
     * Converts a list of message objects with `"role"` and `"content"` keys to a list of token
     * ids. This method is intended for use with chat models, and will read the tokenizer's chat_template attribute to
     * determine the format and control tokens to use when converting.
     *
     * See [here](https://huggingface.co/docs/transformers/chat_templating) for more information.
     *
     * **Example:** Applying a chat template to a conversation.
     *
     * ```javascript
     * import { AutoTokenizer } from "@huggingface/transformers";
     *
     * const tokenizer = await AutoTokenizer.from_pretrained("Xenova/mistral-tokenizer-v1");
     *
     * const chat = [
     *   { "role": "user", "content": "Hello, how are you?" },
     *   { "role": "assistant", "content": "I'm doing great. How can I help you today?" },
     *   { "role": "user", "content": "I'd like to show off how chat templating works!" },
     * ]
     *
     * const text = tokenizer.apply_chat_template(chat, { tokenize: false });
     * // "<s>[INST] Hello, how are you? [/INST]I'm doing great. How can I help you today?</s> [INST] I'd like to show off how chat templating works! [/INST]"
     *
     * const input_ids = tokenizer.apply_chat_template(chat, { tokenize: true, return_tensor: false });
     * // [1, 733, 16289, 28793, 22557, 28725, 910, 460, 368, 28804, 733, 28748, 16289, 28793, 28737, 28742, 28719, 2548, 1598, 28723, 1602, 541, 315, 1316, 368, 3154, 28804, 2, 28705, 733, 16289, 28793, 315, 28742, 28715, 737, 298, 1347, 805, 910, 10706, 5752, 1077, 3791, 28808, 733, 28748, 16289, 28793]
     * ```
     *
     * @param {Message[]} conversation A list of message objects with `"role"` and `"content"` keys,
     * representing the chat history so far.
     * @template {boolean} [TTokenize=true]
     * @template {boolean} [TReturnTensor=true]
     * @template {boolean} [TReturnDict=true]
     * @param {Object} [options] An optional object containing the following properties:
     * @param {string|null} [options.chat_template=null] A Jinja template to use for this conversion. If
     * this is not passed, the model's chat template will be used instead.
     * @param {Object[]} [options.tools=null]
     * A list of tools (callable functions) that will be accessible to the model. If the template does not
     * support function calling, this argument will have no effect. Each tool should be passed as a JSON Schema,
     * giving the name, description and argument types for the tool. See our
     * [chat templating guide](https://huggingface.co/docs/transformers/main/en/chat_templating#automated-function-conversion-for-tool-use)
     * for more information.
     * @param {Record<string, string>[]} [options.documents=null]
     * A list of dicts representing documents that will be accessible to the model if it is performing RAG
     * (retrieval-augmented generation). If the template does not support RAG, this argument will have no
     * effect. We recommend that each document should be a dict containing "title" and "text" keys. Please
     * see the RAG section of the [chat templating guide](https://huggingface.co/docs/transformers/main/en/chat_templating#arguments-for-RAG)
     * for examples of passing documents with chat templates.
     * @param {boolean} [options.add_generation_prompt=false] Whether to end the prompt with the token(s) that indicate
     * the start of an assistant message. This is useful when you want to generate a response from the model.
     * Note that this argument will be passed to the chat template, and so it must be supported in the
     * template for this argument to have any effect.
     * @param {TTokenize} [options.tokenize=true] Whether to tokenize the output. If false, the output will be a string.
     * @param {boolean} [options.padding=false] Whether to pad sequences to the maximum length. Has no effect if tokenize is false.
     * @param {boolean} [options.truncation=false] Whether to truncate sequences to the maximum length. Has no effect if tokenize is false.
     * @param {number|null} [options.max_length=null] Maximum length (in tokens) to use for padding or truncation. Has no effect if tokenize is false.
     * If not specified, the tokenizer's `max_length` attribute will be used as a default.
     * @param {TReturnTensor} [options.return_tensor=true] Whether to return the output as a Tensor or an Array. Has no effect if tokenize is false.
     * @param {TReturnDict} [options.return_dict=true] Whether to return a dictionary with named outputs. Has no effect if tokenize is false.
     * @param {Object} [options.tokenizer_kwargs={}] Additional options to pass to the tokenizer.
     * @returns {ApplyChatTemplateReturn<TTokenize, TReturnTensor, TReturnDict>} The tokenized output.
     */
    apply_chat_template(
        conversation,
        options = /** @type {ApplyChatTemplateOptions<TTokenize, TReturnTensor, TReturnDict>} */ ({}),
    ) {
        let {
            tools = null,
            documents = null,
            chat_template = null,
            add_generation_prompt = false,
            tokenize = /** @type {TTokenize} */ (true),
            padding = false,
            truncation = false,
            max_length = null,
            return_tensor = /** @type {TReturnTensor} */ (true),
            return_dict = /** @type {TReturnDict} */ (true),
            tokenizer_kwargs = {},
            ...kwargs
        } = options;
        chat_template = this.get_chat_template({ chat_template, tools });

        if (typeof chat_template !== 'string') {
            throw Error(`chat_template must be a string, but got ${typeof chat_template}`);
        }

        // Compilation function uses a cache to avoid recompiling the same template
        /** @type {import('@huggingface/jinja').Template} */
        let compiledTemplate = this._compiled_template_cache.get(chat_template);
        if (compiledTemplate === undefined) {
            compiledTemplate = new Template(chat_template);
            this._compiled_template_cache.set(chat_template, compiledTemplate);
        }

        const special_tokens_map = Object.create(null);
        for (const key of SPECIAL_TOKEN_ATTRIBUTES) {
            const value = getTokenFromConfig(this.config, key);
            if (value) {
                special_tokens_map[key] = value;
            }
        }

        const rendered = compiledTemplate.render({
            messages: conversation,
            add_generation_prompt,
            tools,
            documents,
            ...special_tokens_map,
            ...kwargs,
        });

        if (tokenize) {
            const out = this._call(rendered, {
                add_special_tokens: false,
                padding,
                truncation,
                max_length,
                return_tensor,
                ...tokenizer_kwargs,
            });
            return /** @type {ApplyChatTemplateReturn<TTokenize, TReturnTensor, TReturnDict>} */ (
                return_dict ? out : out.input_ids
            );
        }

        return /** @type {ApplyChatTemplateReturn<TTokenize, TReturnTensor, TReturnDict>} */ (rendered);
    }
}

/**
 * Helper function to build translation inputs for an `NllbTokenizer` or `M2M100Tokenizer`.
 * @param {PreTrainedTokenizer} self The tokenizer instance.
 * @param {string|string[]} raw_inputs The text to tokenize.
 * @param {Object} tokenizer_options Options to be sent to the tokenizer
 * @param {Object} generate_kwargs Generation options.
 * @returns {Object} Object to be passed to the model.
 */
export function _build_translation_inputs(self, raw_inputs, tokenizer_options, generate_kwargs) {
    if (!('language_codes' in self) || !Array.isArray(self.language_codes)) {
        throw new Error(
            'Tokenizer must have `language_codes` attribute set and it should be an array of language ids.',
        );
    }
    if (!('languageRegex' in self) || !(self.languageRegex instanceof RegExp)) {
        throw new Error('Tokenizer must have `languageRegex` attribute set and it should be a regular expression.');
    }
    if (!('lang_to_token' in self) || typeof self.lang_to_token !== 'function') {
        throw new Error('Tokenizer must have `lang_to_token` attribute set and it should be a function.');
    }
    const src_lang_token = generate_kwargs.src_lang;
    const tgt_lang_token = generate_kwargs.tgt_lang;

    // Check that the target language is valid:
    if (!self.language_codes.includes(tgt_lang_token)) {
        throw new Error(
            `Target language code "${tgt_lang_token}" is not valid. Must be one of: {${self.language_codes.join(', ')}}`,
        );
    }

    // Allow `src_lang` to be optional. If not set, we'll use the tokenizer's default.
    if (src_lang_token !== undefined) {
        // Check that the source language is valid:
        if (!self.language_codes.includes(src_lang_token)) {
            throw new Error(
                `Source language code "${src_lang_token}" is not valid. Must be one of: {${self.language_codes.join(', ')}}`,
            );
        }

        // In the same way as the Python library, we override the post-processor
        // to force the source language to be first:
        for (const item of self._tokenizer.post_processor.config.single) {
            if ('SpecialToken' in item && self.languageRegex.test(item.SpecialToken.id)) {
                item.SpecialToken.id = self.lang_to_token(src_lang_token);
                break;
            }
        }
        // TODO: Do the same for pair?
    }

    // Override the `forced_bos_token_id` to force the correct language
    generate_kwargs.forced_bos_token_id = self._tokenizer.token_to_id(self.lang_to_token(tgt_lang_token));

    return self._call(raw_inputs, tokenizer_options);
}
