"""IMM5709-only allowlisted XFA writer. Adobe signing/validation stays manual.

Scripts and lookup packets remain byte-identical. Template changes only selected
phone branch presence. Business datasets follow actual indexed template paths.
"""
from io import BytesIO
import hashlib
import json
from pathlib import Path
import re

from lxml import etree as E
import pikepdf

XFA_DATA = 'http://www.xfa.org/schema/xfa-data/1.0/'
FORM_NS = 'http://www.xfa.org/schema/xfa-form/2.8/'
PART = re.compile(r'([^\[\].]+)(?:\[(\d+)\])?')
PHONE_PREFIXES = tuple('form1.Page2.ContactInformation.' + path for path in (
    'q3-4.Phone.', 'q3-4.AltPhone.', 'q5-6.Fax.'))


def parse(raw):
    return E.parse(BytesIO(raw), E.XMLParser(remove_blank_text=False,
                   remove_comments=False, remove_pis=False, resolve_entities=False))


def local(node):
    return E.QName(node).localname if isinstance(node.tag, str) else ''


def parts(path):
    result = []
    for segment in path.split('.'):
        match = PART.fullmatch(segment)
        if not match:
            raise ValueError('Invalid IMM5709 path: ' + path)
        result.append((match[1], int(match[2]) if match[2] else None))
    return result


def resolve(root, path, named=False, create=False, kinds=None):
    """Exact local path; omitted indexes must be unambiguous."""
    parent = root
    for offset, (name, index) in enumerate(parts(path)):
        matches = [c for c in parent if (c.get('name') if named else local(c)) == name]
        if index is None and len(matches) > 1:
            raise ValueError('Ambiguous IMM5709 path: ' + path)
        selected = index if index is not None else 0
        if selected >= len(matches):
            if not create:
                raise ValueError('Unresolved IMM5709 target: ' + path)
            for _ in range(selected + 1 - len(matches)):
                if named:
                    node = E.SubElement(parent, '{' + FORM_NS + '}' + kinds[offset], name=name)
                else:
                    node = E.SubElement(parent, name)
                matches.append(node)
        parent = matches[selected]
    return parent


def form_value(node, value, value_type='text'):
    for old in list(node):
        if local(old) == 'value':
            node.remove(old)
    value_node = E.SubElement(node, '{' + FORM_NS + '}value')
    E.SubElement(value_node, '{' + FORM_NS + '}' + value_type).text = value or None


def checkbox_type(widget):
    if widget.xpath('./*[local-name()="ui"]/*[local-name()="checkButton"]'):
        exports = widget.xpath('./*[local-name()="items"]/*')
        if exports and local(exports[0]) in ('text', 'integer'):
            return local(exports[0])
    return 'text'


def template_chain(template, path):
    """Resolve every named ancestor and preserve actual widget kinds."""
    parent = template
    chain = []
    for name, index in parts(path):
        matches = [c for c in parent if c.get('name') == name]
        if index is None and len(matches) > 1:
            raise ValueError('Ambiguous IMM5709 template target: ' + path)
        selected = index if index is not None else 0
        if selected >= len(matches):
            raise ValueError('Unresolved IMM5709 template target: ' + path)
        parent = matches[selected]
        chain.append(parent)
    if local(parent) not in ('field', 'exclGroup'):
        raise ValueError('Unresolved IMM5709 widget target: ' + path)
    return chain


def write_state(form, path, chain, value):
    node = resolve(form, path, named=True, create=True, kinds=[local(n) for n in chain])
    form_value(node, value, checkbox_type(chain[-1]))
    if local(chain[-1]) == 'exclGroup':
        for widget in chain[-1]:
            if local(widget) != 'field' or not widget.get('name'):
                continue
            saved = next((c for c in node if c.get('name') == widget.get('name')), None)
            if saved is None:
                saved = E.SubElement(node, '{' + FORM_NS + '}field', name=widget.get('name'))
            exports = widget.xpath('./*[local-name()="items"]/*')
            export = (exports[0].text or '') if exports else ''
            form_value(saved, export if export == value else '', checkbox_type(widget))


def load_contract(language):
    return json.loads((Path(__file__).parent / f'imm5709-contract.{language}.json').read_text(encoding='utf-8'))


def fill_imm5709(payload, output_path, blank_path):
    from prefill_imm_xfa_pdf import release_edit_lock
    from imm5708_pdf_state import address_values, contact_address_state, contact_phone_state

    if payload.get('form') != 'IMM5709':
        raise ValueError('IMM5709 writer rejects other forms')
    language = payload.get('language', 'en')
    if language not in ('en', 'fr'):
        raise ValueError('Unsupported IMM5709 language')
    contract = load_contract(language)
    blank = Path(blank_path)
    if hashlib.sha256(blank.read_bytes()).hexdigest() != contract['sha256']:
        raise ValueError('IMM5709 template requires review')
    fields = payload.get('fields', {})
    if not isinstance(fields, dict):
        raise ValueError('IMM5709 payload.fields must be an object')
    for path, value in fields.items():
        meta = contract['fields'].get(path)
        if meta is None:
            raise ValueError('Unallowlisted IMM5709 field: ' + str(path))
        if not isinstance(value, str):
            raise ValueError('Invalid IMM5709 field value: ' + path)
        if len(value) > meta.get('max_chars', 2000):
            raise ValueError('IMM5709 field too long: ' + path)
        if value and meta.get('allowed_values') and value not in meta['allowed_values']:
            raise ValueError('Invalid IMM5709 export: ' + path)
        if value and 'choices' in meta and value not in meta['choices']:
            raise ValueError('Invalid IMM5709 choice: ' + path)
    # IMM5709 shares these widgets with 5708. Reuse pure state transforms only
    # inside this dispatch; templates and other form writers remain unchanged.
    prefix = 'form1.Page2.'
    contact = {k[len(prefix):]: v for k, v in fields.items() if k.startswith(prefix)}
    contact = address_values(contact)
    fields = dict(fields)
    fields.update({prefix + k: v for k, v in contact.items()})
    for path in fields:
        if path not in contract['fields']:
            raise ValueError('Unallowlisted IMM5709 state target: ' + path)
    with pikepdf.open(blank) as pdf:
        xfa = pdf.Root.AcroForm.XFA
        streams = {str(xfa[i]): xfa[i + 1] for i in range(0, len(xfa), 2)}
        template_tree = parse(streams['template'].read_bytes())
        template = template_tree.getroot()
        # Resolve all payload targets before making changes or writing output.
        plans = [(path, value, contract['fields'][path]['template_path'],
                  template_chain(template, contract['fields'][path]['template_path']))
                 for path, value in fields.items()]
        datasets_tree = parse(streams['datasets'].read_bytes())
        datasets = datasets_tree.getroot()
        data = datasets.find('{' + XFA_DATA + '}data')
        if data is None:
            data = E.SubElement(datasets, '{' + XFA_DATA + '}data')
        form_tree = parse(streams['form'].read_bytes()) if 'form' in streams else None
        form = form_tree.getroot() if form_tree is not None else None
        for path, value, saved_path, chain in plans:
            target = resolve(data, path, create=True)
            for old in list(target):
                target.remove(old)
            target.text = value or None
            is_phone = path.startswith(PHONE_PREFIXES)
            if is_phone:
                target.attrib.pop('{' + XFA_DATA + '}dataNode', None)
            else:
                target.set('{' + XFA_DATA + '}dataNode', 'dataValue')
            if form is not None:
                if is_phone:
                    # Official docReady restores saved deltas. Bound phone data
                    # must merge from datasets, like the working IMM5708 writer.
                    node = resolve(form, saved_path, named=True, create=True,
                                   kinds=[local(n) for n in chain])
                    for old in list(node):
                        if local(old) == 'value':
                            node.remove(old)
                else:
                    write_state(form, saved_path, chain, value)
        # Blank can contain stale signature/validated values. Never date drafts.
        for node in list(data.iter()):
            if local(node) in ('ApplicationValidatedFlag', 'FormValidated', 'ApplicationValidationDate', 'DateLastValidated') or (
                local(node) in ('TextField2', 'C1CertificateIssueDate') and node.getparent() is not None and local(node.getparent()) == 'Signature'):
                for descendant in node.iter():
                    descendant.text = None
        if form is not None:
            for node in list(form.iter()):
                if node.get('name') in ('ApplicationValidatedFlag', 'FormValidated', 'ApplicationValidationDate', 'DateLastValidated') or (
                    node.get('name') in ('TextField2', 'C1CertificateIssueDate') and node.getparent() is not None and node.getparent().get('name') == 'Signature'):
                    for descendant in list(node.iter()):
                        if local(descendant) == 'value':
                            descendant.getparent().remove(descendant)
                    for descendant in list(node.iter()):
                        if local(descendant) in ('field', 'exclGroup'):
                            form_value(descendant, '')
            saved = contact_address_state(E.tostring(form_tree), contact)
            saved = contact_phone_state(saved, contact)
            streams['form'].write(saved)
        streams['datasets'].write(E.tostring(datasets_tree))
        # Presence in saved state alone is discarded during Adobe's merge.
        # Byte patch only phone branches; original scripts and PIs stay intact.
        streams['template'].write(contact_phone_state(streams['template'].read_bytes(), contact))
        release_edit_lock(pdf)
        output = Path(output_path)
        output.parent.mkdir(parents=True, exist_ok=True)
        pdf.save(output)
    return output