#!/usr/bin/env python3
"""Fill IRCC XFA (LiveCycle) IMM PDFs by writing the datasets XML packet.

Acrobat Reader still has to open the file and click Validate for 2D barcodes.
This script does not flatten the form and does not run Adobe JavaScript.
"""
from __future__ import annotations

import json
import re
import sys
import xml.etree.ElementTree as ET
from pathlib import Path


MAX_FIELD_CHARS = 2000


def local_name(tag: str) -> str:
    if "}" in tag:
        return tag.rsplit("}", 1)[-1]
    return tag


XFA_DATA_NS = "http://www.xfa.org/schema/xfa-data/1.0/"


def set_element_text(el: ET.Element, value: str) -> None:
    el.text = value
    for child in list(el):
        el.remove(child)
    # IMM 5707 YES/NO is an exclGroup. Mark the node as a single value so
    # Acrobat selects the checkbox instead of treating it as an empty group.
    if local_name(el.tag) == "yesno" and value in {"1", "2"}:
        el.set(f"{{{XFA_DATA_NS}}}dataNode", "dataValue")


def parse_path_part(part: str) -> tuple[str, int | None]:
    match = re.fullmatch(r"([^\[\]]+)\[(\d+)\]", part)
    if match:
        return match.group(1), int(match.group(2))
    return part, None


def children_named(node: ET.Element, name: str) -> list[ET.Element]:
    want = name.lower()
    return [child for child in list(node) if local_name(child.tag).lower() == want]


def _resolve_from(node: ET.Element, parts: list[str]) -> list[ET.Element]:
    """Follow one path. An omitted index uses the first sibling that can finish it.

    IMM 5707 repeats ``PersonalData``: names sit in the first group and DOB,
    country, address, and occupation sit in the second. Always taking sibling 0
    misses that second group. Trying every sibling would also stamp
    ``SectionB.Child`` onto every child row, so the search stops at the first
    sibling whose remainder matches.
    """
    if not parts:
        return [node]
    name, idx = parse_path_part(parts[0])
    named = children_named(node, name)
    if not named:
        return []
    if idx is not None:
        if idx < 0 or idx >= len(named):
            return []
        return _resolve_from(named[idx], parts[1:])
    for candidate in named:
        found = _resolve_from(candidate, parts[1:])
        if found:
            return found
    return []


def find_by_path(root: ET.Element, path: str) -> list[ET.Element]:
    skip = {"datasets", "data", "xfa:datasets", "xfa:data"}
    parts = [p for p in path.replace("/", ".").split(".") if p]
    while parts and parse_path_part(parts[0])[0].lower() in skip:
        parts = parts[1:]
    if not parts:
        return []

    first_name, first_idx = parse_path_part(parts[0])
    starts = [el for el in root.iter() if local_name(el.tag).lower() == first_name.lower()]
    if first_idx is not None:
        starts = starts[first_idx : first_idx + 1]
    matches: list[ET.Element] = []
    for start in starts:
        matches.extend(_resolve_from(start, parts[1:]))
    return matches


def find_by_local_name(root: ET.Element, name: str) -> list[ET.Element]:
    want = name.lower()
    return [el for el in root.iter() if local_name(el.tag).lower() == want]


def release_edit_lock(pdf) -> None:
    """Drop certification locks on the filled copy so Adobe can edit answers.

    Official blanks are DocMDP-certified and Reader-enabled (UR3). Rewriting
    the XFA packet outside Adobe invalidates those signatures, and Reader then
    paints the values as static text. Removing the signatures from this copy
    releases that lock. The blank file on disk is left unchanged.
    """
    root = pdf.Root
    if "/Perms" in root:
        del root["/Perms"]

    acro = root.get("/AcroForm")
    if acro is None:
        return
    if "/SigFlags" in acro:
        del acro["/SigFlags"]

    def clear_signature(field) -> None:
        field_type = str(field.get("/FT", ""))
        field_name = str(field.get("/T", ""))
        if field_type == "/Sig" or "Signature" in field_name:
            for key in ("/V", "/Lock", "/SV"):
                if key in field:
                    del field[key]
        kids = field.get("/Kids")
        if kids:
            for kid in kids:
                clear_signature(kid)

    for field in list(acro.get("/Fields", [])):
        clear_signature(field)

    # PDF.js treats any remaining Fields as AcroForm after SigFlags is removed,
    # so it never loads the XFA pages. Remove only invisible certification
    # widgets from dynamic XFA drafts; retain visible signature/answer fields.
    if acro.get("/XFA") is None or not root.get("/NeedsRendering", False):
        return
    import pikepdf

    removed_widgets = set()

    def retain_field(field, inherited_type="") -> bool:
        field_type = str(field.get("/FT", inherited_type))
        kids = field.get("/Kids")
        if kids is not None and len(kids):
            retained = [kid for kid in kids if retain_field(kid, field_type)]
            field["/Kids"] = pikepdf.Array(retained)
            return bool(retained)
        rect = field.get("/Rect")
        if field_type == "/Sig" and rect is not None and len(rect) == 4:
            if rect[0] == rect[2] or rect[1] == rect[3]:
                removed_widgets.add(field.objgen)
                return False
        return True

    fields = acro.get("/Fields")
    if fields is not None:
        acro["/Fields"] = pikepdf.Array(field for field in fields if retain_field(field))
    if removed_widgets:
        for page in pdf.pages:
            annotations = page.obj.get("/Annots")
            if annotations is not None:
                page.obj["/Annots"] = pikepdf.Array(
                    annotation for annotation in annotations
                    if annotation.objgen not in removed_widgets
                )


def hide_javascript_banner(xml_bytes: bytes) -> bytes:
    """Hide the IRCC template banner that appears when Acrobat JS does not run.

    Official blanks show JavaScriptMessage until form-ready JS hides it. Prefill
    cannot run that script, so the filled draft keeps the red banner unless the
    template presence is set here. This does not enable Validate or barcodes.
    """
    pattern = re.compile(
        br'<(draw|field)\b([^>]*\bname="JavaScriptMessage"[^>]*)>',
        re.IGNORECASE | re.DOTALL,
    )

    def repl(match: re.Match[bytes]) -> bytes:
        inner = match.group(2)
        if re.search(br'\bpresence=', inner, re.I):
            inner = re.sub(br'\bpresence="[^"]*"', b'presence="hidden"', inner, count=1, flags=re.I)
        else:
            inner = inner.rstrip() + b' presence="hidden"'
        return b'<' + match.group(1) + inner + b'>'

    return pattern.sub(repl, xml_bytes)


def yesno_items_as_text(xml_bytes: bytes) -> bytes:
    """Store YES/NO on-values as text, matching the datasets strings we write.

    These checkboxes are authored as integer items (1 = YES, 2 = NO). A
    datasets text node then does not select the box. IMM 1294 checkboxes,
    which do tick, store the same kind of on-value as text.
    """
    yes = re.compile(
        br"(YES</text\s*>\s*</value\s*>\s*</caption\s*>\s*<items\s*>\s*)<integer\s*>1</integer\s*>",
        re.IGNORECASE,
    )
    no = re.compile(
        br"(NO</text\s*>\s*</value\s*>\s*</caption\s*>\s*<items\s*>\s*)<integer\s*>2</integer\s*>",
        re.IGNORECASE,
    )
    xml_bytes = yes.sub(br"\1<text>1</text>", xml_bytes)
    return no.sub(br"\1<text>2</text>", xml_bytes)


# Template order of IMM 5707 exclGroup name="yesno". The child one is shared by
# every child row, so a template value would tick all four rows the same way.
_TEMPLATE_YESNO_PATHS = (
    "SectionA.Applicant.PaddedEntry.MarriageInPerson.MarriageInPerson.yesno",
    "SectionA.Spouse.PaddedEntry.Accompanying.yesno",
    "SectionA.Spouse.PaddedEntry.MarriageInPerson.MarriageInPerson.yesno",
    "SectionA.Parent1.PaddedEntry.Accompanying.yesno",
    "SectionA.Parent2.PaddedEntry.Accompanying.yesno",
    "SectionB.Child.PaddedEntry.Accompanying.yesno",
)


def _yesno_choice(fields: dict[str, str], path: str) -> str | None:
    raw = fields.get(path)
    value = "" if raw is None else str(raw).strip()
    if value in {"1", "2"}:
        return value
    return None


def _som_yesno(path: str) -> str:
    parts: list[str] = []
    for part in path.split("."):
        if not part or part == "page1":
            continue
        name, index = parse_path_part(part)
        parts.append(f"{name}[{0 if index is None else index}]")
    return "page1." + ".".join(parts)


def stamp_template_yesno(xml_bytes: bytes, fields: dict[str, str]) -> bytes:
    """Select Section A YES/NO on the template widget Acrobat actually draws.

    Those boxes are not in the saved form state. Adding a new Section A there
    does not attach to the widget. The child rows are already in that state,
    so they are left to the form packet.
    """
    starts: list[int] = []
    cursor = 0
    while True:
        found = xml_bytes.find(b'name="yesno"', cursor)
        if found < 0:
            break
        starts.append(xml_bytes.rfind(b"<exclGroup", 0, found))
        cursor = found + 1
    if len(starts) != len(_TEMPLATE_YESNO_PATHS):
        return xml_bytes

    inserts: list[tuple[int, bytes]] = []
    for index, path in enumerate(_TEMPLATE_YESNO_PATHS):
        if path.startswith("SectionB.Child"):
            continue
        value = _yesno_choice(fields, path)
        if value is None:
            continue
        excl_start = starts[index]
        group_end = xml_bytes.find(b"</exclGroup", excl_start)
        if group_end < 0:
            continue
        member = b"yes" if value == "1" else b"no"
        marker = xml_bytes.find(b'name="' + member + b'"', excl_start, group_end)
        if marker < 0:
            continue
        blob = b"<value\n><text\n>" + value.encode("ascii") + b"</text\n></value\n>"
        inserts.append((xml_bytes.find(b">", excl_start) + 1, blob))
        inserts.append((xml_bytes.find(b">", marker) + 1, blob))

    for offset, blob in sorted(inserts, key=lambda item: item[0], reverse=True):
        xml_bytes = xml_bytes[:offset] + blob + xml_bytes[offset:]
    return xml_bytes


def retick_yesno_on_open(xml_bytes: bytes, fields: dict[str, str]) -> bytes:
    """Set each YES/NO after Acrobat restores the empty checkbox state."""
    root = b'<subform name="IMM_5707" layout="tb" locale="en_CA" restoreState="auto"\n>'
    if root not in xml_bytes or b'name="prefillYesNo"' in xml_bytes:
        return xml_bytes
    lines: list[str] = []
    for key, raw in fields.items():
        if not isinstance(key, str) or not key.endswith("yesno"):
            continue
        value = _yesno_choice({key: raw}, key)
        if value is None:
            continue
        som = _som_yesno(key)
        lines.append(
            f'var n = this.resolveNode("{som}");'
            f' if (n == null) n = xfa.resolveNode("IMM_5707.{som}");'
            f' if (n != null) n.rawValue = "{value}";'
        )
    if not lines:
        return xml_bytes
    script = "".join(lines).replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
    event = (
        b'<event activity="docReady" name="prefillYesNo">'
        b'<script contentType="application/x-javascript">'
        + script.encode("ascii")
        + b"</script></event>"
    )
    return xml_bytes.replace(root, root + event, 1)


FORM_NS = "http://www.xfa.org/schema/xfa-form/2.8/"


def _form_tag(name: str) -> str:
    return f"{{{FORM_NS}}}{name}"


def _form_named(parent: ET.Element, kind: str, name: str) -> ET.Element | None:
    for child in parent:
        if local_name(child.tag) == kind and child.get("name") == name:
            return child
    return None


def _ensure_form_subform(parent: ET.Element, name: str, index: int) -> ET.Element | None:
    kids = [
        child for child in parent
        if local_name(child.tag) == "subform" and child.get("name") == name
    ]
    if index < len(kids):
        return kids[index]
    # Page sections and child rows are already instantiated. A new Section A
    # here does not bind to the boxes on the page. Nested pieces under an
    # existing child row do, which is how that row's tick is stored.
    if name in {"SectionA", "SectionB", "SectionC", "Instructions", "Child"}:
        return None
    while len(kids) <= index:
        manager = "_"+name
        if not any(local_name(child.tag) == "instanceManager" and child.get("name") == manager for child in parent):
            instance = ET.Element(_form_tag("instanceManager"))
            instance.set("name", manager)
            parent.append(instance)
        created = ET.Element(_form_tag("subform"))
        created.set("name", name)
        parent.append(created)
        kids.append(created)
    return kids[index]


def _set_form_yesno(parent: ET.Element, value: str) -> None:
    group = _form_named(parent, "exclGroup", "yesno")
    if group is None:
        group = ET.SubElement(parent, _form_tag("exclGroup"))
        group.set("name", "yesno")
    for child in list(group):
        if local_name(child.tag) in {"value", "field"}:
            group.remove(child)
    chosen = ET.SubElement(group, _form_tag("value"))
    text = ET.SubElement(chosen, _form_tag("text"))
    text.text = value
    for member, on_value in (("yes", "1"), ("no", "2")):
        field = ET.SubElement(group, _form_tag("field"))
        field.set("name", member)
        if value != on_value:
            continue
        member_value = ET.SubElement(field, _form_tag("value"))
        member_text = ET.SubElement(member_value, _form_tag("text"))
        member_text.text = on_value


def apply_form_yesno(xml_bytes: bytes, fields: dict[str, str]) -> bytes:
    """Persist YES/NO in the saved form state Acrobat restores on open."""
    selected = []
    for key, raw in fields.items():
        if not isinstance(key, str) or not key.endswith("yesno"):
            continue
        value = "" if raw is None else str(raw).strip()
        if value in {"1", "2"}:
            selected.append((key, value))
    if not selected:
        return xml_bytes

    ET.register_namespace("", FORM_NS)
    root = ET.fromstring(xml_bytes)
    if "checksum" in root.attrib:
        del root.attrib["checksum"]
    host = _form_named(root, "subform", "IMM_5707")
    page = None if host is None else _form_named(host, "subform", "page1")
    if page is None:
        return xml_bytes

    for path, value in selected:
        node = page
        for part in [p for p in path.split(".") if p and p != "page1"]:
            name, index = parse_path_part(part)
            if name == "yesno":
                _set_form_yesno(node, value)
                break
            node = _ensure_form_subform(node, name, 0 if index is None else index)
            if node is None:
                break

    return ET.tostring(root, encoding="utf-8", xml_declaration=False)


def apply_values(xml_bytes: bytes, fields: dict[str, str], scalar_fields=()) -> bytes:
    if not fields:
        return xml_bytes

    ET.register_namespace("xfa", "http://www.xfa.org/schema/xfa-data/1.0/")

    # Preserve XML declaration if present.
    declaration = b""
    stripped = xml_bytes.lstrip()
    if stripped.startswith(b"<?xml"):
        end = stripped.find(b"?>")
        if end != -1:
            declaration = stripped[: end + 2] + b"\n"

    root = ET.fromstring(xml_bytes)
    for key, raw in fields.items():
        if not isinstance(key, str) or key.strip() == "":
            continue
        value = "" if raw is None else str(raw)
        if len(value) > MAX_FIELD_CHARS:
            raise ValueError(f"Field {key} exceeds draft writer limit of {MAX_FIELD_CHARS} characters; review continuation requirements")

        if "." in key or "/" in key:
            targets = find_by_path(root, key)
        else:
            targets = find_by_local_name(root, key)
        if not targets:
            raise ValueError(f"Mapped field {key} not found in this template; mapping requires review")
        for el in targets:
            set_element_text(el, value)
            if key in scalar_fields:
                el.set(f"{{{XFA_DATA_NS}}}dataNode", "dataValue")

    body = ET.tostring(root, encoding="utf-8", xml_declaration=False)
    if declaration:
        return declaration + body
    return body


def fill(payload_path: str, output_path: str, blank_path: str) -> Path:
    import pikepdf

    payload = json.loads(Path(payload_path).read_text(encoding="utf-8-sig"))
    if payload.get("form") == "IMM5009":
        from imm5009_pdf import fill_document_request
        return fill_document_request(payload, output_path, blank_path)
    if payload.get("form") == "IMM5283":
        from imm5283_pdf import fill_humanitarian
        return fill_humanitarian(payload, output_path, blank_path)
    if payload.get("form") == "IMM5532":
        from imm5532_pdf import fill_relationship
        return fill_relationship(payload, output_path, blank_path)
    if payload.get("form") == "IMM1344":
        from imm1344_pdf import fill_sponsorship
        return fill_sponsorship(payload, output_path, blank_path)
    if payload.get("form") == "IMM5562":
        from imm5562_pdf import fill_travels
        return fill_travels(payload, output_path, blank_path)
    if payload.get("form") == "IMM5406":
        from imm5406_pdf import fill_family_information
        return fill_family_information(payload, output_path, blank_path)
    if payload.get("form") == "IMM5669":
        from imm5669_pdf import fill_schedule_a
        return fill_schedule_a(payload, output_path, blank_path)
    if payload.get("form") == "IMM0008":
        from imm0008_pdf import fill_imm0008
        return fill_imm0008(payload, output_path, blank_path)
    if payload.get("form") == "IMM5710":
        from imm5710_pdf import fill_imm5710
        return fill_imm5710(payload, output_path, blank_path)
    if payload.get("form") == "IMM5709":
        from imm5709_pdf import fill_imm5709
        return fill_imm5709(payload, output_path, blank_path)
    if payload.get("form") == "IMM5257SCH1":
        from imm5257sch1_pdf import fill_schedule
        return fill_schedule(payload, output_path, blank_path)
    fields = payload.get("fields") or {}
    if not isinstance(fields, dict):
        raise ValueError("payload.fields must be an object")

    blank = Path(blank_path)
    if not blank.is_file():
        raise FileNotFoundError(f"Missing blank PDF: {blank}")

    try:
        pdf = pikepdf.open(blank, password="")
    except pikepdf.PasswordError:
        pdf = pikepdf.open(blank)

    try:
        acro = pdf.Root.get("/AcroForm")
        if acro is None:
            raise ValueError("PDF has no AcroForm (not a fillable IMM XFA form)")
        xfa = acro.get("/XFA")
        if xfa is None:
            raise ValueError("PDF has no XFA packet (AcroForm-only; not an IRCC LiveCycle form)")

        datasets_stream = None
        form_stream = None
        normalized = {str(k): "" if v is None else str(v) for k, v in fields.items()}
        if payload.get("form") == "IMM1294":
            from imm1294_pdf_state import spouse_values
            normalized = spouse_values(normalized)
        if payload.get("form") == "IMM5708":
            from imm5708_pdf_state import address_values, contact_address_state, ADDRESS_SELECTION
            normalized = address_values(normalized)
        if payload.get("form") == "IMM1295":
            from imm1295_pdf_state import address_values, contact_address_state, ADDRESS_SELECTION
            normalized = address_values(normalized, payload.get('language', 'en'))
        for i in range(0, len(xfa), 2):
            name = str(xfa[i]).strip("/").lower()
            stream = xfa[i + 1]
            if name in {"template", "xfa:template"}:
                template = hide_javascript_banner(stream.read_bytes())
                template = yesno_items_as_text(template)
                template = stamp_template_yesno(template, normalized)
                template = retick_yesno_on_open(template, normalized)
                if payload.get("form") == "IMM5708":
                    from imm5708_pdf_state import contact_phone_state
                    template = contact_phone_state(template, normalized)
                    template = contact_address_state(template, normalized)
                if payload.get("form") == "IMM1294":
                    from imm1294_pdf_state import contact_phone_state
                    template = contact_phone_state(template, normalized)
                if payload.get("form") == "IMM1295":
                    from imm1295_pdf_state import contact_phone_state
                    template = contact_phone_state(template, normalized)
                    template = contact_address_state(template, normalized)
                stream.write(template)
            elif name in {"datasets", "xfa:datasets"}:
                datasets_stream = stream
            elif name == "form":
                form_stream = stream
        if datasets_stream is None:
            raise ValueError("XFA datasets packet was not found")

        original = datasets_stream.read_bytes()
        written_values = normalized
        if payload.get("form") == "IMM5476":
            # Official blanks contain XHTML placeholders in these names. Clear only
            # unfilled names; keep payload/coverage limited to actual answer values.
            written_values = dict(normalized)
            for target in ("SectionA.familyName", "SectionB.familyName", "sectionC.familyName", "sectionD.familyName"):
                written_values.setdefault(target, "")
        scalar_fields = (ADDRESS_SELECTION,) if payload.get("form") in {"IMM5708", "IMM1295"} and normalized.get(ADDRESS_SELECTION) in {"Y", "N"} else ()
        datasets = apply_values(original, written_values, scalar_fields)
        datasets_stream.write(datasets)
        if form_stream is not None:
            form_state = apply_form_yesno(form_stream.read_bytes(), normalized)
            if payload.get("form") == "IMM5708":
                from imm5708_pdf_state import contact_phone_state
                form_state = contact_phone_state(form_state, normalized)
                form_state = contact_address_state(form_state, normalized)
            if payload.get("form") == "IMM1294":
                from imm1294_pdf_state import contact_phone_state
                form_state = contact_phone_state(form_state, normalized)
            if payload.get("form") == "IMM1295":
                from imm1295_pdf_state import contact_phone_state
                form_state = contact_phone_state(form_state, normalized)
                form_state = contact_address_state(form_state, normalized)
            form_stream.write(form_state)
        release_edit_lock(pdf)

        out = Path(output_path)
        out.parent.mkdir(parents=True, exist_ok=True)
        pdf.save(out)
        return out
    finally:
        pdf.close()


if __name__ == "__main__":
    if len(sys.argv) < 4:
        print("Usage: python prefill_imm_xfa_pdf.py <payload.json> <output.pdf> <blank.pdf>")
        raise SystemExit(2)
    print(fill(sys.argv[1], sys.argv[2], sys.argv[3]))