import type { HighlightReply, HighlightRequest } from "./highlight-worker.ts";

let worker: Worker | null = null;
let nextId = 0;
const pending = new Map<number, PromiseWithResolvers<string | null>>();

function spawnWorker(): Worker {
  const spawned = new Worker(new URL("./highlight-worker.ts", import.meta.url), { type: "module" });
  spawned.onmessage = ({ data }: MessageEvent<HighlightReply>) => {
    const request = pending.get(data.id)!;
    pending.delete(data.id);
    if ("error" in data) request.reject(new Error(data.error));
    else request.resolve(data.html);
  };
  // also fires if the script fails to load; respawn on next use rather than hang
  spawned.onerror = (event) => {
    spawned.terminate();
    worker = null;
    for (const request of pending.values()) request.reject(new Error(event.message));
    pending.clear();
  };
  return spawned;
}

function splitHighlighted(nodes: Node[]): Node[][] {
  const lines: Node[][] = [[]];
  for (const node of nodes) {
    let parts: Node[][];
    if (node instanceof Text) {
      parts = node.data.split("\n").map((text) => text ? [document.createTextNode(text)] : []);
    } else if (node instanceof Element) {
      parts = splitHighlighted([...node.childNodes]).map((children) => {
        if (children.length === 0) return [];
        const clone = node.cloneNode(false) as Element;
        clone.append(...children);
        return [clone];
      });
    } else {
      parts = splitHighlighted([...node.childNodes]);
    }
    lines.at(-1)!.push(...parts[0]);
    for (const part of parts.slice(1)) lines.push(part);
  }
  return lines;
}

/** `source` highlighted as `language`, line by line; `null` when arborium has no such grammar. */
export async function highlightedLines(language: string, source: string): Promise<Node[][] | null> {
  worker ??= spawnWorker();
  const request = Promise.withResolvers<string | null>();
  const id = nextId++;
  pending.set(id, request);
  worker.postMessage({ id, language, source } satisfies HighlightRequest);
  const html = await request.promise;
  if (html === null) return null;
  const template = document.createElement("template");
  template.innerHTML = html;
  return splitHighlighted([...template.content.childNodes]);
}