"""IMM5406-only incremental writer with allowlisted dynamic family rows."""
import hashlib
import json
import os
from pathlib import Path
import re
from lxml import etree as E

PART = re.compile(r'([^\[\].]+)(?:\[(\d+)\])?')


def restore_partner_on_open(form, contract, fields, state, namespace):
    """Replay the official exit event after host merge, without editing template.

    Setting a dataset value does not fire the marital-status exit handler.
    Reader can recreate a visible Spouse despite its saved presence delta.
    Scope this one-time replay to confirmed answers still matching the PDF;
    subsequent manual changes keep using the untouched official event scripts.
    """
    partner=state.get('partner')
    if partner is None:return
    commands=[]
    for group in ('applicant','spouse'):
        if group=='spouse' and not partner:continue
        path=contract['groups'][group]['fields']['marital_status']['path']
        expected=fields.get(path)
        if expected is None and (group!='applicant' or partner):continue
        commands.append('var status = xfa.form.resolveNode('+json.dumps(path)+');')
        condition='(status.rawValue === null || status.rawValue === "")' if expected is None else 'String(status.rawValue) === '+json.dumps(expected)
        commands.append('if (status && '+condition+') {')
        if expected is not None:commands.append('status.execEvent("exit");')
        if group=='applicant':
            spouse=contract['groups']['spouse']['container']
            commands.append('var partner = xfa.form.resolveNode('+json.dumps(spouse)+');')
            commands.append('if (partner) partner.presence = '+json.dumps('visible' if partner else 'hidden')+';')
        commands.append('}')
    if not commands:return
    root=resolve(form,contract['root'],kinds=['subform'])
    event=E.SubElement(root,'{'+namespace+'}event',name='imm5406RestorePartner',activity='docReady',ref='$host')
    E.SubElement(event,'{'+namespace+'}script',contentType='application/x-javascript',runAt='client').text='\n'.join(commands)


def resolve(root, path, *, create=False, kinds=None, namespace=None):
    node = root
    if not kinds and E.QName(node).localname == 'datasets':
        node = next(c for c in node if isinstance(c.tag,str) and E.QName(c).localname=='data')
    for offset, part in enumerate(path.split('.')):
        match=PART.fullmatch(part)
        if not match:raise ValueError('Invalid IMM5406 path')
        name,index=match[1],int(match[2] or 0)
        same=[n for n in node if isinstance(n.tag,str) and (n.get('name') if kinds else E.QName(n).localname)==name]
        if match[2] is None and len(same)>1:raise ValueError('Ambiguous IMM5406 path: '+path)
        if index>=len(same) and not create:raise ValueError('Unresolved IMM5406 path: '+path)
        while index>=len(same):
            same.append(E.SubElement(node,'{'+namespace+'}'+kinds[offset],name=name) if kinds else E.SubElement(node,name))
        node=same[index]
    return node


def fill_family_information(payload, output_path, blank_path):
    if payload.get('form')!='IMM5406':raise ValueError('IMM5406 writer rejects other forms')
    lang=payload.get('language','en')
    if lang not in ('en','fr'):raise ValueError('Unsupported IMM5406 language')
    contract=json.loads(Path(__file__).with_name(f'imm5406-contract.{lang}.json').read_text(encoding='utf8'))
    blank,output=Path(blank_path),Path(output_path)
    if blank.resolve()==output.resolve():raise ValueError('Cannot overwrite official blank')
    if hashlib.sha256(blank.read_bytes()).hexdigest()!=contract['sha256']:raise ValueError('IMM5406 template requires review')
    fields,counts,state=payload.get('fields'),payload.get('imm5406_rows',{}),payload.get('imm5406_state',{})
    if fields==[]:fields={}
    if not isinstance(fields,dict) or not isinstance(counts,dict) or not isinstance(state,dict):raise ValueError('Invalid IMM5406 payload')
    if set(counts)-{'children','siblings'} or set(state)-{'partner','applicant_married','spouse_married','children','siblings'}:raise ValueError('Unallowlisted IMM5406 state')
    if any(type(v) is not int or v<0 for v in counts.values()) or any(v is not None and type(v) is not bool for v in state.values()):raise ValueError('Invalid IMM5406 state value')
    allowed={}
    for group,definition in contract['groups'].items():
        for index in range(counts.get(group,0) if group in ('children','siblings') else 1):
            for meta in definition['fields'].values():allowed[meta['path'].replace('{index}',str(index))]=meta
    for path in contract['selectors'].values():
        allowed[path]={'kinds':['subform']*(len(path.split('.'))-1)+['field'],'value_type':'integer','max_chars':1,'allowed_values':['0','1']}
    for path,value in fields.items():
        if path not in allowed:raise ValueError('Unallowlisted IMM5406 field: '+str(path))
        meta=allowed[path]
        if not isinstance(value,str) or len(value)>meta['max_chars']:raise ValueError('Invalid IMM5406 value: '+path)
        if meta.get('allowed_values') and value not in meta['allowed_values']:raise ValueError('Invalid IMM5406 export: '+path)
    if os.name=='nt' and not os.environ.get('SystemRoot'):
        import ctypes
        buffer=ctypes.create_unicode_buffer(32768)
        if not ctypes.windll.kernel32.GetWindowsDirectoryW(buffer,len(buffer)):raise OSError('Cannot resolve Windows system directory')
        os.environ['SystemRoot']=buffer.value
    from pyhanko.pdf_utils.incremental_writer import IncrementalPdfFileWriter
    from pyhanko.pdf_utils import generic
    from imm0008_xfa_checksum import form_checksum
    output.parent.mkdir(parents=True,exist_ok=True)
    with blank.open('rb') as source:
        writer=IncrementalPdfFileWriter(source)
        if writer.prev.encrypted:writer.encrypt('')
        xfa=writer.root['/AcroForm']['/XFA']
        packets={str(xfa[i]):xfa[i+1].data for i in range(0,len(xfa),2)}
        parser=E.XMLParser(remove_comments=False,remove_pis=False,resolve_entities=False)
        data=E.fromstring(packets['datasets'],parser)
        ns=contract['form_namespace']
        form=E.fromstring(packets['form'],parser)
        # Preserve the official saved skeleton and canonical Section A/B/C order.
        for group,definition in contract['groups'].items():
            path=definition['container']
            if group in ('children','siblings'):
                path=path.rsplit('.',1)[0]
                resolve(form,path,create=True,kinds=['subform']*len(path.split('.')),namespace=ns)
                selector=contract['selectors']['hideChildren' if group=='children' else 'hideSiblings']
                resolve(form,selector,create=True,kinds=allowed[selector]['kinds'],namespace=ns)
            else:
                resolve(form,path,create=True,kinds=['subform']*len(path.split('.')),namespace=ns)
        page=resolve(form,'IMM_5406.page1',kinds=['subform','subform'])
        section_a=resolve(form,'IMM_5406.page1.SectionA',kinds=['subform']*3)
        page.remove(section_a);page.insert(0,section_a)
        for group in ('children','siblings'):
            definition=contract['groups'][group]
            first=resolve(data,definition['container'].replace('{index}','0'))
            parent=first.getparent();name=E.QName(first).localname
            rows=[n for n in parent if isinstance(n.tag,str) and E.QName(n).localname==name]
            wanted=max(definition['initial'],counts.get(group,0))
            from copy import deepcopy
            while len(rows)<wanted:
                row=deepcopy(first);parent.append(row);rows.append(row)
            saved_first=resolve(form,definition['container'].replace('{index}','0'),create=True,kinds=['subform']*len(definition['container'].split('.')),namespace=ns)
            saved_parent=saved_first.getparent()
            saved_rows=[n for n in saved_parent if n.get('name')==name]
            while len(saved_rows)<wanted:
                saved=deepcopy(saved_first);saved_parent.append(saved);saved_rows.append(saved)
            for index in range(wanted):
                if state.get(group) is not None:
                    path=definition['container'].replace('{index}',str(index))
                    resolve(form,path,create=True,kinds=['subform']*len(path.split('.')),namespace=ns).set('presence','visible' if state[group] else 'hidden')
        for path,value in fields.items():
            meta=allowed[path]
            resolve(data,path,create=False).text=value
            widget=resolve(form,path,create=True,kinds=meta['kinds'],namespace=ns)
            # Adobe forbids a value child on exclGroup; export stays dataset-bound.
            if meta['kinds'][-1]!='exclGroup':
                for old in widget.xpath('./*[local-name()="value"]'):widget.remove(old)
                val=E.SubElement(widget,'{'+ns+'}value')
                E.SubElement(val,'{'+ns+'}'+meta['value_type']).text=value
        visibility=[('partner',contract['groups']['spouse']['container']),
                    ('applicant_married',contract['groups']['applicant']['fields']['marriage_present']['path'].rsplit('.',2)[0]),
                    ('spouse_married',contract['groups']['spouse']['fields']['marriage_present']['path'].rsplit('.',2)[0])]
        for key,path in visibility:
            if state.get(key) is not None:
                resolve(form,path,create=True,kinds=['subform']*len(path.split('.')),namespace=ns).set('presence','visible' if state[key] else 'hidden')
        # A saved subform without its instance manager can be recreated from the
        # template during data merge, losing presence (notably Section A).
        # Keep existing managers; add one before each newly represented group.
        for node in list(form.xpath('//*[local-name()="subform"]')):
            parent=node.getparent()
            if E.QName(parent).localname=='form' or not node.get('name'):continue
            name='_'+node.get('name')
            if not any(E.QName(child).localname=='instanceManager' and child.get('name')==name for child in parent):
                parent.insert(parent.index(node),E.Element('{'+ns+'}instanceManager',name=name))
        restore_partner_on_open(form,contract,fields,state,ns)
        datasets=E.tostring(data)
        form.set('checksum',form_checksum(packets['template'],datasets))
        updates={'datasets':datasets,'form':E.tostring(form)}
        for index in range(0,len(xfa),2):
            name=str(xfa[index])
            if name in updates:
                reference=xfa.raw_get(index+1)
                stream=generic.StreamObject(stream_data=updates[name]);stream.compress()
                writer.objects[(reference.generation,reference.idnum)]=stream
        with output.open('wb') as target:writer.write(target)
    return output