import { saveSettingsDebounced, getCurrentChatId, system_message_types, eventSource, event_types } from "../../../script.js";
import { humanizedDateTime } from "../../RossAscends-mods.js";
import { getApiUrl, extension_settings, getContext } from "../../extensions.js";
import { getFileText, onlyUnique, splitRecursive } from "../../utils.js";
export { MODULE_NAME };

const MODULE_NAME = 'chromadb';

const defaultSettings = {
    strategy: 'original',

    keep_context: 10,
    keep_context_min: 1,
    keep_context_max: 100,
    keep_context_step: 1,

    n_results: 20,
    n_results_min: 0,
    n_results_max: 100,
    n_results_step: 1,

    split_length: 384,
    split_length_min: 64,
    split_length_max: 4096,
    split_length_step: 64,

    file_split_length: 1024,
    file_split_length_min: 512,
    file_split_length_max: 4096,
    file_split_length_step: 128,
};

const postHeaders = {
    'Content-Type': 'application/json',
    'Bypass-Tunnel-Reminder': 'bypass',
};

const chatStateFlags = {};

function invalidateMessageSyncState(messageId) {
    console.log('CHROMADB: invalidating message sync state', messageId);
    const state = getChatSyncState();
    state[messageId] = false;
}

function getChatSyncState() {
    const currentChatId = getCurrentChatId();
    if (!checkChatId(currentChatId)) {
        return;
    }

    const context = getContext();
    const chatState = chatStateFlags[currentChatId] || [];
    chatState.length = context.chat.length;
    for (let i = 0; i < chatState.length; i++) {
        if (chatState[i] === undefined) {
            chatState[i] = false;
        }
    }
    chatStateFlags[currentChatId] = chatState;

    return chatState;
}

async function loadSettings() {
    if (Object.keys(extension_settings.chromadb).length === 0) {
        Object.assign(extension_settings.chromadb, defaultSettings);
    }

    console.log(`loading chromadb strat:${extension_settings.chromadb.strategy}`);
    $("#chromadb_strategy option[value=" + extension_settings.chromadb.strategy + "]").attr(
        "selected",
        "true"
    );
    $('#chromadb_keep_context').val(extension_settings.chromadb.keep_context).trigger('input');
    $('#chromadb_n_results').val(extension_settings.chromadb.n_results).trigger('input');
    $('#chromadb_split_length').val(extension_settings.chromadb.split_length).trigger('input');
    $('#chromadb_file_split_length').val(extension_settings.chromadb.file_split_length).trigger('input');
}

function onStrategyChange() {
    console.log('changing chromadb strat');
    extension_settings.chromadb.strategy = $('#chromadb_strategy').val();

    //$('#chromadb_strategy').select(extension_settings.chromadb.strategy);
    saveSettingsDebounced();
}

function onKeepContextInput() {
    extension_settings.chromadb.keep_context = Number($('#chromadb_keep_context').val());
    $('#chromadb_keep_context_value').text(extension_settings.chromadb.keep_context);
    saveSettingsDebounced();
}

function onNResultsInput() {
    extension_settings.chromadb.n_results = Number($('#chromadb_n_results').val());
    $('#chromadb_n_results_value').text(extension_settings.chromadb.n_results);
    saveSettingsDebounced();
}

function onSplitLengthInput() {
    extension_settings.chromadb.split_length = Number($('#chromadb_split_length').val());
    $('#chromadb_split_length_value').text(extension_settings.chromadb.split_length);
    saveSettingsDebounced();
}

function onFileSplitLengthInput() {
    extension_settings.chromadb.file_split_length = Number($('#chromadb_file_split_length').val());
    $('#chromadb_file_split_length_value').text(extension_settings.chromadb.file_split_length);
    saveSettingsDebounced();
}

function checkChatId(chat_id) {
    if (!chat_id || chat_id.trim() === '') {
        toastr.error('Please select a character and try again.');
        return false;
    }
    return true;
}

async function addMessages(chat_id, messages) {
    const url = new URL(getApiUrl());
    url.pathname = '/api/chromadb';

    const messagesDeepCopy = JSON.parse(JSON.stringify(messages));
    let splittedMessages = [];

    let id = 0;
    messagesDeepCopy.forEach((m, index) => {
        const split = splitRecursive(m.mes, extension_settings.chromadb.split_length);
        splittedMessages.push(...split.map(text => ({
            ...m,
            mes: text,
            send_date: id,
            id: `msg-${id++}`,
            index: index,
            extra: undefined,
        })));
    });

    splittedMessages = filterSyncedMessages(splittedMessages);

    // no messages to add
    if (splittedMessages.length === 0) {
        return { count: 0 };
    }

    const transformedMessages = splittedMessages.map((m) => ({
        id: m.id,
        role: m.is_user ? 'user' : 'assistant',
        content: m.mes,
        date: m.send_date,
        meta: JSON.stringify(m),
    }));

    const addMessagesResult = await fetch(url, {
        method: 'POST',
        headers: postHeaders,
        body: JSON.stringify({ chat_id, messages: transformedMessages }),
    });

    if (addMessagesResult.ok) {
        const addMessagesData = await addMessagesResult.json();

        return addMessagesData; // { count: 1 }
    }

    return { count: 0 };
}

function filterSyncedMessages(splittedMessages) {
    const syncState = getChatSyncState();
    const removeIndices = [];
    const syncedIndices = [];
    for (let i = 0; i < splittedMessages.length; i++) {
        const index = splittedMessages[i].index;

        if (syncState[index]) {
            removeIndices.push(i);
            continue;
        }

        syncedIndices.push(index);
    }

    for (const index of syncedIndices) {
        syncState[index] = true;
    }

    logSyncState(syncState);

    // remove messages that are already synced
    return splittedMessages.filter((_, i) => !removeIndices.includes(i));
}

function logSyncState(syncState) {
    const chat = getContext().chat;
    console.log('CHROMADB: sync state');
    console.table(syncState.map((v, i) => ({ synced: v, name: chat[i].name, message: chat[i].mes })));
}

async function onPurgeClick() {
    const chat_id = getCurrentChatId();
    if (!checkChatId(chat_id)) {
        return;
    }
    const url = new URL(getApiUrl());
    url.pathname = '/api/chromadb/purge';

    const purgeResult = await fetch(url, {
        method: 'POST',
        headers: postHeaders,
        body: JSON.stringify({ chat_id }),
    });

    if (purgeResult.ok) {
        delete chatStateFlags[chat_id];
        toastr.success('ChromaDB context has been successfully cleared');
    }
}

async function onExportClick() {
    const currentChatId = getCurrentChatId();
    if (!checkChatId(currentChatId)) {
        return;
    }
    const url = new URL(getApiUrl());
    url.pathname = '/api/chromadb/export';

    const exportResult = await fetch(url, {
        method: 'POST',
        headers: postHeaders,
        body: JSON.stringify({ currentChatId }),
    });

    if (exportResult.ok) {
        const data = await exportResult.json();
        const blob = new Blob([JSON.stringify(data, null, 2)], { type: 'application/json' });
        const href = URL.createObjectURL(blob);
        const link = document.createElement('a');
        link.href = href;
        link.download = currentChatId + '.json';
        document.body.appendChild(link);
        link.click();
        document.body.removeChild(link);
    } else {
        toastr.error('An error occurred while attempting to download the data');
    }
}

async function onSelectImportFile(e) {
    const file = e.target.files[0];
    const currentChatId = getCurrentChatId();
    if (!checkChatId(currentChatId)) {
        return;
    }

    if (!file) {
        return;
    }

    try {
        toastr.info('This may take some time, depending on the file size', 'Processing...');

        const text = await getFileText(file);
        const imported = JSON.parse(text);

        imported.chat_id = currentChatId;

        const url = new URL(getApiUrl());
        url.pathname = '/api/chromadb/import';

        const importResult = await fetch(url, {
            method: 'POST',
            headers: postHeaders,
            body: JSON.stringify(imported),
        });

        if (importResult.ok) {
            const importResultData = await importResult.json();

            toastr.success(`Number of chunks: ${importResultData.count}`, 'Injected successfully!');
            return importResultData;
        } else {
            throw new Error();
        }
    }
    catch (error) {
        console.log(error);
        toastr.error('Something went wrong while importing the data');
    }
    finally {
        e.target.form.reset();
    }
}

async function queryMessages(chat_id, query) {
    const url = new URL(getApiUrl());
    url.pathname = '/api/chromadb/query';

    const queryMessagesResult = await fetch(url, {
        method: 'POST',
        headers: postHeaders,
        body: JSON.stringify({ chat_id, query, n_results: extension_settings.chromadb.n_results }),
    });

    if (queryMessagesResult.ok) {
        const queryMessagesData = await queryMessagesResult.json();

        return queryMessagesData;
    }

    return [];
}

async function onSelectInjectFile(e) {
    const file = e.target.files[0];
    const currentChatId = getCurrentChatId();
    if (!checkChatId(currentChatId)) {
        return;
    }
    if (!file) {
        return;
    }

    try {
        toastr.info('This may take some time, depending on the file size', 'Processing...');
        const text = await getFileText(file);

        const split = splitRecursive(text, extension_settings.chromadb.file_split_length).filter(onlyUnique);

        const messages = split.map(m => ({
            id: `${file.name}-${split.indexOf(m)}`,
            role: 'system',
            content: m,
            date: Date.now(),
            meta: JSON.stringify({
                name: file.name,
                is_user: false,
                is_name: false,
                is_system: false,
                send_date: humanizedDateTime(),
                mes: m,
                extra: {
                    type: system_message_types.NARRATOR,
                }
            }),
        }));

        const url = new URL(getApiUrl());
        url.pathname = '/api/chromadb';

        const addMessagesResult = await fetch(url, {
            method: 'POST',
            headers: postHeaders,
            body: JSON.stringify({ chat_id: currentChatId, messages: messages }),
        });

        if (addMessagesResult.ok) {
            const addMessagesData = await addMessagesResult.json();

            toastr.success(`Number of chunks: ${addMessagesData.count}`, 'Injected successfully!');
            return addMessagesData;
        } else {
            throw new Error();
        }
    }
    catch (error) {
        console.log(error);
        toastr.error('Something went wrong while injecting the data');
    }
    finally {
        e.target.form.reset();
    }
}

window.chromadb_interceptGeneration = async (chat) => {
    const currentChatId = getCurrentChatId();
    const selectedStrategy = extension_settings.chromadb.strategy;
    if (currentChatId) {
        const messagesToStore = chat.slice(0, -extension_settings.chromadb.keep_context);

        if (messagesToStore.length > 0) {
            await addMessages(currentChatId, messagesToStore);

            const lastMessage = chat[chat.length - 1];

            if (lastMessage) {
                const queriedMessages = await queryMessages(currentChatId, lastMessage.mes);

                queriedMessages.sort((a, b) => a.date - b.date);

                const newChat = [];

                if (selectedStrategy === 'ross') {
                    //adds chroma to the end of chat and allows Generate() to cull old messages naturally.
                    const context = getContext();
                    const charname = context.name2;
                    newChat.push(
                        {
                            is_name: false,
                            is_user: false,
                            mes: `[Use these past chat exchanges to inform ${charname}'s next response:`,
                            name: "system",
                            send_date: 0,
                        }
                    );
                    newChat.push(...queriedMessages.map(m => JSON.parse(m.meta)));
                    newChat.push(
                        {
                            is_name: false,
                            is_user: false,
                            mes: `]\n`,
                            name: "system",
                            send_date: 0,
                        }
                    );
                    chat.splice(chat.length, 0, ...newChat);
                }

                if (selectedStrategy === 'original') {
                    //removes .length # messages from the start of 'kept messages'
                    //replaces them with chromaDB results (with no separator)
                    newChat.push(...queriedMessages.map(m => JSON.parse(m.meta)));
                    chat.splice(0, messagesToStore.length, ...newChat);

                }
                console.log('ChromaDB chat after injection', chat);
            }
        }
    }
}

jQuery(async () => {
    const settingsHtml = `
    <div class="chromadb_settings">
        <div class="inline-drawer">
            <div class="inline-drawer-toggle inline-drawer-header">
            <b>Smart Context</b>
            <div class="inline-drawer-icon fa-solid fa-circle-chevron-down down"></div>
        </div>
        <div class="inline-drawer-content">
            <p>This extension rearranges the messages in the current chat to keep more relevant information in the context. Adjust the sliders below based on average amount of messages in your prompt (refer to the chat cut-off line).</p>
            <span>Memory Injection Strategy</span>
            <select id="chromadb_strategy">
                <option value="original">Replace non-kept chat items with memories</option>
                <option value="ross">Add memories after chat with a header tag</option>
            </select>
            <label for="chromadb_keep_context">How many original chat messages to keep: (<span id="chromadb_keep_context_value"></span>) messages</label>
            <input id="chromadb_keep_context" type="range" min="${defaultSettings.keep_context_min}" max="${defaultSettings.keep_context_max}" step="${defaultSettings.keep_context_step}" value="${defaultSettings.keep_context}" />
            <label for="chromadb_n_results">Maximum number of ChromaDB 'memories' to inject: (<span id="chromadb_n_results_value"></span>) messages</label>
            <input id="chromadb_n_results" type="range" min="${defaultSettings.n_results_min}" max="${defaultSettings.n_results_max}" step="${defaultSettings.n_results_step}" value="${defaultSettings.n_results}" />
            <label for="chromadb_split_length">Max length for each 'memory' pulled from the current chat history: (<span id="chromadb_split_length_value"></span>) characters</label>
            <input id="chromadb_split_length" type="range" min="${defaultSettings.split_length_min}" max="${defaultSettings.split_length_max}" step="${defaultSettings.split_length_step}" value="${defaultSettings.split_length}" />
            <label for="chromadb_file_split_length">Max length for each 'memory' pulled from imported text files: (<span id="chromadb_file_split_length_value"></span>) characters</label>
            <input id="chromadb_file_split_length" type="range" min="${defaultSettings.file_split_length_min}" max="${defaultSettings.file_split_length_max}" step="${defaultSettings.file_split_length_step}" value="${defaultSettings.file_split_length}" />
            <div class="flex-container spaceEvenly">
                <div id="chromadb_inject" title="Upload custom textual data to use in the context of the current chat" class="menu_button">
                    <i class="fa-solid fa-file-arrow-up"></i>
                    <span>Inject Data (TXT file)</span>
                </div>
                <div id="chromadb_export" title="Export all of the current chromadb data for this current chat" class="menu_button">
                    <i class="fa-solid fa-file-export"></i>
                    <span>Export</span>
                </div>
                <div id="chromadb_import" title="Import a full chromadb export for this current chat" class="menu_button">
                    <i class="fa-solid fa-file-import"></i>
                    <span>Import</span>
                </div>
                <div id="chromadb_purge" title="Force purge all the data related to the current chat from the database" class="menu_button">
                    <i class="fa-solid fa-broom"></i>
                    <span>Purge Chat from the DB</span>
                </div>
            </div>
            <small><i>Since ChromaDB state is not persisted to disk by default, you'll need to inject text data every time the Extras API server is restarted.</i></small>
        </div>
        <form><input id="chromadb_inject_file" type="file" accept="text/plain" hidden></form>
        <form><input id="chromadb_import_file" type="file" accept="application/json" hidden></form>
    </div>`;

    $('#extensions_settings').append(settingsHtml);
    $('#chromadb_strategy').on('change', onStrategyChange);
    $('#chromadb_keep_context').on('input', onKeepContextInput);
    $('#chromadb_n_results').on('input', onNResultsInput);
    $('#chromadb_split_length').on('input', onSplitLengthInput);
    $('#chromadb_file_split_length').on('input', onFileSplitLengthInput);
    $('#chromadb_inject').on('click', () => $('#chromadb_inject_file').trigger('click'));
    $('#chromadb_import').on('click', () => $('#chromadb_import_file').trigger('click'));
    $('#chromadb_inject_file').on('change', onSelectInjectFile);
    $('#chromadb_import_file').on('change', onSelectImportFile);
    $('#chromadb_purge').on('click', onPurgeClick);
    $('#chromadb_export').on('click', onExportClick);
    await loadSettings();

    // Not sure if this is needed, but it's here just in case
    eventSource.on(event_types.MESSAGE_DELETED, getChatSyncState);
    eventSource.on(event_types.MESSAGE_RECEIVED, getChatSyncState);
    eventSource.on(event_types.MESSAGE_SENT, getChatSyncState);
    // Will make the sync state update when a message is edited or swiped
    eventSource.on(event_types.MESSAGE_EDITED, invalidateMessageSyncState);
    eventSource.on(event_types.MESSAGE_SWIPED, invalidateMessageSyncState);
});