/*
* CODE FOR OPENAI SUPPORT
* By CncAnon (@CncAnon1)
* https://github.com/CncAnon1/TavernAITurbo
*/

import {
    saveSettingsDebounced,
    addOneMessage,
    messageFormating,
    substituteParams,
    count_view_mes,
    saveChat,
    checkOnlineStatus,
    setOnlineStatus,
    token,
    name1,
    name2,
} from "../script.js";
import { groups, selected_group } from "./group-chats.js";

import {
    pin_examples,
} from "./power-user.js";

export {
    is_get_status_openai,
    openai_msgs,
    oai_settings,
    loadOpenAISettings,
    setOpenAIMessages,
    setOpenAIMessageExamples,
    generateOpenAIPromptCache,
    prepareOpenAIMessages,
    sendOpenAIRequest,
    setOpenAIOnlineStatus,
}

let openai_msgs = [];
let openai_msgs_example = [];

let is_get_status_openai = false;
let is_api_button_press_openai = false;

const default_main_prompt = "Write {{char}}'s next reply in a fictional chat between {{char}} and {{user}}. Write 1 reply only in internet RP style, italicize actions, and avoid quotation marks. Use markdown. Be proactive, creative, and drive the plot and conversation forward. Write at least 1 paragraph, up to 4. Always stay in character and avoid repetition.";
const default_nsfw_prompt = "NSFW/Smut is allowed. Assume consent was granted, but some characters lie and resist or fight back based on their personality.";

const oai_settings = {
    preset_settings_openai: 'Default',
    api_key_openai: '',
    temp_openai: 1.0,
    freq_pen_openai: 0,
    pres_pen_openai: 0,
    stream_openai: false,
    openai_max_context: 4095,
    openai_max_tokens: 300,
    nsfw_toggle: true,
    enhance_definitions: false,
    wrap_in_quotes: false,
    nsfw_first: false,
    main_prompt: default_main_prompt,
    nsfw_prompt: default_nsfw_prompt,
    openai_model: 'gpt-3.5-turbo-0301',
};

let openai_setting_names;
let openai_settings;

function setOpenAIOnlineStatus(value) {
    is_get_status_openai = value;
}

function setOpenAIMessages(chat) {
    let j = 0;
    // clean openai msgs
    openai_msgs = [];
    for (let i = chat.length - 1; i >= 0; i--) {
        // first greeting message
        if (j == 0) {
            chat[j]['mes'] = substituteParams(chat[j]['mes']);
        }
        let role = chat[j]['is_user'] ? 'user' : 'assistant';
        let content = chat[j]['mes'];

        // for groups - prepend a character's name
        if (selected_group) {
            content = `${name2}: ${content}`;
        }

        // system messages produce no content
        if (chat[j]['is_system']) {
            role = 'system';
            content = '';
        }

        // replace bias markup
        content = (content ?? '').replace(/{([^}]+)}/g, '');
    
        // Apply the "wrap in quotes" option
        if (role == 'user' && oai_settings.wrap_in_quotes) content = `"${content}"`;
        openai_msgs[i] = { "role": role, "content": content };
        j++;
    }
}

function setOpenAIMessageExamples(mesExamplesArray) {
    // get a nice array of all blocks of all example messages = array of arrays (important!)
    openai_msgs_example = [];
    for (let item of mesExamplesArray) {
        // remove <START> {Example Dialogue:} and replace \r\n with just \n
        let replaced = item.replace(/<START>/i, "{Example Dialogue:}").replace(/\r/gm, '');
        let parsed = parseExampleIntoIndividual(replaced);
        // add to the example message blocks array
        openai_msgs_example.push(parsed);
    }
}

function generateOpenAIPromptCache(charPersonality, topAnchorDepth, anchorTop, anchorBottom) {
    openai_msgs = openai_msgs.reverse();
    let is_add_personality = false;
    openai_msgs.forEach(function (msg, i, arr) {//For added anchors and others
        let item = msg["content"];
        if (i === openai_msgs.length - topAnchorDepth && count_view_mes >= topAnchorDepth && !is_add_personality) {
            is_add_personality = true;
            if ((anchorTop != "" || charPersonality != "")) {
                if (anchorTop != "") charPersonality += ' ';
                item = `[${name2} is ${charPersonality}${anchorTop}]\n${item}`;
            }
        }
        if (i >= openai_msgs.length - 1 && count_view_mes > 8 && $.trim(item).substr(0, (name1 + ":").length) == name1 + ":") {//For add anchor in end
            item = anchorBottom + "\n" + item;
        }
        msg["content"] = item;
        openai_msgs[i] = msg;
    });
}

function parseExampleIntoIndividual(messageExampleString) {
    let result = []; // array of msgs
    let tmp = messageExampleString.split("\n");
    let cur_msg_lines = [];
    let in_user = false;
    let in_bot = false;
    // DRY my cock and balls
    function add_msg(name, role) {
        // join different newlines (we split them by \n and join by \n)
        // remove char name
        // strip to remove extra spaces
        let parsed_msg = cur_msg_lines.join("\n").replace(name + ":", "").trim();

        if (selected_group && role == 'assistant') {
            parsed_msg = `${name}: ${parsed_msg}`;
        }

        result.push({ "role": role, "content": parsed_msg });
        cur_msg_lines = [];
    }
    // skip first line as it'll always be "This is how {bot name} should talk"
    for (let i = 1; i < tmp.length; i++) {
        let cur_str = tmp[i];
        // if it's the user message, switch into user mode and out of bot mode
        // yes, repeated code, but I don't care
        if (cur_str.indexOf(name1 + ":") === 0) {
            in_user = true;
            // we were in the bot mode previously, add the message
            if (in_bot) {
                add_msg(name2, "assistant");
            }
            in_bot = false;
        } else if (cur_str.indexOf(name2 + ":") === 0) {
            in_bot = true;
            // we were in the user mode previously, add the message
            if (in_user) {
                add_msg(name1, "user");
            }
            in_user = false;
        }
        // push the current line into the current message array only after checking for presence of user/bot
        cur_msg_lines.push(cur_str);
    }
    // Special case for last message in a block because we don't have a new message to trigger the switch
    if (in_user) {
        add_msg(name1, "user");
    } else if (in_bot) {
        add_msg(name2, "assistant");
    }
    return result;
}

function formatWorldInfo(value) {
    if (!value) {
        return '';
    }

    // placeholder if we would want to apply some formatting
    return `[Details of the fictional world the RP set in:\n${value}\n]`;
}

function prepareOpenAIMessages(name2, storyString, worldInfoBefore, worldInfoAfter, extensionPrompt, bias) {
    let this_max_context = oai_settings.openai_max_context;
    let nsfw_toggle_prompt = "";
    let enhance_definitions_prompt = "";

    if (oai_settings.nsfw_toggle) {
        nsfw_toggle_prompt = oai_settings.nsfw_prompt;
    } else {
        nsfw_toggle_prompt = "Avoid writing a NSFW/Smut reply. Creatively write around it NSFW/Smut scenarios in character.";
    }
    // Experimental but kinda works
    if (oai_settings.enhance_definitions) {
        enhance_definitions_prompt = "If you have more knowledge of " + name2 + ", add to the character's lore and personality to enhance them but keep the Character Sheet's definitions absolute.";
    }

    let whole_prompt = [];
    // If it's toggled, NSFW prompt goes first.
    if (oai_settings.nsfw_first) {
        whole_prompt = [nsfw_toggle_prompt, oai_settings.main_prompt, enhance_definitions_prompt, "\n\n", formatWorldInfo(worldInfoBefore), storyString, formatWorldInfo(worldInfoAfter), extensionPrompt]
    }
    else {
        whole_prompt = [oai_settings.main_prompt, nsfw_toggle_prompt, enhance_definitions_prompt, "\n\n", formatWorldInfo(worldInfoBefore), storyString, formatWorldInfo(worldInfoAfter), extensionPrompt]
    }

    // Join by a space and replace placeholders with real user/char names
    storyString = substituteParams(whole_prompt.join(" "))

    let prompt_msg = { "role": "system", "content": storyString }
    let examples_tosend = [];
    let openai_msgs_tosend = [];

    // todo: static value, maybe include in the initial context calculation
    let new_chat_msg = { "role": "system", "content": "[Start a new chat]" };
    let start_chat_count = countTokens([new_chat_msg]);
    let total_count = countTokens([prompt_msg], true) + start_chat_count;

    if (bias) {
        let bias_msg = { "role": "system", "content": bias };
        openai_msgs.push(bias_msg);
    }
    
    if (selected_group) {
        const groupMembers = groups.find(x => x.id === selected_group)?.members;
        const names = Array.isArray(groupMembers) ? groupMembers.join(', ') : '';
        new_chat_msg.content = `[Start a new group chat. Group members: ${names}]`;
        let group_nudge = { "role": "system", "content": `[Write the next reply only as ${name2}]` };
        openai_msgs.push(group_nudge);
    }

    // The user wants to always have all example messages in the context
    if (pin_examples) {
        // first we send *all* example messages
        // we don't check their token size since if it's bigger than the context, the user is fucked anyway
        // and should've have selected that option (maybe have some warning idk, too hard to add)
        for (const element of openai_msgs_example) {
            // get the current example block with multiple user/bot messages
            let example_block = element;
            // add the first message from the user to tell the model that it's a new dialogue
            // TODO: instead of role user content use role system name example_user
            // message from the user so the model doesn't confuse the context (maybe, I just think that this should be done)
            if (example_block.length != 0) {
                examples_tosend.push(new_chat_msg);
            }
            for (const example of example_block) {
                // add all the messages from the example
                examples_tosend.push(example);
            }
        }
        total_count += countTokens(examples_tosend);
        // go from newest message to oldest, because we want to delete the older ones from the context
        for (let j = openai_msgs.length - 1; j >= 0; j--) {
            let item = openai_msgs[j];
            let item_count = countTokens(item);
            // If we have enough space for this message, also account for the max assistant reply size
            if ((total_count + item_count) < (this_max_context - oai_settings.openai_max_tokens)) {
                openai_msgs_tosend.push(item);
                total_count += item_count;
            }
            else {
                // early break since if we still have more messages, they just won't fit anyway
                break;
            }
        }
    } else {
        for (let j = openai_msgs.length - 1; j >= 0; j--) {
            let item = openai_msgs[j];
            let item_count = countTokens(item);
            // If we have enough space for this message, also account for the max assistant reply size
            if ((total_count + item_count) < (this_max_context - oai_settings.openai_max_tokens)) {
                openai_msgs_tosend.push(item);
                total_count += item_count;
            }
            else {
                // early break since if we still have more messages, they just won't fit anyway
                break;
            }
        }

        console.log(total_count);

        for (const example of openai_msgs_example) {
            // get the current example block with multiple user/bot messages
            let example_block = example;

            for (let k = 0; k < example_block.length; k++) {
                if (example_block.length == 0) { continue; }
                let example_count = countTokens(example_block[k]);
                // add all the messages from the example
                if ((total_count + example_count + start_chat_count) < (this_max_context - oai_settings.openai_max_tokens)) {
                    if (k == 0) {
                        examples_tosend.push(new_chat_msg);
                        total_count += start_chat_count;
                    }
                    examples_tosend.push(example_block[k]);
                    total_count += example_count;
                }
                else { break; }
            }
        }
    }
    // reverse the messages array because we had the newest at the top to remove the oldest,
    // now we want proper order
    openai_msgs_tosend.reverse();
    openai_msgs_tosend = [prompt_msg, ...examples_tosend, new_chat_msg, ...openai_msgs_tosend]

    console.log("We're sending this:")
    console.log(openai_msgs_tosend);
    console.log(`Calculated the total context to be ${total_count} tokens`);
    return openai_msgs_tosend;
}

async function sendOpenAIRequest(openai_msgs_tosend) {
    const generate_data = {
        "messages": openai_msgs_tosend,
        "model": oai_settings.openai_model,
        "temperature": parseFloat(oai_settings.temp_openai),
        "frequency_penalty": parseFloat(oai_settings.freq_pen_openai),
        "presence_penalty": parseFloat(oai_settings.pres_pen_openai),
        "max_tokens": oai_settings.openai_max_tokens,
        "stream": false, //oai_settings.stream_openai,
    };

    const generate_url = '/generate_openai';
    // TODO: fix streaming
    const streaming = oai_settings.stream_openai;
    const last_view_mes = count_view_mes;

    const response = await fetch(generate_url, {
        method: 'POST',
        body: JSON.stringify(generate_data),
        headers: {
            'Content-Type': 'application/json',
            "X-CSRF-Token": token,
        }
    });

    const data = await response.json();

    if (data.error) {
        throw new Error(data);
    }

    return data.choices[0]["message"]["content"];
}

// Unused
function onStream(e, resolve, reject, last_view_mes) {
    let end = false;
    if (!oai_settings.stream_openai)
        return;
    let response = e.currentTarget.response;
    if (response == "{\"error\":true}") {
        reject('', 'error');
    }

    let eventList = response.split("\n");
    let getMessage = "";
    for (let event of eventList) {
        if (!event.startsWith("data"))
            continue;
        if (event == "data: [DONE]") {
            chat[chat.length - 1]['mes'] = getMessage;
            $("#send_but").css("display", "block");
            $("#loading_mes").css("display", "none");
            saveChat();
            end = true;
            break;
        }
        let data = JSON.parse(event.substring(6));
        // the first and last messages are undefined, protect against that
        getMessage += data.choices[0]["delta"]["content"] || "";
    }

    if ($("#chat").children().filter(`[mesid="${last_view_mes}"]`).length == 0) {
        chat[chat.length] = {};
        chat[chat.length - 1]['name'] = name2;
        chat[chat.length - 1]['is_user'] = false;
        chat[chat.length - 1]['is_name'] = false;
        chat[chat.length - 1]['send_date'] = Date.now();
        chat[chat.length - 1]['mes'] = "";
        addOneMessage(chat[chat.length - 1]);
    }

    let messageText = messageFormating($.trim(getMessage), name1);
    $("#chat").children().filter(`[mesid="${last_view_mes}"]`).children('.mes_block').children('.mes_text').html(messageText);

    let $textchat = $('#chat');
    $textchat.scrollTop($textchat[0].scrollHeight);

    if (end) {
        resolve();
    }
}

function countTokens(messages, full = false) {
    if (!Array.isArray(messages)) {
        messages = [messages];
    }
    let token_count = -1;
    jQuery.ajax({
        async: false,
        type: 'POST', // 
        url: `/tokenize_openai?model=${oai_settings.openai_model}`,
        data: JSON.stringify(messages),
        dataType: "json",
        contentType: "application/json",
        success: function (data) {
            token_count = data.token_count;
        }
    });
    if (!full) token_count -= 2;
    return token_count;
}

function loadOpenAISettings(data, settings) {
    if (settings.api_key_openai != undefined) {
        oai_settings.api_key_openai = settings.api_key_openai;
        $("#api_key_openai").val(oai_settings.api_key_openai);
    }

    openai_setting_names = data.openai_setting_names;
    openai_settings = data.openai_settings;
    openai_settings = data.openai_settings;
    openai_settings.forEach(function (item, i, arr) {
        openai_settings[i] = JSON.parse(item);
    });

    $("#settings_perset_openai").empty();
    let arr_holder = {};
    openai_setting_names.forEach(function (item, i, arr) {
        arr_holder[item] = i;
        $('#settings_perset_openai').append(`<option value=${i}>${item}</option>`);

    });
    openai_setting_names = arr_holder;

    oai_settings.preset_settings_openai = settings.preset_settings_openai;
    $(`#settings_perset_openai option[value=${openai_setting_names[oai_settings.preset_settings_openai]}]`).attr('selected', true);

    oai_settings.temp_openai = settings.temp_openai ?? 0.9;
    oai_settings.freq_pen_openai = settings.freq_pen_openai ?? 0.7;
    oai_settings.pres_pen_openai = settings.pres_pen_openai ?? 0.7;
    oai_settings.stream_openai = settings.stream_openai ?? true;
    oai_settings.openai_max_context = settings.openai_max_context ?? 4095;
    oai_settings.openai_max_tokens = settings.openai_max_tokens ?? 300;

    if (settings.nsfw_toggle !== undefined) oai_settings.nsfw_toggle = !!settings.nsfw_toggle;
    if (settings.keep_example_dialogue !== undefined) oai_settings.keep_example_dialogue = !!settings.keep_example_dialogue;
    if (settings.enhance_definitions !== undefined) oai_settings.enhance_definitions = !!settings.enhance_definitions;
    if (settings.wrap_in_quotes !== undefined) oai_settings.wrap_in_quotes = !!settings.wrap_in_quotes;
    if (settings.nsfw_first !== undefined) oai_settings.nsfw_first = !!settings.nsfw_first;
    if (settings.openai_model !== undefined) oai_settings.openai_model = settings.openai_model;

    $('#stream_toggle').prop('checked', oai_settings.stream_openai);

    $('#openai_max_context').val(oai_settings.openai_max_context);
    $('#openai_max_context_counter').html(`${oai_settings.openai_max_context} Tokens`);

    $('#openai_max_tokens').val(oai_settings.openai_max_tokens);
    $(`#model_openai_select option[value="${oai_settings.openai_model}"`).attr('selected', true);

    $('#nsfw_toggle').prop('checked', oai_settings.nsfw_toggle);
    $('#keep_example_dialogue').prop('checked', oai_settings.keep_example_dialogue);
    $('#enhance_definitions').prop('checked', oai_settings.enhance_definitions);
    $('#wrap_in_quotes').prop('checked', oai_settings.wrap_in_quotes);
    $('#nsfw_first').prop('checked', oai_settings.nsfw_first);

    if (settings.main_prompt !== undefined) oai_settings.main_prompt = settings.main_prompt;
    if (settings.nsfw_prompt !== undefined) oai_settings.nsfw_prompt = settings.nsfw_prompt;
    $('#main_prompt_textarea').val(oai_settings.main_prompt);
    $('#nsfw_prompt_textarea').val(oai_settings.nsfw_prompt);

    $('#temp_openai').val(oai_settings.temp_openai);
    $('#temp_counter_openai').text(Number(oai_settings.temp_openai).toFixed(2));

    $('#freq_pen_openai').val(oai_settings.freq_pen_openai);
    $('#freq_pen_counter_openai').text(Number(oai_settings.freq_pen_openai).toFixed(2));

    $('#pres_pen_openai').val(oai_settings.pres_pen_openai);
    $('#pres_pen_counter_openai').text(Number(oai_settings.pres_pen_openai).toFixed(2));
}

async function getStatusOpen() {
    if (is_get_status_openai) {
        let data = { key: oai_settings.api_key_openai };

        jQuery.ajax({
            type: 'POST', // 
            url: '/getstatus_openai', // 
            data: JSON.stringify(data),
            beforeSend: function () { },
            cache: false,
            dataType: "json",
            contentType: "application/json",
            success: function (data) {
                if (!('error' in data))
                    setOnlineStatus('Valid');
                resultCheckStatusOpen();
            },
            error: function (jqXHR, exception) {
                setOnlineStatus('no_connection');
                console.log(exception);
                console.log(jqXHR);
                resultCheckStatusOpen();
            }
        });
    } else {
        setOnlineStatus('no_connection');
    }
}

function resultCheckStatusOpen() {
    is_api_button_press_openai = false;
    checkOnlineStatus();
    $("#api_loading_openai").css("display", 'none');
    $("#api_button_openai").css("display", 'inline-block');
}

$(document).ready(function () {
    $(document).on('input', '#temp_openai', function () {
        oai_settings.temp_openai = $(this).val();
        $('#temp_counter_openai').text(Number($(this).val()).toFixed(2));
        saveSettingsDebounced();
    });

    $(document).on('input', '#freq_pen_openai', function () {
        oai_settings.freq_pen_openai = $(this).val();
        $('#freq_pen_counter_openai').text(Number($(this).val()).toFixed(2));
        saveSettingsDebounced();
    });

    $(document).on('input', '#pres_pen_openai', function () {
        oai_settings.pres_pen_openai = $(this).val();
        $('#pres_pen_counter_openai').text(Number($(this).val()));
        saveSettingsDebounced();

    });

    $(document).on('input', '#openai_max_context', function () {
        oai_settings.openai_max_context = parseInt($(this).val());
        $('#openai_max_context_counter').text(`${$(this).val()} Tokens`);
        saveSettingsDebounced();
    });

    $(document).on('input', '#openai_max_tokens', function () {
        oai_settings.openai_max_tokens = parseInt($(this).val());
        saveSettingsDebounced();
    });

    $("#model_openai_select").change(function() {
        const value = $(this).val();
        oai_settings.openai_model = value;
        saveSettingsDebounced();
    });

    $('#stream_toggle').change(function () {
        oai_settings.stream_openai = !!$('#stream_toggle').prop('checked');
        saveSettingsDebounced();
    });

    $('#nsfw_toggle').change(function () {
        oai_settings.nsfw_toggle = !!$('#nsfw_toggle').prop('checked');
        saveSettingsDebounced();
    });

    $('#enhance_definitions').change(function () {
        oai_settings.enhance_definitions = !!$('#enhance_definitions').prop('checked');
        saveSettingsDebounced();
    });

    $('#wrap_in_quotes').change(function () {
        oai_settings.wrap_in_quotes = !!$('#wrap_in_quotes').prop('checked');
        saveSettingsDebounced();
    });

    $('#nsfw_first').change(function () {
        oai_settings.nsfw_first = !!$('#nsfw_first').prop('checked');
        saveSettingsDebounced();
    });

    $("#settings_perset_openai").change(function () {
        oai_settings.preset_settings_openai = $('#settings_perset_openai').find(":selected").text();

        const preset = openai_settings[openai_setting_names[preset_settings_openai]];
        oai_settings.temp_openai = preset.temperature;
        oai_settings.freq_pen_openai = preset.frequency_penalty;
        oai_settings.pres_pen_openai = preset.presence_penalty;

        // probably not needed
        $('#temp_counter_openai').text(oai_settings.temp_openai);
        $('#freq_pen_counter_openai').text(oai_settings.freq_pen_openai);
        $('#pres_pen_counter_openai').text(oai_settings.pres_pen_openai);

        $('#temp_openai').val(oai_settings.temp_openai).trigger('input');
        $('#freq_pen_openai').val(oai_settings.freq_pen_openai).trigger('input');
        $('#pres_pen_openai').val(oai_settings.pres_pen_openai).trigger('input');

        saveSettingsDebounced();
    });

    $("#api_button_openai").click(function (e) {
        e.stopPropagation();
        if ($('#api_key_openai').val() != '') {
            $("#api_loading_openai").css("display", 'inline-block');
            $("#api_button_openai").css("display", 'none');
            oai_settings.api_key_openai = $.trim($('#api_key_openai').val());
            saveSettingsDebounced();
            is_get_status_openai = true;
            is_api_button_press_openai = true;
            getStatusOpen();
        }
    });
});