const DATASETS_URL = "./data/datasets.json";
const PANEL_TITLE = "Dataset";
const LEGACY_PANEL_TITLE = "dataset";
const PANEL_TITLES = new Set([PANEL_TITLE, LEGACY_PANEL_TITLE]);
const LINK_CLASS = "nv3-dataset-link";
const ROW_CLASS = "nv3-dataset-link-row";
const TOOLTIP_LINK_CLASS = "nv3-dataset-tooltip-link";
const TOOLTIP_ITEM_CLASS = "nv3-dataset-tooltip-item";
const COLOR_LINK_CLASS = "nv3-dataset-color-link";
const COLOR_LABEL_CLASS = "nv3-dataset-color-label";
const COLOR_LABEL_TEXT_CLASS = "nv3-dataset-color-label-text";
const COLOR_LEGEND_SELECTOR = ".nv3-color-legend-scroll";
const COLOR_LEGEND_FIELD_ATTR = "data-nv3-field";
const DATASET_FIELD = "dataset";
const STARTUP_RETRY_MS = 15000;
const STARTUP_RETRY_INTERVAL_MS = 250;

let datasetList = [];
let datasetsByName = new Map();
let observer = null;
let scanQueued = false;
let startupRetryTimer = null;
let startupRetryDeadline = 0;

const normalizeText = (value) => (value ?? "").replace(/\s+/g, " ").trim();

function injectStyles() {
  if (document.getElementById("nv3-dataset-link-styles")) return;

  const style = document.createElement("style");
  style.id = "nv3-dataset-link-styles";
  style.textContent = `
    .${ROW_CLASS} {
      position: relative;
    }

    .${LINK_CLASS} {
      display: inline-flex;
      flex: 0 0 auto;
      width: 1rem;
      height: 1rem;
      margin-left: 0.35rem;
      align-items: center;
      justify-content: center;
      border-radius: 0.25rem;
      color: rgb(100 116 139);
      opacity: 0.8;
      text-decoration: none;
      vertical-align: -0.125rem;
    }

    .${LINK_CLASS}:hover {
      background: rgb(226 232 240);
      color: rgb(37 99 235);
      opacity: 1;
    }

    .dark .${LINK_CLASS}:hover {
      background: rgb(51 65 85);
      color: rgb(147 197 253);
    }

    .${LINK_CLASS} svg {
      width: 0.75rem;
      height: 0.75rem;
      pointer-events: none;
    }

    .${COLOR_LABEL_CLASS} {
      display: inline-flex !important;
      align-items: center;
      max-width: 100%;
      min-width: 0;
    }

    .${COLOR_LABEL_TEXT_CLASS} {
      min-width: 0;
      overflow: hidden;
      text-overflow: ellipsis;
      white-space: nowrap;
    }

    .${COLOR_LINK_CLASS} {
      margin-left: 0.25rem;
    }

    .${TOOLTIP_LINK_CLASS} {
      display: inline-flex;
      flex: 0 0 auto;
      width: 1rem;
      height: 1rem;
      margin-left: 0.2rem;
      align-items: center;
      justify-content: center;
      border-radius: 0.25rem;
      color: rgb(100 116 139);
      opacity: 0.8;
      text-decoration: none;
      vertical-align: -0.125rem;
    }

    .${TOOLTIP_LINK_CLASS}:hover {
      background: rgb(226 232 240);
      color: rgb(37 99 235);
      opacity: 1;
    }

    .dark .${TOOLTIP_LINK_CLASS}:hover {
      background: rgb(51 65 85);
      color: rgb(147 197 253);
    }

    .${TOOLTIP_LINK_CLASS} svg {
      width: 0.75rem;
      height: 0.75rem;
      flex: 0 0 auto;
      pointer-events: none;
    }

    .${TOOLTIP_ITEM_CLASS} {
      display: inline;
      white-space: normal;
    }
  `;
  document.head.append(style);
}

function makeIcon() {
  const svg = document.createElementNS("http://www.w3.org/2000/svg", "svg");
  svg.setAttribute("viewBox", "0 0 24 24");
  svg.setAttribute("fill", "none");
  svg.setAttribute("stroke", "currentColor");
  svg.setAttribute("stroke-width", "2");
  svg.setAttribute("stroke-linecap", "round");
  svg.setAttribute("stroke-linejoin", "round");
  svg.setAttribute("aria-hidden", "true");

  const paths = [
    "M18 13v6a2 2 0 0 1-2 2H5a2 2 0 0 1-2-2V8a2 2 0 0 1 2-2h6",
    "M15 3h6v6",
    "M10 14 21 3",
  ];

  for (const d of paths) {
    const path = document.createElementNS("http://www.w3.org/2000/svg", "path");
    path.setAttribute("d", d);
    svg.append(path);
  }

  return svg;
}

function datasetForText(text) {
  const normalized = normalizeText(text);
  return datasetsByName.get(normalized) ?? null;
}

function titleTextNodes() {
  if (!document.body) return [];

  const nodes = [];
  const walker = document.createTreeWalker(document.body, NodeFilter.SHOW_TEXT, {
    acceptNode(node) {
      if (isInsideTooltip(node.parentElement)) return NodeFilter.FILTER_REJECT;
      return PANEL_TITLES.has(normalizeText(node.nodeValue))
        ? NodeFilter.FILTER_ACCEPT
        : NodeFilter.FILTER_REJECT;
    },
  });

  while (walker.nextNode()) nodes.push(walker.currentNode);
  return nodes;
}

function datasetHitCount(text) {
  return datasetList.reduce(
    (count, dataset) => count + (text.includes(dataset.display_name) ? 1 : 0),
    0,
  );
}

function normalizePanelTitle(titleNode) {
  if (normalizeText(titleNode.nodeValue) === LEGACY_PANEL_TITLE) {
    titleNode.nodeValue = PANEL_TITLE;
  }
}

function findPanelFromTitle(titleNode) {
  for (
    let element = titleNode.parentElement, depth = 0;
    element && element !== document.body && depth < 12;
    element = element.parentElement, depth += 1
  ) {
    const text = normalizeText(element.textContent);
    if (![...PANEL_TITLES].some((title) => text.includes(title))) continue;
    if (text.length > 5000) continue;
    if (datasetHitCount(text) > 0) return element;
  }

  return null;
}

function findRowForLabel(labelElement, displayName) {
  let best = labelElement;

  for (
    let element = labelElement;
    element && element !== document.body;
    element = element.parentElement
  ) {
    const text = normalizeText(element.textContent);
    if (!text.includes(displayName)) continue;
    if ([...PANEL_TITLES].some((title) => text.includes(title))) break;
    if (text.length > displayName.length + 160) break;
    best = element;
  }

  return best;
}

function createLink(dataset) {
  if (!dataset.url) return null;

  const link = document.createElement("a");
  link.className = LINK_CLASS;
  link.href = dataset.url;
  link.target = "_blank";
  link.rel = "noopener noreferrer";
  link.title = `Open ${dataset.source_repo ?? dataset.display_name} on Hugging Face`;
  link.setAttribute("aria-label", link.title);
  link.dataset.datasetId = dataset.id;
  link.append(makeIcon());

  for (const eventName of ["pointerdown", "mousedown", "click"]) {
    link.addEventListener(eventName, (event) => {
      event.stopPropagation();
    });
  }

  return link;
}

function createColorLegendLink(dataset) {
  const link = createLink(dataset);
  if (link) link.classList.add(COLOR_LINK_CLASS);
  return link;
}

function createTooltipLink(dataset) {
  if (!dataset.url) return null;

  const link = document.createElement("a");
  link.className = TOOLTIP_LINK_CLASS;
  link.href = dataset.url;
  link.target = "_blank";
  link.rel = "noopener noreferrer";
  link.title = `Open ${dataset.source_repo ?? dataset.display_name} on Hugging Face`;
  link.setAttribute("aria-label", link.title);
  link.dataset.datasetId = dataset.id;
  link.append(makeIcon());

  for (const eventName of ["pointerdown", "mousedown", "click"]) {
    link.addEventListener(eventName, (event) => {
      event.stopPropagation();
    });
  }

  return link;
}

function createTooltipItem(dataset) {
  const item = document.createElement("span");
  item.className = TOOLTIP_ITEM_CLASS;
  item.dataset.datasetId = dataset.id;
  item.append(document.createTextNode(dataset.display_name));

  const link = createTooltipLink(dataset);
  if (link) item.append(link);

  return item;
}

function attachLink(labelElement, dataset) {
  const row = findRowForLabel(labelElement, dataset.display_name);
  if (!row || row.querySelector(`.${LINK_CLASS}[data-dataset-id="${dataset.id}"]`)) {
    return;
  }

  const link = createLink(dataset);
  if (!link) return;

  row.classList.add(ROW_CLASS);
  row.append(link);
}

function enhancePanel(panel) {
  const walker = document.createTreeWalker(panel, NodeFilter.SHOW_TEXT, {
    acceptNode(node) {
      if (node.parentElement?.closest(`.${LINK_CLASS}`)) return NodeFilter.FILTER_REJECT;
      return datasetForText(node.nodeValue)
        ? NodeFilter.FILTER_ACCEPT
        : NodeFilter.FILTER_REJECT;
    },
  });

  const matches = [];
  while (walker.nextNode()) matches.push(walker.currentNode);

  for (const node of matches) {
    const dataset = datasetForText(node.nodeValue);
    if (dataset && node.parentElement) {
      attachLink(node.parentElement, dataset);
    }
  }

  for (const element of panel.querySelectorAll("[title]")) {
    const dataset = datasetForText(element.getAttribute("title"));
    if (dataset) attachLink(element, dataset);
  }
}

function labelTextForColorLegend(labelElement) {
  const title = normalizeText(labelElement.getAttribute("title"));
  if (title) return title;

  return normalizeText(
    [...labelElement.childNodes]
      .filter(
        (node) =>
          node.nodeType === Node.TEXT_NODE ||
          (node.nodeType === Node.ELEMENT_NODE &&
            node.classList.contains(COLOR_LABEL_TEXT_CLASS)),
      )
      .map((node) => (node.nodeType === Node.TEXT_NODE ? node.nodeValue : node.textContent))
      .join(" ") || labelElement.textContent,
  );
}

function enhanceColorLegendLabel(labelElement, dataset) {
  for (const link of labelElement.querySelectorAll(`.${COLOR_LINK_CLASS}`)) {
    if (link.dataset.datasetId !== dataset.id) link.remove();
  }

  labelElement.classList.add(COLOR_LABEL_CLASS);

  let text = labelElement.querySelector(`.${COLOR_LABEL_TEXT_CLASS}`);
  if (!text) {
    text = document.createElement("span");
    text.className = COLOR_LABEL_TEXT_CLASS;
    text.dataset.nv3Wrapped = "1";

    const nodes = [...labelElement.childNodes].filter(
      (node) =>
        !(
          node.nodeType === Node.ELEMENT_NODE &&
          node.classList.contains(COLOR_LINK_CLASS)
        ),
    );
    labelElement.prepend(text);
    for (const node of nodes) text.append(node);
  }

  if (!normalizeText(text.textContent)) {
    text.textContent = labelTextForColorLegend(labelElement) || dataset.display_name;
  }

  if (!labelElement.querySelector(`.${COLOR_LINK_CLASS}[data-dataset-id="${dataset.id}"]`)) {
    const link = createColorLegendLink(dataset);
    if (link) labelElement.append(link);
  }
}

function restoreColorLegendLabel(labelElement) {
  for (const link of labelElement.querySelectorAll(`.${COLOR_LINK_CLASS}`)) link.remove();

  if (!labelElement.classList.contains(COLOR_LABEL_CLASS)) return;

  const text = labelElement.querySelector(`.${COLOR_LABEL_TEXT_CLASS}`);
  if (text?.dataset.nv3Wrapped === "1") {
    while (text.firstChild) labelElement.insertBefore(text.firstChild, text);
    text.remove();
  } else if (text) {
    const restoredText = normalizeText(labelElement.getAttribute("title")) || text.textContent;
    labelElement.textContent = restoredText;
  }

  labelElement.classList.remove(COLOR_LABEL_CLASS);
}

function isDatasetColorLegend(table) {
  return normalizeText(table.getAttribute(COLOR_LEGEND_FIELD_ATTR)) === DATASET_FIELD;
}

function enhanceColorLegends() {
  for (const table of document.querySelectorAll(COLOR_LEGEND_SELECTOR)) {
    const isDatasetLegend = isDatasetColorLegend(table);

    for (const row of table.querySelectorAll("tr")) {
      const labelElement =
        row.querySelector("td:nth-child(2) div") ?? row.querySelector("td:nth-child(2)");
      if (!labelElement) continue;

      if (!isDatasetLegend) {
        restoreColorLegendLabel(labelElement);
        continue;
      }

      const dataset = datasetForText(labelTextForColorLegend(labelElement));
      if (dataset) enhanceColorLegendLabel(labelElement, dataset);
    }
  }
}

function datasetMatchInText(text, startIndex) {
  let match = null;

  for (const dataset of datasetList) {
    const index = text.indexOf(dataset.display_name, startIndex);
    if (index < 0) continue;
    if (
      match == null ||
      index < match.index ||
      (index === match.index && dataset.display_name.length > match.dataset.display_name.length)
    ) {
      match = { index, dataset };
    }
  }

  return match;
}

function textIncludesDataset(text) {
  return datasetMatchInText(text, 0) != null;
}

function enhanceTooltipTextNode(node) {
  if (
    node.parentElement?.closest(`.${LINK_CLASS}, .${TOOLTIP_LINK_CLASS}, .${TOOLTIP_ITEM_CLASS}`)
  ) {
    return;
  }

  const text = node.nodeValue ?? "";
  let index = 0;
  const fragment = document.createDocumentFragment();

  while (index < text.length) {
    const match = datasetMatchInText(text, index);
    if (match == null) {
      fragment.append(document.createTextNode(text.slice(index)));
      break;
    }

    if (match.index > index) {
      fragment.append(document.createTextNode(text.slice(index, match.index)));
    }

    fragment.append(createTooltipItem(match.dataset));
    index = match.index + match.dataset.display_name.length;
  }

  node.replaceWith(fragment);
}

function isTooltipContainer(element) {
  return (
    element.classList.contains("backdrop-blur-sm") &&
    element.classList.contains("overflow-y-scroll")
  );
}

function tooltipContainers() {
  return [...document.querySelectorAll(".embedding-atlas-root div")].filter(isTooltipContainer);
}

function isInsideTooltip(element) {
  return Boolean(element?.closest(".backdrop-blur-sm.overflow-y-scroll"));
}

function enhanceTooltip(container) {
  const walker = document.createTreeWalker(container, NodeFilter.SHOW_TEXT, {
    acceptNode(node) {
      if (
        node.parentElement?.closest(
          `.${LINK_CLASS}, .${TOOLTIP_LINK_CLASS}, .${TOOLTIP_ITEM_CLASS}`,
        )
      ) {
        return NodeFilter.FILTER_REJECT;
      }
      return textIncludesDataset(node.nodeValue ?? "")
        ? NodeFilter.FILTER_ACCEPT
        : NodeFilter.FILTER_REJECT;
    },
  });

  const matches = [];
  while (walker.nextNode()) matches.push(walker.currentNode);

  for (const node of matches) enhanceTooltipTextNode(node);
}

function scan() {
  scanQueued = false;
  if (datasetList.length === 0) return;

  const panels = new Set(
    titleTextNodes()
      .map((titleNode) => {
        normalizePanelTitle(titleNode);
        return findPanelFromTitle(titleNode);
      })
      .filter(Boolean),
  );

  for (const panel of panels) enhancePanel(panel);
  enhanceColorLegends();
  for (const tooltip of tooltipContainers()) enhanceTooltip(tooltip);
}

function queueScan() {
  if (scanQueued) return;
  scanQueued = true;
  window.requestAnimationFrame(scan);
}

function startStartupRetries() {
  startupRetryDeadline = Date.now() + STARTUP_RETRY_MS;
  if (startupRetryTimer != null) return;

  startupRetryTimer = window.setInterval(() => {
    queueScan();
    if (Date.now() >= startupRetryDeadline) {
      window.clearInterval(startupRetryTimer);
      startupRetryTimer = null;
    }
  }, STARTUP_RETRY_INTERVAL_MS);
}

async function init() {
  try {
    const response = await fetch(DATASETS_URL, { cache: "no-store" });
    if (!response.ok) throw new Error(`HTTP ${response.status}`);
    const data = await response.json();
    datasetList = Array.isArray(data.datasets) ? data.datasets : [];
    datasetsByName = new Map(
      datasetList.map((dataset) => [normalizeText(dataset.display_name), dataset]),
    );
  } catch (error) {
    console.warn("Dataset links unavailable:", error);
    return;
  }

  injectStyles();
  queueScan();
  startStartupRetries();

  window.addEventListener("load", queueScan, { once: true });
  window.addEventListener("nemotron-atlas-ready", () => {
    queueScan();
    startStartupRetries();
  });

  observer = new MutationObserver(queueScan);
  observer.observe(document.body, {
    childList: true,
    subtree: true,
    characterData: true,
  });
}

if (document.readyState === "loading") {
  document.addEventListener("DOMContentLoaded", init, { once: true });
} else {
  init();
}