import { Marked, type Token, type TokensList } from "marked";
import { code } from "./code.tsx";

const escape = (text: string) =>
  text.replace(
    /[&<>"']/g,
    (char) => `&${{ "&": "amp", "<": "lt", ">": "gt", '"': "quot", "'": "#39" }[char]};`,
  );

// model output is untrusted: raw html shows as text, and links only go where links should
const marked = new Marked({
  gfm: true,
  tokenizer: {
    del: () => undefined,
  },
  renderer: {
    html: ({ text }) => escape(text),
    link({ href, tokens }) {
      const text = this.parser.parseInline(tokens);
      if (!/^(https?:|mailto:|#|\/)/i.test(href)) return text;
      return `<a href="${escape(href)}" target="_blank" rel="noopener noreferrer">${text}</a>`;
    },
  },
});

/** One top-level block's nodes; its code is highlighted only once `settled`, since a fence being written still grows. */
function draw(token: Token, links: TokensList["links"], settled: boolean): ChildNode[] {
  const template = document.createElement("template");
  template.innerHTML = marked.parser(Object.assign([token], { links }));
  for (const block of template.content.querySelectorAll("pre > code")) {
    const source = block.textContent ?? "";
    if (!settled) {
      // a growing fence redraws whole on every revision, so it stays one text node until it settles
      const pre = document.createElement("pre");
      pre.className = "code bare";
      const text = document.createElement("span");
      text.className = "code-text";
      text.textContent = source.replace(/\n$/, "");
      pre.append(text);
      block.parentElement!.replaceWith(pre);
      continue;
    }
    const language = block.className.match(/language-(\S+)/)?.[1];
    block.parentElement!.replaceWith(
      code(source, { gutter: false, ...(language ? { language } : {}) }),
    );
  }
  return [...template.content.childNodes];
}

type Drawn = { raw: string; settled: boolean; nodes: ChildNode[] };

/**
 * Markdown that may still be being written: each `update` draws again only the top-level blocks that changed, which
 * while it streams is the last. `final` says no more is coming.
 */
export function markdownView() {
  const element = document.createElement("div");
  element.className = "markdown";
  let drawn: Drawn[] = [];
  let shown: { text: string; final: boolean } | null = null;
  let links = "";
  return {
    element,
    update(text: string, final: boolean) {
      if (shown?.text === text && shown.final === final) return;
      shown = { text, final };
      const tokens = marked.lexer(text);
      // a reference link resolves against definitions anywhere in the text, even those written later
      const known = JSON.stringify(tokens.links);
      if (known !== links) {
        links = known;
        drawn = [];
        element.replaceChildren();
      }
      const next = tokens.map((token, index): Drawn => {
        const settled = final || index < tokens.length - 1;
        const before = drawn[index];
        if (before?.raw === token.raw && before.settled === settled) return before;
        return { raw: token.raw, settled, nodes: draw(token, tokens.links, settled) };
      });
      let first = 0;
      while (first < next.length && next[first] === drawn[first]) first++;
      for (const { nodes } of drawn.slice(first)) for (const node of nodes) node.remove();
      element.append(...next.slice(first).flatMap(({ nodes }) => nodes));
      drawn = next;
    },
  };
}

export function markdown(text: string): HTMLElement {
  const view = markdownView();
  view.update(text, true);
  return view.element;
}