// ==UserScript==
// @name         Perplexity Model Selection
// @namespace    https://greasyfork.org/en/users/688917
// @version      0.3
// @description  Adds model selection buttons to Perplexity AI using jQuery
// @author       dpgc, lyh16, mall-fluffy-bongo
// @match        https://www.perplexity.ai/*
// @license      MIT
// @run-at       document-end
// @downloadURL none
// ==/UserScript==

(function() {
    'use strict';

    // Check if jQuery is loaded on the page
    if (typeof jQuery === 'undefined') {
        var script = document.createElement('script');
        script.src = 'https://ajax.googleapis.com/ajax/libs/jquery/3.5.1/jquery.min.js';
        script.type = 'text/javascript';
        document.getElementsByTagName('head')[0].appendChild(script);

        script.onload = function() {
            setup();
        };
    } else {
        setup();
    }

    function createModelSelectorElement(buttonText) {
        var $button = $('<button/>', {
            type: 'button',
            class: 'model-selector md:hover:bg-offsetPlus text-textOff dark:text-textOffDark md:hover:text-textMain dark:md:hover:bg-offsetPlusDark dark:md:hover:text-textMainDark font-sans focus:outline-none outline-none outline-transparent transition duration-300 ease-in-out font-sans select-none items-center relative group/button justify-center text-center items-center rounded-full cursor-point active:scale-95 origin-center whitespace-nowrap inline-flex text-sm px-2 font-medium h-8'
        });

        const $svg = $(`
        <svg aria-hidden="true" focusable="false" data-prefix="far" data-icon="bars-filter" class="svg-inline--fa fa-bars-filter fa-fw fa-1x mr-1" role="img" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 448 512"><path fill="currentColor" d="M0 88C0 74.7 10.7 64 24 64H424c13.3 0 24 10.7 24 24s-10.7 24-24 24H24C10.7 112 0 101.3 0 88zM64 248c0-13.3 10.7-24 24-24H360c13.3 0 24 10.7 24 24s-10.7 24-24 24H88c-13.3 0-24-10.7-24-24zM288 408c0 13.3-10.7 24-24 24H184c-13.3 0-24-10.7-24-24s10.7-24 24-24h80c13.3 0 24 10.7 24 24z"></path></svg>
        `)
        var $textDiv = $(`<div class="model-selector-text text-align-center relative truncate">${buttonText}</div>`);
        var $buttonContentDiv = $('<div/>', {
            class: 'flex items-center leading-none justify-center gap-1',
        }).append($svg).append($textDiv);

        $button.append($buttonContentDiv);
        var $wrapperDiv = $('<div class="model-selector-wrapper mr-2"/>').append($('<span/>').append($button));

        return {
            $element: $wrapperDiv,
            setModelName: (modelName) => {
                // $textDiv.text(`${buttonText} (${modelName})`);
                $textDiv.text(`${modelName} `);
            }
        }
    }

    function createSelectionPopover(sourceElement) {
        const createSelectionElement = (input) => {
            const {name, onClick} = input;
            const $element = $(`
            <div class="md:h-full">
                <div class="md:h-full">
                    <div class="relative cursor-pointer md:hover:bg-offsetPlus py-md px-sm md:p-sm rounded md:hover:dark:bg-offsetPlusDark transition-all duration-300 md:h-full -ml-sm md:ml-0 select-none rounded">
                        <div class="flex items-center justify-between relative">
                            <div class="flex items-center gap-x-xs default font-sans text-sm font-medium text-textMain dark:text-textMainDark selection:bg-superDuper selection:text-textMain">
                                <span>${name}</span>
                            </div>
                        </div>
                    </div>
                </div>
            </div>
            `);

            $element.click(onClick);
            return $element;
        }

        const popoverHTML = `<div class="flex justify-center items-center">
            <div class="ease-in-out duration-150 transition">
                <div class="absolute left-0 top-0 z-30">
                    <div data-tag="popper" data-popper-reference-hidden="false" data-popper-escaped="false" data-popper-placement="bottom-end" style="position: absolute; inset: 0px 0px auto auto;">
                        <div class="border animate-in ease-in-out fade-in zoom-in-95 duration-150 rounded shadow-sm p-xs border-borderMain/50 ring-borderMain/50 divide-borderMain/50 dark:divide-borderMainDark/50 dark:ring-borderMainDark/50 dark:border-borderMainDark/50 bg-background dark:bg-backgroundDark">
                            <div data-tag="menu" class="min-w-[160px] max-w-[250px] border-borderMain/50 ring-borderMain/50 divide-borderMain/50 dark:divide-borderMainDark/50 dark:ring-borderMainDark/50 dark:border-borderMainDark/50 bg-transparent">
                                <!-- Put elements here -->
                            </div>
                        </div>
                    </div>
                </div>
            </div>
        </div>`;

        const $popover = $(popoverHTML);
        const $popper = $popover.find('[data-tag="popper"]');
        const $menuContaienr = $popover.find('[data-tag="menu"]');

        if (sourceElement) {
            const {top, left, width, height} = sourceElement.getBoundingClientRect();
            const offset = 10;
            const popperWidth = $popper.outerWidth();
            $popper.css('transform', `translate(${left + (width + popperWidth * 2)}px, ${top + height + offset}px)`);
        }

        return {
            $element: $popover,
            addSelection: (input) => {
                const $selection = createSelectionElement(input);
                $menuContaienr.append($selection);
            }
        }

    }

    async function fetchSettings() {
        const url = 'https://www.perplexity.ai/p/api/v1/user/settings';
        const response = await fetch(url);
        if (!response.ok) throw new Error('Failed to fetch settings');
        return await response.json();
    }

    function setupSelection() {
        let selector = '';
        // const currentURL = window.location.href;
        // if (currentURL === 'https://www.perplexity.ai/') {
        //     selector = '.flex.bg-background.dark\\:bg-offsetDark.rounded-l-lg.col-start-1.row-start-2.-ml-2';
        // } else if (currentURL.startsWith('https://www.perplexity.ai/search/')) {
        //     selector = '.pointer-events-none.fixed.z-10.grid-cols-12.gap-xl.px-sm.py-sm.md\\:bottom-lg.md\\:grid.md\\:px-0.bottom-\\[64px\\].border-borderMain\\/50.ring-borderMain\\/50.divide-borderMain\\/50.dark\\:divide-borderMainDark\\/50.dark\\:ring-borderMainDark\\/50.dark\\:border-borderMainDark\\/50.bg-transparent';
        // } else {
        //     return;
        // }
        selector = '.flex.bg-background.dark\\:bg-offsetDark.rounded-l-lg.col-start-1.row-start-2.-ml-2';

        const $focusElement = $('div:contains("Focus")').closest(selector);
        if (!$focusElement.length) return;

        if ($focusElement.data('state') === 'injected') return;
        $focusElement.data('state', 'injected');

        const aiModels = [
            {
                "name": "Default",
                "code": "turbo"
            },
            {
                "name": "Experimental",
                "code": "experimental"
            },
            {
                "name": "GPT-4 Turbo",
                "code": "gpt4"
            },
            {
                "name": "Claude 3 Opus",
                "code": "claude3opus"
            },
            {
                "name": "Claude 3 Sonnet",
                "code": "claude2"
            },
            {
                "name": "Mistral Large",
                "code": "mistral"
            }
        ];

        const imageModels = [
            {
                "name": "Playground v.2.5",
                "code": "default"
            },
            {
                "name": "DALL-E 3",
                "code": "dall-e-3"
            },
            {
                "name": "Stable Diffusion XL",
                "code": "sdxl"
            }
        ];

        const aiModelSelector = createModelSelectorElement("Chat Model");
        const imageModelSelector = createModelSelectorElement("Image Model");

        let latestSettings = undefined;
        const getCurrentModel = () => {
            return latestSettings?.["default_model"];
        }
        const getCurrentImageModel = () => {
            return latestSettings?.["default_image_generation_model"];
        }
        const updateFromSettings = () => {
            fetchSettings().then((settings) => {
                latestSettings = settings;
                const aiModelCode = getCurrentModel();
                const aiModelName = aiModels.find(m => m.code === aiModelCode)?.name;
                if (aiModelName) aiModelSelector.setModelName(aiModelName);

                const imageModelCode = getCurrentImageModel();
                const imageModelName = imageModels.find(m => m.code === imageModelCode)?.name;
                if (imageModelName) imageModelSelector.setModelName(imageModelName);
            });
        };
        updateFromSettings();

        const setModel = async (model, isImageModel) => {
            const el = $focusElement[0];
            const fiberKey = Object.keys(el).find(k => k.startsWith('__reactFiber'));
            if (!fiberKey) throw new Error('Failed to find key of React Fiber');
            const fiber = el[fiberKey];
            const settingsKey = isImageModel ? 'default_image_generation_model' : 'default_model';
            return await fiber.child.sibling.memoizedProps.socket.emitWithAck('save_user_settings', {
                [settingsKey]: model,
                source: "default",
                version: "2.5"
            })
        }

        aiModelSelector.$element.click(async () => {
            const {$element: $popover, addSelection} = createSelectionPopover(aiModelSelector.$element[0]);
            $('main').append($popover);
            const closePopover = () => {
                $popover.remove();
                $(document).off('click', closePopover);
            }
            for (const model of aiModels) {
                addSelection({
                    name: model.name,
                    onClick: async () => {
                        await setModel(model.code, false);
                        updateFromSettings();
                        closePopover();
                    }
                });
            }

            setTimeout(() => {
                $(document).on('click', closePopover);
                $popover.on('click', (e) => e.stopPropagation());
            }, 500);
        });

        imageModelSelector.$element.click(async () => {
            const {$element: $popover, addSelection} = createSelectionPopover(imageModelSelector.$element[0]);
            $('main').append($popover);
            const closePopover = () => {
                $popover.remove();
                $(document).off('click', closePopover);
            }
            for (const model of imageModels) {
                addSelection({
                    name: model.name,
                    onClick: async () => {
                        await setModel(model.code, true);
                        updateFromSettings();
                        closePopover();
                    }
                });
            }

            setTimeout(() => {
                $(document).on('click', closePopover);
                $popover.on('click', (e) => e.stopPropagation());
            }, 500);
        });

        $focusElement.append(aiModelSelector.$element);
        $focusElement.append(imageModelSelector.$element);

        // Add CSS styles for responsive layout
        $('<style>')
            .prop('type', 'text/css')
            .html(`
                .model-selector-wrapper {
                    margin-right: 12px; /* Add right margin to create space between buttons */
                }
                @media (max-width: 768px) {
                    .model-selector-wrapper {
                        display: block;
                        margin-right: 0;
                        margin-bottom: 8px;
                    }
                    .model-selector {
                        width: 100%;
                    }
                    .model-selector-text {
                        max-width: 120px;
                    }
                }
            `)
            .appendTo('head');
    }

    function setup() {
        setupSelection();
        setInterval(() => {
            setupSelection();
            console.log('run');
        }, 500);
    }
})();