import {
    animation_duration,
    callPopup,
    chat,
    cleanUpMessage,
    event_types,
    eventSource,
    Generate,
    getGeneratingApi,
    is_send_press,
    isStreamingEnabled,
} from '../script.js';
import { debounce, delay, getStringHash } from './utils.js';
import { decodeTextTokens, getTokenizerBestMatch } from './tokenizers.js';
import { power_user } from './power-user.js';

const TINTS = 4;
const MAX_MESSAGE_LOGPROBS = 100;

/**
 * Tuple of a candidate token and its logarithm of probability of being chosen
 * @typedef {[string, number]} Candidate - (token, logprob)
 */

/**
 * Logprob data for a single message
 * @typedef {Object} MessageLogprobData
 * @property {number} created - timestamp of when the message was generated
 * @property {number} hash - hash of the message object
 * @property {number} messageId - ID of the source message
 * @property {number} swipeId - ID of the source swipe on the source message
 * @property {string} api - API used to generate the message
 * @property {TokenLogprobs[]} messageLogprobs Logprob data for each token, by
 * its index in the message
 * @property {string | null} continueFrom - the 'continue' prefix used to
 * generate the message, if any
 */

/**
 * Logprob data for a single token
 * @typedef {Object} TokenLogprobs
 * @property {string} token - A token generated by the model
 * @property {Candidate[]} topLogprobs - Array of top candidate tokens
 */

let state = {
    /** @type {TokenLogprobs | null} */
    selectedTokenLogprobs: null,
    /** @type {Map<number, MessageLogprobData>} */
    messageLogprobs: new Map(),
};

/**
 * renderAlternativeTokensView renders the Token Probabilities UI and all
 * subviews with the active message's logprobs data. If the message has no token
 * logprobs, a zero-state is rendered.
 */
function renderAlternativeTokensView() {
    const view = $('#logprobs_generation_output');
    if (!view.is(':visible')) {
        return;
    }
    view.empty();
    state.selectedTokenLogprobs = null;
    renderTopLogprobs();

    const { messageLogprobs, continueFrom } = getActiveMessageLogprobData() || {};
    const usingSmoothStreaming = isStreamingEnabled() && power_user.smooth_streaming;
    if (!messageLogprobs?.length || usingSmoothStreaming) {
        const emptyState = $('<div></div>');
        const noTokensMsg = usingSmoothStreaming
            ? 'Token probabilities are not available when using Smooth Streaming.'
            : 'No token probabilities available for the current message.';
        const msg = power_user.request_token_probabilities
            ? noTokensMsg
            : '<span>Enable <b>Request token probabilities</b> in the User Settings menu to use this feature.</span>';
        emptyState.html(msg);
        emptyState.addClass('logprobs_empty_state');
        view.append(emptyState);
        return;
    }

    const prefix = continueFrom || '';
    const tokenSpans = [];

    if (prefix) {
        const prefixSpan = $('<span></span>');
        prefixSpan.text(prefix);
        prefixSpan.html(prefixSpan.html().replace(/\n/g, '<br>'));
        prefixSpan.addClass('logprobs_output_prefix');
        prefixSpan.attr('title', 'Select to reroll the last \'Continue\' generation');
        prefixSpan.click(onPrefixClicked);
        addKeyboardProps(prefixSpan);
        tokenSpans.push(...withVirtualWhitespace(prefix, prefixSpan));
    }

    messageLogprobs.forEach((tokenData, i) => {
        const { token } = tokenData;
        const span = $('<span></span>');
        const text = toVisibleWhitespace(token);
        span.text(text);
        span.addClass('logprobs_output_token');
        span.addClass('logprobs_tint_' + (i % TINTS));
        span.click(() => onSelectedTokenChanged(tokenData, span));
        addKeyboardProps(span);
        tokenSpans.push(...withVirtualWhitespace(token, span));
    });

    view.append(tokenSpans);

    // scroll past long prior context
    if (prefix) {
        view.find('.logprobs_output_token').first()[0].scrollIntoView();
    }
}

function addKeyboardProps(element) {
    element.attr('role', 'button');
    element.attr('tabindex', '0');
    element.keydown(function (e) {
        if (e.key === 'Enter' || e.key === ' ') {
            element.click();
        }
    });
}

/**
 * renderTopLogprobs renders the top logprobs subview with the currently
 * selected token highlighted. If no token is selected, the subview is hidden.
 */
function renderTopLogprobs() {
    const view = $('.logprobs_candidate_list');
    const hint = $('#logprobs_top_logprobs_hint').hide();
    view.empty();

    if (!state.selectedTokenLogprobs) {
        return;
    }

    const { token: selectedToken, topLogprobs } = state.selectedTokenLogprobs;

    let sum = 0;
    const nodes = [];
    const candidates = topLogprobs
        .sort(([, logA], [, logB]) => logB - logA)
        .map(([text, log]) => {
            if (log <= 0) {
                const probability = Math.exp(log);
                sum += probability;
                return [text, probability, log];
            }
            else {
                return [text, log, null];
            }
        });
    candidates.push(['<others>', 1 - sum, 0]);

    let matched = false;
    for (const [token, probability, log] of candidates) {
        const container = $('<button class="flex-container flexFlowColumn logprobs_top_candidate"></button>');
        const tokenNormalized = String(token).replace(/^▁/g, ' ');

        if (token === selectedToken || tokenNormalized === selectedToken) {
            matched = true;
            container.addClass('selected');
        }

        const tokenText = $('<span></span>').text(`${toVisibleWhitespace(token)}`);
        const percentText = $('<span></span>').text(`${(probability * 100).toFixed(2)}%`);
        container.append(tokenText, percentText);
        if (log) {
            container.attr('title', `logarithm: ${log}`);
        }
        addKeyboardProps(container);
        if (token !== '<others>') {
            container.click(() => onAlternativeClicked(state.selectedTokenLogprobs, token));
        } else {
            container.prop('disabled', true);
        }
        nodes.push(container);
    }

    // Highlight the <others> node if the selected token was not included in the
    // top logprobs
    if (!matched) {
        nodes[nodes.length - 1].css('background-color', 'rgba(255, 0, 0, 0.1)');
    }

    view.append(nodes);
}

/**
 * onSelectedTokenChanged is called when the user clicks on a token in the
 * token output view. It updates the selected token state and re-renders the
 * top logprobs view, or deselects the token if it was already selected.
 * @param {TokenLogprobs} logprobs - logprob data for the selected token
 * @param {Element} span - target span node that was clicked
 */
function onSelectedTokenChanged(logprobs, span) {
    $('.logprobs_output_token.selected').removeClass('selected');
    if (state.selectedTokenLogprobs === logprobs) {
        state.selectedTokenLogprobs = null;
    } else {
        state.selectedTokenLogprobs = logprobs;
        $(span).addClass('selected');
    }
    renderTopLogprobs();
}

/**
 * onAlternativeClicked is called when the user clicks on an alternative token
 * in the top logprobs view. It will create a new swipe message and prefill it
 * with all text up to the selected token, followed by the chosen alternative.
 * Then it requests a `continue` completion from the model with the new prompt.
 * @param {TokenLogprobs} tokenLogprobs - logprob data for selected alternative
 * @param {string} alternative - selected alternative token's text
 */
function onAlternativeClicked(tokenLogprobs, alternative) {
    if (!checkGenerateReady()) {
        return;
    }

    if (getGeneratingApi() === 'openai') {
        return callPopup(`<h3>Feature unavailable</h3><p>Due to API limitations, rerolling a token is not supported with OpenAI. Try switching to a different API.</p>`, 'text');
    }

    const { messageLogprobs, continueFrom } = getActiveMessageLogprobData();
    const replaceIndex = messageLogprobs.findIndex(x => x === tokenLogprobs);

    const tokens = messageLogprobs.slice(0, replaceIndex + 1).map(({ token }) => token);
    tokens[replaceIndex] = String(alternative).replace(/^▁/g, ' ');

    const prefix = continueFrom || '';
    const prompt = prefix + tokens.join('');
    const messageId = chat.length - 1;
    createSwipe(messageId, prompt);

    $('.swipe_right:last').click(); // :see_no_evil:

    Generate('continue').then(_ => void _);
}

/**
 * onPrefixClicked is called when the user clicks on the carried-over prefix
 * in the token output view. It allows them to reroll the last 'continue'
 * completion with none of the output generated from it, in case they don't
 * like the results.
 */
function onPrefixClicked() {
    if (!checkGenerateReady()) {
        return;
    }

    const { continueFrom } = getActiveMessageLogprobData();
    const messageId = chat.length - 1;
    const prefix = continueFrom || '';
    createSwipe(messageId, prefix);
    $('.swipe_right:last').click();
    Generate('continue').then(_ => void _);
}

function checkGenerateReady() {
    if (is_send_press) {
        toastr.warning(`Please wait for the current generation to complete.`);
        return false;
    }
    return true;
}


/**
 * onToggleLogprobsPanel is called when the user performs an action that toggles
 * the logprobs view, such as clicking the Token Probabilities menu item or the
 * close button.
 */
function onToggleLogprobsPanel() {
    const logprobsViewer = $('#logprobsViewer');

    // largely copied from CFGScale toggle
    if (logprobsViewer.css('display') === 'none') {
        logprobsViewer.addClass('resizing');
        logprobsViewer.css('display', 'flex');
        logprobsViewer.css('opacity', 0.0);
        renderAlternativeTokensView();
        logprobsViewer.transition({
            opacity: 1.0,
            duration: animation_duration,
        }, async function () {
            await delay(50);
            logprobsViewer.removeClass('resizing');
        });
    } else {
        logprobsViewer.addClass('resizing');
        logprobsViewer.transition({
                opacity: 0.0,
                duration: animation_duration,
            },
            async function () {
                await delay(50);
                logprobsViewer.removeClass('resizing');
            });
        setTimeout(function () {
            logprobsViewer.hide();
        }, animation_duration);
    }
}

/**
 * createSwipe appends a new swipe to the target chat message with the given
 * text.
 * @param {number} messageId - target chat message ID
 * @param {string} prompt - initial prompt text which will be continued
 */
function createSwipe(messageId, prompt) {
    // need to call `cleanUpMessage` on our new prompt, because we were working
    // with raw model output and our new prompt is missing trimming/macro replacements
    const cleanedPrompt = cleanUpMessage(prompt, false, false);

    const msg = chat[messageId];
    const newSwipeInfo = {
        send_date: msg.send_date,
        gen_started: msg.gen_started,
        gen_finished: msg.gen_finished,
        extra: { ...structuredClone(msg.extra), from_logprobs: new Date().getTime() },
    };

    msg.swipes = msg.swipes || [];
    msg.swipe_info = msg.swipe_info || [];

    // Add our new swipe, then make sure the active swipe is the one just before
    // it. The call to `swipe_right` will switch to it immediately.
    msg.swipes.push(cleanedPrompt);
    msg.swipe_info.push(newSwipeInfo);
    msg.swipe_id = Math.max(0, msg.swipes.length - 2);
}

/**
 * toVisibleWhitespace receives input text and replaces spaces with &middot; and
 * newlines with ↵.
 * @param {string} input
 * @returns {string}
 */
function toVisibleWhitespace(input) {
    return input.replace(/ /g, '·').replace(/▁/g, '·').replace(/\n/g, '↵');
}

/**
 * withVirtualWhitespace inserts line breaks and a zero-width space before and
 * after the span node if its token begins or ends with whitespace in order to
 * allow text to wrap despite whitespace characters being replaced with a dot.
 * @param {string} text - token text being evaluated for whitespace
 * @param {Element} span - target span node to be wrapped
 * @returns {Element[]} array of nodes to be appended to the DOM
 */
function withVirtualWhitespace(text, span) {
    const result = [span];
    if (text.match(/^\s/)) {
        result.unshift(document.createTextNode('\u200b'));
    }
    if (text.match(/\s$/)) {
        result.push($(document.createTextNode('\u200b')));
    }
    if (text.match(/^▁/)) {
        result.unshift(document.createTextNode('\u200b'));
    }
    // line breaks are trickier. we don't currently handle consecutive line
    // breaks or line breaks occuring in between non-whitespace characters, but
    // tokenizers generally don't produce those anyway.

    // matches leading line break, at least one character, and trailing line break
    if (text.match(/^\n(?:.|\n)+\n$/)) {
        result.unshift($('<br>'));
        result.push($('<br>'));
    } else if (text.match(/^\n/)) {
        result.unshift($('<br>'));
    } else if (text.match(/\n$/)) {
        result.push($('<br>'));
    }
    return result;
}

/**
 * saveLogprobsForActiveMessage receives an array of TokenLogprobs objects
 * representing the top logprobs for each token in a message and associates it
 * with the active message.
 *
 * **Ensure the active message has been updated and rendered before calling
 * this function or the logprobs data will be saved to the wrong message.**
 * @param {TokenLogprobs[]} logprobs - array of logprobs data for each token
 * @param {string | null} continueFrom  - for 'continue' generations, the prompt
 */
export function saveLogprobsForActiveMessage(logprobs, continueFrom) {
    if (!logprobs) {
        // non-streaming APIs could return null data
        return;
    }

    convertTokenIdLogprobsToText(logprobs);

    const msgId = chat.length - 1;
    /** @type {MessageLogprobData} */
    const data = {
        created: new Date().getTime(),
        api: getGeneratingApi(),
        messageId: msgId,
        swipeId: chat[msgId].swipe_id,
        messageLogprobs: logprobs,
        continueFrom,
        hash: getMessageHash(chat[msgId]),
    }

    state.messageLogprobs.set(data.hash, data);

    // Clean up old logprobs data
    const oldLogprobs = Array.from(state.messageLogprobs.values())
        .sort((a, b) => b.created - a.created)
        .slice(MAX_MESSAGE_LOGPROBS);
    for (const oldData of oldLogprobs) {
        state.messageLogprobs.delete(oldData.hash);
    }
}

function getMessageHash(message) {
    // We don't use the swipe ID as a hash component because it's not stable,
    // deleting a swipe will change the ID of all subsequent swipes.
    const hashParams = {
        name: message.name,
        mid: chat.indexOf(message),
        text: message.mes,
    };
    return getStringHash(JSON.stringify(hashParams));
}

/**
 * getActiveMessageLogprobData returns the logprobs data for the active chat
 * message.
 * @returns {MessageLogprobData || null}
 */
function getActiveMessageLogprobData() {
    const hash = getMessageHash(chat[chat.length - 1]);
    return state.messageLogprobs.get(hash) || null;
}

/**
 * convertLogprobTokenIdsToText mutates the given logprobs data's topLogprobs
 * field keyed by token text instead of token ID. This is only necessary for
 * APIs which only return token IDs in their logprobs data; for others this
 * function is a no-op.
 * @param {TokenLogprobs[]} input - logprobs data with numeric token IDs
 */
function convertTokenIdLogprobsToText(input) {
    const api = getGeneratingApi();
    if (api !== 'novel') {
        return input;
    }

    const tokenizerId = getTokenizerBestMatch(api);

    // Flatten unique token IDs across all logprobs
    const tokenIds = Array.from(new Set(input.flatMap(logprobs =>
        logprobs.topLogprobs.map(([token]) => token).concat(logprobs.token)
    )));

    // Submit token IDs to tokenizer to get token text, then build ID->text map
    const { chunks } = decodeTextTokens(tokenizerId, tokenIds);
    const tokenIdText = new Map(tokenIds.map((id, i) => [id, chunks[i]]));

    // Fixup logprobs data with token text
    input.forEach(logprobs => {
        logprobs.token = tokenIdText.get(logprobs.token);
        logprobs.topLogprobs = logprobs.topLogprobs.map(([token, logprob]) =>
            [tokenIdText.get(token), logprob]
        );
    });
}

export function initLogprobs() {
    const debouncedRender = debounce(renderAlternativeTokensView, 500);
    $('#logprobsViewerClose').click(onToggleLogprobsPanel);
    $('#option_toggle_logprobs').click(onToggleLogprobsPanel);
    eventSource.on(event_types.CHAT_CHANGED, debouncedRender);
    eventSource.on(event_types.CHARACTER_MESSAGE_RENDERED, debouncedRender);
    eventSource.on(event_types.IMPERSONATE_READY, debouncedRender);
    eventSource.on(event_types.MESSAGE_DELETED, debouncedRender);
    eventSource.on(event_types.MESSAGE_EDITED, debouncedRender);
    eventSource.on(event_types.MESSAGE_SWIPED, debouncedRender);
}