ultisuite-client/lib/drive/docx-header-footer-export.ts

import type { DocPageHeaderFooter, DocPageSetup } from "@/lib/drive/doc-page-setup"
import type { TipTapJSON } from "@/lib/drive/richtext-import"
import { prepareTablesForDocxExport } from "@/lib/drive/docs-table-export"

type DocxArchive = Record<string, Uint8Array>

const W_NS = "http://schemas.openxmlformats.org/wordprocessingml/2006/main"
const R_NS = "http://schemas.openxmlformats.org/officeDocument/2006/relationships"
const REL_NS = "http://schemas.openxmlformats.org/package/2006/relationships"
const CT_NS = "http://schemas.openxmlformats.org/package/2006/content-types"

function decodeXml(bytes: Uint8Array | undefined): string {
  if (!bytes) return ""
  return new TextDecoder().decode(bytes)
}

function encodeXml(text: string): Uint8Array {
  return new TextEncoder().encode(text)
}

function escapeXml(text: string): string {
  return text
    .replace(/&/g, "&amp;")
    .replace(/</g, "&lt;")
    .replace(/>/g, "&gt;")
    .replace(/"/g, "&quot;")
}

async function tipTapToBodyXml(content: TipTapJSON): Promise<string> {
  try {
    const { generateDOCX } = await import("@docen/export-docx")
    const prepared = prepareTablesForDocxExport(content)
    const buf = await generateDOCX(prepared, { outputType: "uint8array" })
    const { unzipSync } = await import("fflate")
    const archive = unzipSync(buf as Uint8Array) as DocxArchive
    const docXml = decodeXml(archive["word/document.xml"])
    const bodyMatch = docXml.match(/<w:body>([\s\S]*?)<\/w:body>/i)
    if (bodyMatch?.[1]) {
      return bodyMatch[1]
        .replace(/<w:sectPr\b[\s\S]*?<\/w:sectPr>/gi, "")
        .replace(/<w:sectPr\b[^>]*\/>/gi, "")
    }
  } catch {
    /* fallback below */
  }

  const paragraphs: string[] = []
  const walk = (node: unknown) => {
    if (!node || typeof node !== "object") return
    const record = node as TipTapJSON
    if (record.type === "paragraph" || record.type === "heading") {
      const texts: string[] = []
      const visit = (n: unknown) => {
        if (!n || typeof n !== "object") return
        const r = n as TipTapJSON
        if (r.type === "text" && typeof r.text === "string") texts.push(escapeXml(r.text))
        if (Array.isArray(r.content)) r.content.forEach(visit)
      }
      if (Array.isArray(record.content)) record.content.forEach(visit)
      paragraphs.push(
        `<w:p><w:r><w:t xml:space="preserve">${texts.join("")}</w:t></w:r></w:p>`
      )
    }
    if (Array.isArray(record.content)) record.content.forEach(walk)
  }
  walk(content)
  return paragraphs.join("")
}

function wrapHeaderFooterPart(bodyInner: string, kind: "hdr" | "ftr"): string {
  const tag = kind === "hdr" ? "header" : "footer"
  return `<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<w:${tag} xmlns:w="${W_NS}" xmlns:r="${R_NS}">${bodyInner}</w:${tag}>`
}

async function regionToPartXml(
  region: DocPageHeaderFooter | null | undefined,
  kind: "hdr" | "ftr",
  pageNumbers?: DocPageSetup["pageNumbers"]
): Promise<string | null> {
  if (!region?.content) return null
  let bodyInner = await tipTapToBodyXml(region.content as TipTapJSON)

  if (
    pageNumbers?.enabled &&
    pageNumbers.placement === (kind === "hdr" ? "header" : "footer")
  ) {
    const align =
      pageNumbers.align === "center"
        ? "center"
        : pageNumbers.align === "right"
          ? "right"
          : "left"
    const pageField = `<w:fldSimple w:instr=" PAGE "><w:r><w:t>1</w:t></w:r></w:fldSimple>`
    bodyInner += `<w:p><w:pPr><w:jc w:val="${align}"/></w:pPr><w:r>${pageField}</w:r></w:p>`
  }

  if (!bodyInner.trim()) return null
  return wrapHeaderFooterPart(bodyInner, kind)
}

function addRelationship(relsXml: string, id: string, type: string, target: string): string {
  if (relsXml.includes(`Id="${id}"`)) return relsXml
  const insert = `<Relationship Id="${id}" Type="${type}" Target="${target}"/>`
  return relsXml.replace("</Relationships>", `${insert}</Relationships>`)
}

function addOverride(contentTypesXml: string, partName: string, contentType: string): string {
  if (contentTypesXml.includes(`PartName="/${partName}"`)) return contentTypesXml
  const insert = `<Override PartName="/${partName}" ContentType="${contentType}"/>`
  return contentTypesXml.replace("</Types>", `${insert}</Types>`)
}

function patchDocumentSectPr(
  documentXml: string,
  refs: { header?: string; footer?: string; headerFirst?: string; footerFirst?: string },
  differentFirst: boolean
): string {
  const stripRefs = (input: string) =>
    input
      .replace(/<w:headerReference\b[^>]*\/?>/gi, "")
      .replace(/<w:footerReference\b[^>]*\/?>/gi, "")
      .replace(/<w:titlePg\b[^>]*\/?>/gi, "")

  const buildRefs = () => {
    const parts: string[] = []
    if (refs.header) parts.push(`<w:headerReference w:type="default" r:id="${refs.header}"/>`)
    if (refs.footer) parts.push(`<w:footerReference w:type="default" r:id="${refs.footer}"/>`)
    if (differentFirst && refs.headerFirst) {
      parts.push(`<w:headerReference w:type="first" r:id="${refs.headerFirst}"/>`)
    }
    if (differentFirst && refs.footerFirst) {
      parts.push(`<w:footerReference w:type="first" r:id="${refs.footerFirst}"/>`)
    }
    if (differentFirst) parts.push("<w:titlePg/>")
    return parts.join("")
  }

  const refsXml = buildRefs()

  if (/<w:sectPr\b[^>]*>[\s\S]*?<\/w:sectPr>/i.test(documentXml)) {
    return documentXml.replace(/<w:sectPr\b[^>]*>([\s\S]*?)<\/w:sectPr>/i, (_m, inner) => {
      return `<w:sectPr>${stripRefs(inner)}${refsXml}</w:sectPr>`
    })
  }

  return documentXml.replace(/<\/w:body>/i, `<w:sectPr>${refsXml}</w:sectPr></w:body>`)
}

/** Inject header/footer parts and references into a DOCX archive buffer. */
export async function patchDocxHeaderFooter(
  buffer: Uint8Array,
  pageSetup: DocPageSetup
): Promise<Uint8Array> {
  const { unzipSync, zipSync } = await import("fflate")
  const archive = { ...(unzipSync(buffer) as DocxArchive) }

  const headerXml = await regionToPartXml(pageSetup.header, "hdr", pageSetup.pageNumbers)
  const footerXml = await regionToPartXml(pageSetup.footer, "ftr", pageSetup.pageNumbers)
  const headerFirstXml = pageSetup.headerFooterDifferentFirstPage
    ? await regionToPartXml(pageSetup.headerFirstPage, "hdr", pageSetup.pageNumbers)
    : null
  const footerFirstXml = pageSetup.headerFooterDifferentFirstPage
    ? await regionToPartXml(pageSetup.footerFirstPage, "ftr", pageSetup.pageNumbers)
    : null

  if (!headerXml && !footerXml && !headerFirstXml && !footerFirstXml) {
    return buffer
  }

  let relsXml =
    decodeXml(archive["word/_rels/document.xml.rels"]) ||
    `<?xml version="1.0" encoding="UTF-8"?><Relationships xmlns="${REL_NS}"></Relationships>`

  let contentTypesXml =
    decodeXml(archive["[Content_Types].xml"]) ||
    `<?xml version="1.0" encoding="UTF-8"?><Types xmlns="${CT_NS}"></Types>`

  const refs: { header?: string; footer?: string; headerFirst?: string; footerFirst?: string } = {}

  if (headerXml) {
    archive["word/header1.xml"] = encodeXml(headerXml)
    relsXml = addRelationship(
      relsXml,
      "rIdHeader1",
      "http://schemas.openxmlformats.org/officeDocument/2006/relationships/header",
      "header1.xml"
    )
    contentTypesXml = addOverride(
      contentTypesXml,
      "word/header1.xml",
      "application/vnd.openxmlformats-officedocument.wordprocessingml.header+xml"
    )
    refs.header = "rIdHeader1"
  }

  if (footerXml) {
    archive["word/footer1.xml"] = encodeXml(footerXml)
    relsXml = addRelationship(
      relsXml,
      "rIdFooter1",
      "http://schemas.openxmlformats.org/officeDocument/2006/relationships/footer",
      "footer1.xml"
    )
    contentTypesXml = addOverride(
      contentTypesXml,
      "word/footer1.xml",
      "application/vnd.openxmlformats-officedocument.wordprocessingml.footer+xml"
    )
    refs.footer = "rIdFooter1"
  }

  if (headerFirstXml) {
    archive["word/header2.xml"] = encodeXml(headerFirstXml)
    relsXml = addRelationship(
      relsXml,
      "rIdHeader2",
      "http://schemas.openxmlformats.org/officeDocument/2006/relationships/header",
      "header2.xml"
    )
    contentTypesXml = addOverride(
      contentTypesXml,
      "word/header2.xml",
      "application/vnd.openxmlformats-officedocument.wordprocessingml.header+xml"
    )
    refs.headerFirst = "rIdHeader2"
  }

  if (footerFirstXml) {
    archive["word/footer2.xml"] = encodeXml(footerFirstXml)
    relsXml = addRelationship(
      relsXml,
      "rIdFooter2",
      "http://schemas.openxmlformats.org/officeDocument/2006/relationships/footer",
      "footer2.xml"
    )
    contentTypesXml = addOverride(
      contentTypesXml,
      "word/footer2.xml",
      "application/vnd.openxmlformats-officedocument.wordprocessingml.footer+xml"
    )
    refs.footerFirst = "rIdFooter2"
  }

  archive["word/_rels/document.xml.rels"] = encodeXml(relsXml)
  archive["[Content_Types].xml"] = encodeXml(contentTypesXml)

  const documentXml = decodeXml(archive["word/document.xml"])
  archive["word/document.xml"] = encodeXml(
    patchDocumentSectPr(documentXml, refs, Boolean(pageSetup.headerFooterDifferentFirstPage))
  )

  return zipSync(archive)
}