#!/usr/bin/env python3
"""Unify all tracked change authors in a docx to 'WB'.

Usage: python unify-author-wb.py <input.docx> [output.docx]
If output is omitted, overwrites input.

Covers: w:ins, w:del, rPrChange, pPrChange, sectPrChange,
        tblPrChange, trPrChange, tcPrChange.
Also fixes XML declaration (single→double quotes) for OnlyOffice compatibility.
"""
import sys, os, zipfile, re
from lxml import etree

WNS = '{http://schemas.openxmlformats.org/wordprocessingml/2006/main}'

CHANGE_TAGS = ('ins', 'del', 'rPrChange', 'pPrChange',
               'sectPrChange', 'tblPrChange', 'trPrChange', 'tcPrChange')

def unify_author(src_path, out_path=None):
    if out_path is None:
        out_path = src_path
    tmp_path = out_path + '.tmp'

    zin = zipfile.ZipFile(src_path, 'r')
    doc_xml = zin.read('word/document.xml')
    tree = etree.fromstring(doc_xml)
    body = tree.find(f'{WNS}body')

    changed = 0
    for tag_suffix in CHANGE_TAGS:
        for elem in body.iter(f'{WNS}{tag_suffix}'):
            author = elem.get(f'{WNS}author')
            if author and author != 'WB':
                elem.set(f'{WNS}author', 'WB')
                changed += 1

    # Serialize + fix XML declaration
    doc_bytes = etree.tostring(tree, xml_declaration=True, encoding='UTF-8', standalone=True)
    doc_str = doc_bytes.decode('utf-8')
    doc_str = doc_str.replace(
        "<?xml version='1.0' encoding='UTF-8' standalone='yes'?>",
        '<?xml version="1.0" encoding="UTF-8" standalone="yes"?>')

    with zipfile.ZipFile(tmp_path, 'w', zipfile.ZIP_DEFLATED) as zout:
        for item in zin.namelist():
            if item == 'word/document.xml':
                zout.writestr(item, doc_str.encode('utf-8'))
            else:
                zout.writestr(item, zin.read(item))
    zin.close()
    os.replace(tmp_path, out_path)

    # Verify
    z = zipfile.ZipFile(out_path)
    vdoc = z.read('word/document.xml')
    vtree = etree.fromstring(vdoc)
    vbody = vtree.find(f'{WNS}body')
    remaining = set()
    for tag_suffix in CHANGE_TAGS:
        for elem in vbody.iter(f'{WNS}{tag_suffix}'):
            a = elem.get(f'{WNS}author', '')
            if a != 'WB':
                remaining.add(a)
    z.close()

    print(f"✅ {changed} author attributes → WB")
    if remaining:
        print(f"⚠️ Remaining non-WB authors: {remaining}")
    else:
        print(f"   All authors = WB")
    print(f"   Output: {out_path} ({os.path.getsize(out_path):,} bytes)")

if __name__ == '__main__':
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(1)
    src = sys.argv[1]
    out = sys.argv[2] if len(sys.argv) > 2 else None
    unify_author(src, out)