// ==UserScript==
// @name         Perplexity Model Selection
// @namespace    https://greasyfork.org/en/users/688917
// @version      0.11
// @description  Adds model selection buttons to Perplexity AI using jQuery
// @author       dpgc, lyh16, mall-fluffy-bongo, RoyRiv3r
// @match        https://www.perplexity.ai/*
// @license      MIT
// @run-at       document-end
// @downloadURL https://update.greasyfork.icu/scripts/490634/Perplexity%20Model%20Selection.user.js
// @updateURL https://update.greasyfork.icu/scripts/490634/Perplexity%20Model%20Selection.meta.js
// ==/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 focusAreas = ["Focus", "Academic", "Writing", "Wolfram|Alpha", "YouTube", "Reddit"];
    const focusSelectors = focusAreas.map(text => `div:contains("${text}")`).join(', ');
    const $focusElement = $(focusSelectors).closest(selector);

    if (!$focusElement.length) return;

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

    const aiModels = [
      {
        name: "Default",
        code: "turbo",
      },
      {
        name: "Claude 3.5 Sonnet",
        code: "claude2",
      },
      {
        name: "Sonar Large",
        code: "experimental",
      },
      {
        name: "GPT-4o",
        code: "gpt4o",
      },
      {
        name: "Claude 3 Opus",
        code: "claude3opus",
      },
      {
        name: "Sonar Huge",
        code: "llama_x_large",
      },
      {
        name: "Grok 2",
        code: "grok",
      },
      {
        name: "Claude 3.5 Haiku",
        code: "claude35haiku",
      },
      {
        name: "O1",
        code: "o1",
      },
    ];

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

    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 findFiberNodeWithSocket = (fiber) => {
      if (!fiber) return null;

      if (fiber.memoizedProps && fiber.memoizedProps.socket) {
        return fiber;
      }

      return (
        findFiberNodeWithSocket(fiber.child) ||
        findFiberNodeWithSocket(fiber.sibling)
      );
    };

    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 targetFiber = findFiberNodeWithSocket(fiber);
      if (!targetFiber)
        throw new Error("Failed to find fiber node with socket property");

      const settingsKey = isImageModel
        ? "default_image_generation_model"
        : "default_model";
      return await targetFiber.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);
  }
})();