// reads docs/*.md and renders them. the markdown stays the only source: this adds heading
// ids, turns links between docs into site urls, highlights code, and fails the build on a
// link to a doc or heading that doesn't exist.

import fs from "node:fs";
import path from "node:path";
import { markdownToHast } from "satteri";
import GithubSlugger from "github-slugger";
import { visit, SKIP } from "unist-util-visit";
import { toString } from "hast-util-to-string";
import { toHtml } from "hast-util-to-html";
import { createHighlighter } from "shiki";
import katex from "katex";
import type { Root, Element, ElementContent } from "hast";
import { placements, SETTINGS_PAGES, type Placement } from "./nav";
import { buildSettings, type SettingsSection } from "./settings";

export const DOCS_DIR = path.resolve(process.env.HYDRANT_DOCS_DIR ?? path.join(process.cwd(), "../docs"));

export interface Heading {
  depth: number;
  text: string;
  id: string;
}

export interface Doc {
  /** path under docs/ without .md, e.g. "api/filter" */
  id: string;
  title: string;
  url: string;
  place: Placement;
  headings: Heading[];
  /** set on settings pages (configuration), in page order */
  settings?: SettingsSection[];
  html: string;
  /** the rendered tree, for pages that restructure their own content */
  tree: Root;
}

/** README is a section's own page, everything else keeps its path */
export function urlFor(id: string): string {
  if (id === "README") return "/overview/";
  if (id.endsWith("/README")) return `/${id.slice(0, -"/README".length)}/`;
  return `/${id}/`;
}

function listDocs(dir: string, prefix = ""): string[] {
  return fs.readdirSync(dir, { withFileTypes: true }).flatMap((e) => {
    const rel = prefix ? `${prefix}/${e.name}` : e.name;
    if (e.isDirectory()) return listDocs(path.join(dir, e.name), rel);
    return e.name.endsWith(".md") ? [rel.slice(0, -3)] : [];
  });
}

export function frontmatter(src: string): { title?: string; body: string } {
  const m = /^---\r?\n([\s\S]*?)\r?\n---\r?\n?/.exec(src);
  if (!m) return { body: src };
  const title = /^title:\s*(.+?)\s*$/m.exec(m[1])?.[1];
  return { title, body: src.slice(m[0].length) };
}

const isElement = (n: unknown, tag?: string): n is Element =>
  (n as Element)?.type === "element" && (!tag || (n as Element).tagName === tag);

/** a doc without frontmatter (the specs) is titled by its first h1, which then leaves the body */
function takeTitle(tree: Root): string | undefined {
  const i = tree.children.findIndex((n) => isElement(n, "h1"));
  if (i < 0) return undefined;
  const title = toString(tree.children[i] as Element).replace(/^Spec:\s*/, "");
  tree.children.splice(i, 1);
  return title;
}

function addHeadingIds(tree: Root): Heading[] {
  const slugger = new GithubSlugger();
  const out: Heading[] = [];
  visit(tree, "element", (node: Element) => {
    const m = /^h([1-6])$/.exec(node.tagName);
    if (!m) return;
    const text = toString(node);
    const id = slugger.slug(text);
    node.properties.id = id;
    out.push({ depth: Number(m[1]), text, id });
    return SKIP;
  });
  return out;
}

const EXTERNAL = /^[a-z][a-z0-9+.-]*:/i;

function rewriteLinks(tree: Root, from: string, ids: Map<string, Set<string>>, errors: string[]) {
  visit(tree, "element", (node: Element) => {
    if (node.tagName !== "a") return;
    const href = String(node.properties.href ?? "");
    if (!href || EXTERNAL.test(href)) return;
    const [file, hash = ""] = href.split("#");
    let target = from;
    if (file) {
      if (!file.endsWith(".md")) {
        errors.push(`${from}.md links to ${href}, which isn't a doc`);
        return;
      }
      target = path.posix.normalize(path.posix.join(path.posix.dirname(from), file.slice(0, -3)));
    }
    const anchors = ids.get(target);
    if (!anchors) {
      errors.push(`${from}.md links to ${href}, but there's no docs/${target}.md`);
      return;
    }
    if (hash && !anchors.has(hash)) {
      errors.push(`${from}.md links to ${href}, but docs/${target}.md has no heading #${hash}`);
      return;
    }
    node.properties.href = file ? urlFor(target) + (hash ? `#${hash}` : "") : `#${hash}`;
  });
}

type Highlighter = Awaited<ReturnType<typeof createHighlighter>>;

function highlightCode(tree: Root, hl: Highlighter) {
  const langs = new Set(hl.getLoadedLanguages());
  visit(tree, "element", (node: Element, index, parent) => {
    if (node.tagName !== "pre" || !parent || index === undefined) return;
    const code = node.children.find((c): c is Element => isElement(c, "code"));
    if (!code) return;
    const cls = ((code.properties.className as string[] | undefined) ?? []).find((c) => c.startsWith("language-"));
    const lang = cls?.slice("language-".length);
    const text = toString(code).replace(/\n$/, "");
    const highlighted = hl.codeToHast(text, {
      lang: lang && langs.has(lang) ? lang : "text",
      theme: "hydrant",
    });
    const pre = highlighted.children[0] as Element;
    // shiki writes its own `class` too, which would come out as a second class attribute
    delete pre.properties.class;
    pre.properties.className = ["code"];
    delete pre.properties.style;
    delete pre.properties.tabindex;
    if (lang) pre.properties.dataLang = lang;
    parent.children[index] = { type: "element", tagName: "div", properties: { className: ["codebox"] }, children: [pre] };
    return SKIP;
  });
}

/** a paragraph that is all $$ math is a display block on github and tangled, even written
 *  `$$a\nb$$`. the parser only knows the fenced `$$\na\nb\n$$` form (and otherwise reads to
 *  the end of the doc looking for the closing fence), so fence those first. code is left alone */
function fenceMath(md: string): string {
  return md
    .split(/(^```[\s\S]*?^```)/m)
    .map((part, i) => (i % 2 ? part : part.replace(/(^|\n\n)\$\$(?!\n)((?:(?!\n\n)[\s\S])+?)\$\$(?=\n\n|\n?$)/g, "$1$$$$\n$2\n$$$$")))
    .join("");
}

/** $$ math as mathml, which browsers draw themselves, so there are no math fonts or css to
 *  ship */
function renderMath(tree: Root): string[] {
  const errors: string[] = [];
  visit(tree, "element", (node: Element, index, parent) => {
    const cls = (node.properties.className as string[] | undefined) ?? [];
    if (node.tagName !== "code" || !cls.includes("language-math") || !parent || index === undefined) return;
    const alone = isElement(parent, "pre");
    let html: string;
    try {
      html = katex.renderToString(toString(node), { displayMode: alone, output: "mathml", throwOnError: true });
    } catch (e) {
      errors.push(`math that doesn't render: ${(e as Error).message}`);
      return;
    }
    const raw = { type: "raw", value: html } as unknown as ElementContent;
    if (!alone) {
      parent.children[index] = raw;
      return;
    }
    Object.assign(parent, { tagName: "div", properties: { className: ["math"] }, children: [raw] });
  });
  return errors;
}

/** wide tables scroll inside the paper instead of stretching it */
function wrapTables(tree: Root) {
  visit(tree, "element", (node: Element, index, parent) => {
    if (node.tagName !== "table" || !parent || index === undefined) return;
    parent.children[index] = { type: "element", tagName: "div", properties: { className: ["tablewrap"] }, children: [node] };
    return SKIP;
  });
}

// code colours are the palette's --code-* variables. shiki only takes real colours, so each
// role gets a stand-in (#000001, #000002, ...) that's swapped for its variable on the way out.
const ROLES: { role: string; scope: string[]; fontStyle?: string }[] = [
  { role: "text", scope: [] },
  { role: "comment", scope: ["comment", "punctuation.definition.comment"], fontStyle: "italic" },
  { role: "string", scope: ["string", "string.quoted", "markup.inline.raw"] },
  { role: "constant", scope: ["constant.numeric", "constant.language", "constant.character"] },
  { role: "keyword", scope: ["keyword", "storage", "storage.type", "keyword.control"], fontStyle: "bold" },
  { role: "function", scope: ["entity.name.function", "support.function"] },
  { role: "type", scope: ["entity.name.type", "support.type", "entity.name.class"] },
  { role: "property", scope: ["support.type.property-name", "meta.object-literal.key", "entity.name.tag"] },
  { role: "parameter", scope: ["variable.parameter", "variable.other.option"] },
];
const standIn = (i: number) => `#${(i + 1).toString(16).padStart(6, "0")}`;
const THEME = {
  name: "hydrant",
  type: "light" as const,
  colors: { "editor.foreground": standIn(0) },
  tokenColors: ROLES.slice(1).map(({ scope, fontStyle }, i) => ({ scope, settings: { foreground: standIn(i + 1), fontStyle } })),
  colorReplacements: Object.fromEntries(ROLES.map(({ role }, i) => [standIn(i), `var(--code-${role})`])),
};

let cache: Promise<Map<string, Doc>> | undefined;

/** built once per build; in dev every request reads the docs again, so edits show on reload */
export function loadDocs(): Promise<Map<string, Doc>> {
  if (import.meta.env.DEV) return build();
  cache ??= build();
  return cache;
}

async function build(): Promise<Map<string, Doc>> {
  const nav = placements();
  const errors: string[] = [];
  const found = listDocs(DOCS_DIR).sort();
  for (const id of found) if (!nav.has(id)) errors.push(`docs/${id}.md isn't in site/src/lib/nav.ts, so it has no place on the site`);
  for (const id of nav.keys()) if (!found.includes(id)) errors.push(`site/src/lib/nav.ts lists ${id}, but there's no docs/${id}.md`);

  const parsed = found.filter((id) => nav.has(id)).map((id) => {
    const { title: fmTitle, body } = frontmatter(fs.readFileSync(path.join(DOCS_DIR, `${id}.md`), "utf8"));
    // only $$ is math, so a lone $ in prose stays a dollar sign
    const tree = markdownToHast(fenceMath(body), { features: { math: { singleDollarTextMath: false } } }) as Root;
    for (const e of renderMath(tree)) errors.push(`docs/${id}.md: ${e}`);
    const title = fmTitle ?? takeTitle(tree) ?? id;
    const headings = addHeadingIds(tree);
    let settings: SettingsSection[] | undefined;
    try {
      settings = SETTINGS_PAGES.includes(id) ? buildSettings(tree, id) : undefined;
    } catch (e) {
      errors.push((e as Error).message);
    }
    return { id, title, tree, headings, settings };
  });
  // a setting is linkable like a heading: configuration.md#FIREHOSE_WORKERS
  const ids = new Map(parsed.map((d) => [d.id, new Set([...d.headings.map((h) => h.id), ...(d.settings ?? []).flatMap((s) => s.settings.map((x) => x.key))])]));
  for (const d of parsed) rewriteLinks(d.tree, d.id, ids, errors);
  if (errors.length) throw new Error(`the docs don't build:\n  ${errors.join("\n  ")}`);

  const hl = await createHighlighter({ themes: [THEME], langs: ["bash", "json", "rust", "go", "toml"] });
  const out = new Map<string, Doc>();
  for (const d of parsed) {
    highlightCode(d.tree, hl);
    wrapTables(d.tree);
    out.set(d.id, { ...d, url: urlFor(d.id), place: nav.get(d.id)!, html: toHtml(d.tree, { allowDangerousHtml: true }) });
  }
  return out;
}