"""
Standalone function to add comments to a .docx file using zipfile + lxml.
Use AFTER ContractEditor.save() since ContractEditor rewrites document.xml.

Usage:
    from add_comments_to_docx import add_comments_to_docx
    # comments = [(anchor_text, comment_text), ...]
    placed = add_comments_to_docx('/tmp/【修】contract.docx', comments)
"""
import zipfile, io
from datetime import datetime
from lxml import etree

W = 'http://schemas.openxmlformats.org/wordprocessingml/2006/main'
R_NS = 'http://schemas.openxmlformats.org/officeDocument/2006/relationships'

def qn(tag):
    return f'{{{W}}}{tag}'

def add_comments_to_docx(filepath, comments, author='WB'):
    """Add comments to a docx file.
    
    Args:
        filepath: path to docx (modified in place)
        comments: list of (anchor_text, comment_text) tuples
        author: comment author name (default 'WB')
    
    Returns:
        int: number of comments successfully placed
    """
    with open(filepath, 'rb') as f:
        data = f.read()
    zin = zipfile.ZipFile(io.BytesIO(data))
    buf = io.BytesIO()
    zout = zipfile.ZipFile(buf, 'w', zipfile.ZIP_DEFLATED)
    doc_xml = zin.read('word/document.xml')
    doc_tree = etree.fromstring(doc_xml)
    body = doc_tree.find(qn('body'))
    
    nsmap = {'w': W, 'r': R_NS}
    comments_xml = etree.Element(qn('comments'), nsmap=nsmap)
    
    comment_id = 200
    placed = 0
    
    for anchor_text, comment_text in comments:
        cid = str(comment_id)
        comment_id += 1
        
        # Create comment element
        comment_el = etree.SubElement(comments_xml, qn('comment'))
        comment_el.set(qn('id'), cid)
        comment_el.set(qn('author'), author)
        comment_el.set(qn('date'), datetime.now().strftime('%Y-%m-%dT%H:%M:%SZ'))
        cp = etree.SubElement(comment_el, qn('p'))
        cr = etree.SubElement(cp, qn('r'))
        ct = etree.SubElement(cr, qn('t'))
        ct.set('{http://www.w3.org/XML/1998/namespace}space', 'preserve')
        ct.text = comment_text
        
        # Find anchor in document
        for p in body.iter(qn('p')):
            runs = p.findall(f'.//{qn("r")}')
            full = ''.join(''.join(t.text or '' for t in r.findall(qn('t'))) for r in runs)
            if anchor_text in full:
                # Insert commentRangeStart at beginning of paragraph
                crs = etree.Element(qn('commentRangeStart'))
                crs.set(qn('id'), cid)
                p.insert(0, crs)
                
                # Append commentRangeEnd + reference
                cre = etree.Element(qn('commentRangeEnd'))
                cre.set(qn('id'), cid)
                p.append(cre)
                
                ref_run = etree.SubElement(p, qn('r'))
                ref_rpr = etree.SubElement(ref_run, qn('rPr'))
                ref_style = etree.SubElement(ref_rpr, qn('rStyle'))
                ref_style.set(qn('val'), 'CommentReference')
                ref_cr = etree.SubElement(ref_run, qn('commentReference'))
                ref_cr.set(qn('id'), cid)
                
                placed += 1
                break
    
    # Rewrite zip
    for item in zin.namelist():
        if item == 'word/document.xml':
            zout.writestr(item, etree.tostring(doc_tree, xml_declaration=True, encoding='UTF-8', standalone=True))
        elif item == '[Content_Types].xml':
            ct_xml = zin.read(item)
            ct_tree = etree.fromstring(ct_xml)
            if not any('comments.xml' in (el.get('PartName') or '') for el in ct_tree):
                override = etree.SubElement(ct_tree, 'Override')
                override.set('PartName', '/word/comments.xml')
                override.set('ContentType', 'application/vnd.openxmlformats-officedocument.wordprocessingml.comments+xml')
            zout.writestr(item, etree.tostring(ct_tree, xml_declaration=True, encoding='UTF-8', standalone=True))
        elif item == 'word/_rels/document.xml.rels':
            rels_xml = zin.read(item)
            rels_tree = etree.fromstring(rels_xml)
            if not any('comments.xml' in (el.get('Target') or '') for el in rels_tree):
                rel = etree.SubElement(rels_tree, 'Relationship')
                rel.set('Id', f'rId{len(rels_tree) + 10}')
                rel.set('Type', 'http://schemas.openxmlformats.org/officeDocument/2006/relationships/comments')
                rel.set('Target', 'comments.xml')
            zout.writestr(item, etree.tostring(rels_tree, xml_declaration=True, encoding='UTF-8', standalone=True))
        else:
            zout.writestr(item, zin.read(item))
    
    zout.writestr('word/comments.xml', etree.tostring(comments_xml, xml_declaration=True, encoding='UTF-8', standalone=True))
    zout.close()
    zin.close()
    
    with open(filepath, 'wb') as f:
        f.write(buf.getvalue())
    return placed