#!/usr/bin/env python3
"""Embed the Deeper Signals brand fonts into a generated booklet docx.

Why this is needed
==================
docx-js (the JS library we use to generate the document body) does not
write font embedding entries. Without embedded fonts, Word users on
machines that don't have Manrope / Space Grotesk / Roboto installed will
see the document with substituted fonts and the layout will look wrong.

The original Deeper Signals booklet template embeds the TTF files inside
the docx package. This script recreates that embedding so client-facing
documents render consistently on any machine.

The embedding format follows Word's documented obfuscation: the first 32
bytes of the TTF are XOR-ed with a 16-byte GUID, repeated. We write a
deterministic zero GUID (which mirrors what the original template uses)
so the obfuscation is effectively a no-op but Word still treats the font
as a valid embedded TTF.

Usage:
    python embed_fonts.py <input.docx> <fonts_dir> [output.docx]

If output.docx is omitted, the input file is overwritten in place.
"""

from __future__ import annotations

import argparse
import os
import shutil
import sys
import tempfile
import zipfile
from pathlib import Path

# Font family -> (regular.ttf, bold.ttf) inside the fonts_dir.
# The original template embeds: Manrope, Manrope Medium, Manrope SemiBold,
# Space Grotesk Medium, Space Grotesk SemiBold, and Roboto (with italic +
# bold-italic variants for Roboto).
FONT_FAMILIES = [
    {
        "name": "Roboto",
        "regular": "Roboto-regular.ttf",
        "bold": "Roboto-bold.ttf",
        "italic": "Roboto-italic.ttf",
        "boldItalic": "Roboto-boldItalic.ttf",
    },
    {
        "name": "Manrope",
        "regular": "Manrope-regular.ttf",
        "bold": "Manrope-bold.ttf",
    },
    {
        "name": "Manrope Medium",
        "regular": "ManropeMedium-regular.ttf",
        "bold": "ManropeMedium-bold.ttf",
    },
    {
        "name": "Manrope SemiBold",
        "regular": "ManropeSemiBold-regular.ttf",
        "bold": "ManropeSemiBold-bold.ttf",
    },
    {
        "name": "Space Grotesk Medium",
        "regular": "SpaceGroteskMedium-regular.ttf",
        "bold": "SpaceGroteskMedium-bold.ttf",
    },
    {
        "name": "Space Grotesk SemiBold",
        "regular": "SpaceGroteskSemiBold-regular.ttf",
        "bold": "SpaceGroteskSemiBold-bold.ttf",
    },
]

# Zero GUID: the obfuscation XOR mask is all zeros, so the obfuscated TTF
# bytes are identical to the original. Word still accepts this as valid.
FONT_KEY = "{00000000-0000-0000-0000-000000000000}"


def obfuscate_ttf(ttf_bytes: bytes, guid: str = FONT_KEY) -> bytes:
    """Apply the Word font obfuscation transform (XOR first 32 bytes).

    With a zero GUID the transform is a no-op, which is what the original
    template uses. The function is parameterised so a real GUID can be
    plugged in later if a non-shared key is ever required.
    """
    # Parse 16 hex bytes from the GUID, reverse order per Word spec
    hex_str = guid.replace("{", "").replace("}", "").replace("-", "")
    raw = bytes.fromhex(hex_str)
    # Word reads bytes in reverse for each segment; for the zero GUID this
    # doesn't matter but kept here for correctness.
    mask = raw[::-1]

    out = bytearray(ttf_bytes)
    for i in range(min(32, len(out))):
        out[i] ^= mask[i % 16]
    return bytes(out)


def build_font_table_xml(fonts_meta: list[dict]) -> str:
    """Render the fontTable.xml that pairs font names with relationship ids."""
    parts = [
        '<?xml version="1.0" encoding="utf-8"?>',
        '<w:fonts xmlns:r="http://schemas.openxmlformats.org/officeDocument/2006/relationships"'
        ' xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">',
        '  <w:font w:name="Arial"/>',
    ]
    for fm in fonts_meta:
        parts.append(f'  <w:font w:name="{fm["name"]}">')
        if "regular_rid" in fm:
            parts.append(
                f'    <w:embedRegular w:fontKey="{FONT_KEY}" r:id="{fm["regular_rid"]}" w:subsetted="0"/>'
            )
        if "bold_rid" in fm:
            parts.append(
                f'    <w:embedBold w:fontKey="{FONT_KEY}" r:id="{fm["bold_rid"]}" w:subsetted="0"/>'
            )
        if "italic_rid" in fm:
            parts.append(
                f'    <w:embedItalic w:fontKey="{FONT_KEY}" r:id="{fm["italic_rid"]}" w:subsetted="0"/>'
            )
        if "boldItalic_rid" in fm:
            parts.append(
                f'    <w:embedBoldItalic w:fontKey="{FONT_KEY}" r:id="{fm["boldItalic_rid"]}" w:subsetted="0"/>'
            )
        parts.append("  </w:font>")
    parts.append("</w:fonts>")
    return "\n".join(parts)


def build_font_table_rels_xml(rel_entries: list[tuple[str, str]]) -> str:
    """Render fontTable.xml.rels pairing rIds to embedded font filenames."""
    parts = [
        '<?xml version="1.0" encoding="utf-8"?>',
        '<Relationships xmlns="http://schemas.openxmlformats.org/package/2006/relationships">',
    ]
    for rid, target in rel_entries:
        parts.append(
            f'  <Relationship Id="{rid}"'
            f' Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/font"'
            f' Target="{target}"/>'
        )
    parts.append("</Relationships>")
    return "\n".join(parts)


def ensure_embed_setting(settings_bytes: bytes) -> bytes:
    """Insert <w:embedTrueTypeFonts w:val="1"/> into settings.xml if absent.

    The element is positioned according to the OOXML schema. The schema
    sequence places embedSystemFonts / embedTrueTypeFonts immediately
    after displayBackgroundShape and before saveSubsetFonts. We place
    the element either right after displayBackgroundShape (if present)
    or as the first child of the settings root.
    """
    text = settings_bytes.decode("utf-8")
    if "embedTrueTypeFonts" in text:
        return settings_bytes

    inject = '\n  <w:embedTrueTypeFonts w:val="1"/>'

    # Prefer insertion immediately after <w:displayBackgroundShape/>.
    marker = "<w:displayBackgroundShape/>"
    idx = text.find(marker)
    if idx != -1:
        anchor = idx + len(marker)
        return (text[:anchor] + inject + text[anchor:]).encode("utf-8")

    # Fall back: insert right after the opening <w:settings ...> tag.
    open_tag_end = text.find(">", text.find("<w:settings"))
    if open_tag_end == -1:
        return settings_bytes
    return (text[: open_tag_end + 1] + inject + text[open_tag_end + 1 :]).encode("utf-8")


def ensure_font_table_rel(rels_bytes: bytes) -> bytes:
    """Make sure the main document rels file links to fontTable.xml.

    docx-js typically already includes a fontTable relationship; if it
    doesn't we add one so the embedded fonts can be located by Word.
    """
    text = rels_bytes.decode("utf-8")
    if "fontTable.xml" in text:
        return rels_bytes
    # Find an unused rId number.
    used = []
    for tok in text.split('Id="'):
        if tok and tok[0:3] == "rId":
            try:
                used.append(int(tok[3:].split('"', 1)[0]))
            except ValueError:
                pass
    next_id = (max(used) + 1) if used else 100
    new_rel = (
        f'  <Relationship Id="rId{next_id}"'
        f' Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/fontTable"'
        f' Target="fontTable.xml"/>\n'
    )
    insert_at = text.rfind("</Relationships>")
    return (text[:insert_at] + new_rel + text[insert_at:]).encode("utf-8")


def embed_fonts(input_path: Path, fonts_dir: Path, output_path: Path) -> None:
    if not input_path.exists():
        raise FileNotFoundError(input_path)
    if not fonts_dir.exists():
        raise FileNotFoundError(fonts_dir)

    with tempfile.TemporaryDirectory() as tmp:
        tmp_dir = Path(tmp)
        # Extract original zip
        with zipfile.ZipFile(input_path, "r") as zin:
            zin.extractall(tmp_dir)

        # Build font binary blobs and metadata
        fonts_meta: list[dict] = []
        rel_entries: list[tuple[str, str]] = []
        rid_counter = 1
        embedded_dir = tmp_dir / "word" / "fonts"
        embedded_dir.mkdir(parents=True, exist_ok=True)

        for family in FONT_FAMILIES:
            meta: dict = {"name": family["name"]}
            for variant_key, rid_key in [
                ("regular", "regular_rid"),
                ("bold", "bold_rid"),
                ("italic", "italic_rid"),
                ("boldItalic", "boldItalic_rid"),
            ]:
                if variant_key not in family:
                    continue
                src = fonts_dir / family[variant_key]
                if not src.exists():
                    # Skip silently if a particular variant is absent; the
                    # font will still embed in its other variants.
                    continue
                rid = f"rId{rid_counter}"
                rid_counter += 1
                target_name = f"font{rid_counter}.odttf"
                obf = obfuscate_ttf(src.read_bytes())
                (embedded_dir / target_name).write_bytes(obf)
                meta[rid_key] = rid
                rel_entries.append((rid, f"fonts/{target_name}"))
            fonts_meta.append(meta)

        # Write fontTable.xml and its rels file
        (tmp_dir / "word" / "fontTable.xml").write_text(
            build_font_table_xml(fonts_meta), encoding="utf-8"
        )
        rels_dir = tmp_dir / "word" / "_rels"
        rels_dir.mkdir(exist_ok=True)
        (rels_dir / "fontTable.xml.rels").write_text(
            build_font_table_rels_xml(rel_entries), encoding="utf-8"
        )

        # Patch settings.xml to enable font embedding
        settings_path = tmp_dir / "word" / "settings.xml"
        if settings_path.exists():
            settings_path.write_bytes(ensure_embed_setting(settings_path.read_bytes()))
        else:
            settings_path.write_text(
                '<?xml version="1.0" encoding="utf-8"?>\n'
                '<w:settings xmlns:w="http://schemas.openxmlformats.org/wordprocessingml/2006/main">\n'
                '  <w:embedTrueTypeFonts w:val="1"/>\n'
                "</w:settings>",
                encoding="utf-8",
            )

        # Make sure the document rels point at the fontTable
        doc_rels_path = tmp_dir / "word" / "_rels" / "document.xml.rels"
        if doc_rels_path.exists():
            doc_rels_path.write_bytes(ensure_font_table_rel(doc_rels_path.read_bytes()))

        # Ensure Content_Types declares the obfuscated font extension
        ct_path = tmp_dir / "[Content_Types].xml"
        ct_text = ct_path.read_text(encoding="utf-8")
        if "odttf" not in ct_text:
            inject = '<Default Extension="odttf" ContentType="application/vnd.openxmlformats-officedocument.obfuscatedFont"/>'
            # Place after the <Types ...> opening tag
            open_end = ct_text.find(">", ct_text.find("<Types")) + 1
            ct_text = ct_text[:open_end] + inject + ct_text[open_end:]
            ct_path.write_text(ct_text, encoding="utf-8")

        # Repack everything as a docx
        with zipfile.ZipFile(output_path, "w", zipfile.ZIP_DEFLATED) as zout:
            for root, _, files in os.walk(tmp_dir):
                for f in files:
                    full = Path(root) / f
                    arc = full.relative_to(tmp_dir).as_posix()
                    zout.write(full, arc)


def main() -> int:
    p = argparse.ArgumentParser(description="Embed Deeper Signals brand fonts into a generated booklet.")
    p.add_argument("input", type=Path, help="Path to the generated docx.")
    p.add_argument("fonts_dir", type=Path, help="Directory containing TTF font files.")
    p.add_argument("output", type=Path, nargs="?", default=None, help="Output docx path (default: overwrite input).")
    args = p.parse_args()

    target = args.output or args.input
    if target == args.input:
        # Write to a temp file then move, so a mid-flight failure doesn't
        # corrupt the original.
        with tempfile.NamedTemporaryFile(delete=False, suffix=".docx") as tf:
            tmp_out = Path(tf.name)
        embed_fonts(args.input, args.fonts_dir, tmp_out)
        shutil.move(str(tmp_out), str(target))
    else:
        embed_fonts(args.input, args.fonts_dir, target)

    print(f"Embedded fonts into {target}")
    return 0


if __name__ == "__main__":
    sys.exit(main())
