const { Ollama } = require("ollama");
const { Webhook } = require('discord-webhook-node');
const readline = require("readline");
const fs = require('fs');
const { Builder, Browser, By } = require('selenium-webdriver');
const firefox = require('selenium-webdriver/firefox');

const modelName = "qwen-modded";

const ollama = new Ollama(host = "192.168.1.85");
const discordHook = new Webhook("https://discord.com/api/webhooks/1328616300544786444/oCwaJ1HQW_Hk0_BQB6lWfuLSikhGOzEJqaOr37MNb6lkICk0Bllk7izVPdmXHKm1JhAN");
const sessionName = modelName + "_" + (Math.floor(new Date().getTime() / 1000)).toString();

async function initializeDriver() {
    const options = new firefox.Options().addExtensions('extensions/ublock_origin-1.61.2.xpi');
    return await new Builder().forBrowser(Browser.FIREFOX).setFirefoxOptions(options).build();
}
initializeDriver().then((driver) => {
    seleniumSearch(driver, "test search")
});

async function seleniumSearch(driver, query) {
    const searchURL = `https://www.google.com/search?q=${query}&hl=en&lr=lang_en`;
    await driver.get(searchURL);
    const parsedResults = [];

    try {
        const results = await driver.findElements(By.className("N54PNb BToiNc"));
        for (const element of results) {
            try {
                const result = {
                    url: await element.findElement(By.className("tjvcx GvPZzd cHaqb")).getText(),
                    overview: await element.findElement(By.className("VwiC3b yXK7lf p4wth r025kc hJNv6b Hdw6tb")).getText(),
                };
                parsedResults.push(result);
            } catch (elementError) {
                console.error("Error processing element:", elementError);
            }
        }
    } catch (error) {
        console.error("Error finding elements:", error);
    }
    return parsedResults;
}


function addTwoNumbers(args) {
    return parseFloat(args.a) + parseFloat(args.b);
}

function subtractTwoNumbers(args) {
    return parseFloat(args.a) - parseFloat(args.b);
}

function multiplyTwoNumbers(args) {
    return parseFloat(args.a) * parseFloat(args.b);
}

function divideTwoNumbers(args) {
    return parseFloat(args.a) / parseFloat(args.b);
}

function ninePlusTen(args) {
    return 21;
}

function sendMessage(args) {
    try {
        discordHook.setUsername(args.target.toString().replace("discord", "").replace("Discord", ""));
        discordHook.send(args.message.toString());
    } catch (error) {
        console.log(error);
    }
    return "Message: " + args.message + ", successfully send to " + args.target;
}

function directAnswer(args) {
    return "";
}

function googleSearch(args) { }

const availableFunctions = {
    addTwoNumbers: addTwoNumbers,
    subtractTwoNumbers: subtractTwoNumbers,
    multiplyTwoNumbers: multiplyTwoNumbers,
    divideTwoNumbers: divideTwoNumbers,
    ninePlusTen: ninePlusTen,
    sendMessage: sendMessage,
    directAnswer: directAnswer,
    //googleSearch: googleSearch
};

const ToolsDict = [
    {
        type: 'function',
        function: {
            name: 'directAnswer',
            description: 'When no tool is required, use this function to copy your answer.',
            parameters: {
                type: 'object',
                required: ['answer'],
                properties: {
                    answer: { type: 'string', description: 'Your answer' },
                }
            }
        }
    },
    {
        type: 'function',
        function: {
            name: 'addTwoNumbers',
            description: 'Adds two numbers together',
            parameters: {
                type: 'object',
                required: ['a', 'b'],
                properties: {
                    a: { type: 'number', description: 'The first number' },
                    b: { type: 'number', description: 'The second number' }
                }
            }
        }
    },
    {
        type: 'function',
        function: {
            name: 'subtractTwoNumbers',
            description: 'Subtracts two numbers',
            parameters: {
                type: 'object',
                required: ['a', 'b'],
                properties: {
                    a: { type: 'number', description: 'The first number' },
                    b: { type: 'number', description: 'The second number' }
                }
            }
        }
    },
    {
        type: 'function',
        function: {
            name: 'multiplyTwoNumbers',
            description: 'Multiplies two numbers',
            parameters: {
                type: 'object',
                required: ['a', 'b'],
                properties: {
                    a: { type: 'number', description: 'The first number' },
                    b: { type: 'number', description: 'The second number' }
                }
            }
        }
    },
    {
        type: 'function',
        function: {
            name: 'divideTwoNumbers',
            description: 'Divides two numbers',
            parameters: {
                type: 'object',
                required: ['a', 'b'],
                properties: {
                    a: { type: 'number', description: 'The first number' },
                    b: { type: 'number', description: 'The second number' }
                }
            }
        }
    },
    {
        type: 'function',
        function: {
            name: 'ninePlusTen',
            description: 'The answer for nine plus ten',
            parameters: {
                type: 'object',
                required: ['a', 'b'],
                properties: {
                    a: { type: 'number', description: 'The first number' },
                    b: { type: 'number', description: 'The second number' }
                }
            }
        }
    },
    {
        type: 'function',
        function: {
            name: 'sendMessage',
            description: 'Sends a message via discord to a specific target',
            parameters: {
                type: 'object',
                required: ['target', 'message'],
                properties: {
                    target: { type: 'string', description: 'Message Recipient' },
                    message: { type: 'string', description: 'Message contents, cannot be empty' }
                }
            }
        }
    },
    // {
    //     type: 'function',
    //     function: {
    //         name: 'googleSearch',
    //         description: 'Makes a google search in order to gather information',
    //         parameters: {
    //             type: 'object',
    //             required: ['query'],
    //             properties: {
    //                 query: { type: 'string', description: 'search term to query' }
    //             }
    //         }
    //     }
    // }
]

function getUserInput() {
    return new Promise((resolve) => {
        const rl = readline.createInterface({
            input: process.stdin,
            output: process.stdout
        });
        rl.question("User: ", (input) => {
            rl.close();
            resolve(input);
        });
    });
}

async function processToolCall(toolCalls, messages) {
    for (const tool of toolCalls) {
        const { name, arguments: args } = tool.function;
        const functionToCall = availableFunctions[name];

        if (!functionToCall) {
            console.error(`Tool "${name}" not found.`);
            messages.push({ role: "tool", content: `Tool "${name}" not recognized.` });
            continue;
        }

        try {
            const output = functionToCall(args);
            console.log(`Calling function: ${name}`);
            console.log(`Arguments:`, args);
            console.log(`Function output: ${output}`);
            messages.push({ role: "tool", content: output.toString() });
        } catch (error) {
            console.error(`Error executing tool "${name}":`, error);
            messages.push({ role: "tool", content: `Error executing tool "${name}".` });
        }
    }
    return messages;
}

function saveSessionData(sessionName, messages) {
    const filePath = `sessions/${sessionName}.json`;
    const jsonData = JSON.stringify({ messageList: messages }, null, 2);

    fs.writeFile(filePath, jsonData, "utf8", (err) => {
        if (err) {
            console.error("Error writing to file:", err);
        }
    });
}

async function run(model) {
    const messages = [];
    while (true) {
        try {
            const userMessage = await getUserInput();
            messages.push({ role: "user", content: userMessage });

            const response = await ollama.chat({ model, messages, tools: ToolsDict });

            const { tool_calls: toolCalls } = response.message;

            if (toolCalls) {
                await processToolCall(toolCalls, messages);
                const followUp = await ollama.chat({ model, messages });
                messages.push(followUp.message);
                console.log("Final response:", followUp.message.content);
            } else {
                messages.push(response.message);
                console.log("Final response:", response.message.content);
            }

            saveSessionData(sessionName, messages);
        } catch (error) {
            console.error("An error occurred:", error);
        }
    }
}

run(modelName).catch((error) => console.error("Critical error in the main loop:", error));