import { Handler, registerHandlers } from "pagedjs";
import { resolveBreakTokenElement } from "./paged-break-token";

interface PagedBreakToken {
  node?: Node;
}

interface PagedChunker {
  source: ParentNode;
}

interface PagedHandlerContext {
  chunker: PagedChunker;
}

function elementAncestors(
  element: Element,
  selector: string
) {
  const ancestors: Element[] = [];
  let current = element.parentElement;

  while (current) {
    if (current.matches(selector)) {
      ancestors.unshift(current);
    }
    current = current.parentElement;
  }

  return ancestors;
}

class RepeatTableHeadersHandler extends Handler {
  declare chunker: PagedChunker;
  private splitTableRefs: string[] = [];

  constructor(
    chunker: unknown,
    polisher: unknown,
    caller: unknown
  ) {
    super(chunker, polisher, caller);
  }

  afterPageLayout(
    pageElement: HTMLElement,
    _page: unknown,
    breakToken?: PagedBreakToken
  ) {
    this.splitTableRefs = [];
    const element = resolveBreakTokenElement(breakToken?.node);
    if (!element) {
      return;
    }

    const tables = elementAncestors(element, "table");
    if (element.matches("table")) {
      tables.push(element);
    }

    this.splitTableRefs = Array.from(
      new Set(
        tables
          .map((table) => table.getAttribute("data-ref"))
          .filter((ref): ref is string => Boolean(ref))
      )
    );

    for (const ref of this.splitTableRefs) {
      const renderedTable =
        pageElement.querySelector<HTMLElement>(
          `table[data-ref="${CSS.escape(ref)}"]`
        );
      if (!renderedTable?.querySelector("tbody > tr")) {
        renderedTable?.setAttribute(
          "data-empty-split-table",
          "true"
        );
      }
    }
  }

  layout(rendered: HTMLElement) {
    for (const ref of this.splitTableRefs) {
      const renderedTable =
        rendered.querySelector<HTMLTableElement>(
          `table[data-ref="${CSS.escape(ref)}"]`
        );
      if (
        !renderedTable ||
        renderedTable.hasAttribute("data-repeated-header")
      ) {
        continue;
      }

      const sourceTable =
        this.chunker.source.querySelector<HTMLTableElement>(
          `table[data-ref="${CSS.escape(ref)}"]`
        );
      if (!sourceTable) {
        continue;
      }

      const firstChild = renderedTable.firstChild;
      for (const colgroup of sourceTable.querySelectorAll("colgroup")) {
        renderedTable.insertBefore(
          colgroup.cloneNode(true),
          firstChild
        );
      }

      if (!renderedTable.querySelector("thead")) {
        const sourceHeader = sourceTable.querySelector("thead");
        if (sourceHeader) {
          renderedTable.insertBefore(
            sourceHeader.cloneNode(true),
            renderedTable.firstChild
          );
        }
      }

      renderedTable.setAttribute("data-repeated-header", "true");
    }
  }
}

const handlerRegistration = globalThis as typeof globalThis & {
  __mdToPdfTableHandlerRegistered?: boolean;
};

if (!handlerRegistration.__mdToPdfTableHandlerRegistered) {
  registerHandlers(RepeatTableHeadersHandler);
  handlerRegistration.__mdToPdfTableHandlerRegistered = true;
}

export { RepeatTableHeadersHandler };