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

SOURCE = Path(r'C:\Users\云\Desktop\SOP-QC-YL')
sys.stdout.reconfigure(encoding='utf-8')
OUT = Path('scripts/sop-qc')
OUT.mkdir(exist_ok=True)
W = '{http://schemas.openxmlformats.org/wordprocessingml/2006/main}'

def clean(value):
    return re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f]', '', value).replace('\r', '\n').strip()

def word_text(node):
    output = []
    for run in node.iter(W+'r'):
        text = ''.join(n.text or '' if n.tag in (W+'t', W+'delText') else '\t' if n.tag == W+'tab' else '\n' if n.tag in (W+'br', W+'cr') else '' for n in run.iter())
        alignment = run.find(W+'rPr/'+W+'vertAlign')
        if alignment is not None and alignment.get(W+'val') == 'superscript' and text.strip():
            text = '^('+text+')'
        output.append(text)
    return clean(''.join(output))

records = []
for file in sorted(SOURCE.iterdir(), key=lambda p: p.name):
    if file.name.startswith('~$') or file.suffix.lower() not in ('.doc', '.docx'):
        continue
    blocks = []
    if file.suffix.lower() == '.docx':
        with zipfile.ZipFile(file) as z:
            body = ET.fromstring(z.read('word/document.xml')).find(W+'body')
            for node in body:
                if node.tag == W+'p':
                    blocks.append({'p': word_text(node)})
                elif node.tag == W+'tbl':
                    blocks.append({'table': [[ '\n'.join(word_text(p) for p in cell.findall(W+'p')).strip() for cell in row.findall(W+'tc')] for row in node.findall(W+'tr')]})
    else:
        body = ET.parse(OUT/'extracted'/(file.name+'.xml')).getroot()
        for node in body:
            if node.tag == 'p':
                blocks.append({'p': clean(node.text or '')})
            elif node.tag == 'table':
                blocks.append({'table': [[clean('\n'.join(p.text or '' for p in cell)) for cell in row] for row in node]})
    text = '\n'.join(b.get('p', '\n'.join(' | '.join(row) for row in b.get('table', []))) for b in blocks)
    records.append({'file': file.name, 'sha256': hashlib.sha256(file.read_bytes()).hexdigest(), 'blocks': blocks, 'text': text})
    (OUT/'extracted'/(file.name+'.txt')).write_text(text, encoding='utf-8')
(OUT/'documents.json').write_text(json.dumps(records, ensure_ascii=False, indent=2), encoding='utf-8')
print(json.dumps({'documents': len(records), 'tables': sum(sum('table' in b for b in d['blocks']) for d in records)}, ensure_ascii=False))