#!/usr/bin/env python3
"""
contract_preprocess.py — 合同预处理：检测并切割非审查图片内容

用途：在workflow审查前，检测合同末尾的纯图片附件（如招标公告截图、中标通知书等），
     切割出来保存，审查完后再还原。

判断逻辑：
1. 扫描文件结构：文字段落数 vs 图片段落数
2. 全文/大部分是图片（扫描件合同）→ 不切割，标记需OCR
3. 正文文字+末尾图片附件 → 切割末尾图片区域
4. 切割点：从最后一个"纯文字附件"结束后，到第一个"纯图片附件"开始

输出：
- {basename}_stripped.docx — 去掉图片附件的版本（供workflow处理）
- {basename}_cutdata.json — 切割信息（供还原用）
"""

import zipfile, json, os, sys, re
from lxml import etree

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


def analyze_contract(docx_path):
    """Analyze contract structure, return analysis dict"""
    with zipfile.ZipFile(docx_path) as z:
        doc = etree.fromstring(z.read('word/document.xml'))
        media_files = {n: z.getinfo(n).file_size for n in z.namelist() if n.startswith('word/media/')}

    body = doc.find(f'{{{W}}}body')
    paras = body.findall(f'{{{W}}}p')

    paragraphs = []
    total_text_chars = 0
    total_img_paras = 0

    for i, p in enumerate(paras):
        texts = p.findall(f'.//{{{W}}}t')
        text = ''.join(t.text or '' for t in texts).strip()

        has_img = any('drawing' in (e.tag if isinstance(e.tag, str) else '') for e in p.iter())

        blips = list(p.iter(f'{{{A_NS}}}blip'))
        img_rids = [b.get(f'{{{R_NS}}}embed', '') for b in blips]

        total_text_chars += len(text)
        if has_img:
            total_img_paras += 1

        paragraphs.append({
            'idx': i,
            'text': text,
            'text_len': len(text),
            'has_img': has_img,
            'img_rids': img_rids,
            'is_appendix_heading': bool(re.match(r'^附件[一二三四五六七八九十\d]+[：:、]', text)),
        })

    return {
        'total_paras': len(paras),
        'total_text_chars': total_text_chars,
        'total_img_paras': total_img_paras,
        'media_files': media_files,
        'total_media_bytes': sum(media_files.values()),
        'paragraphs': paragraphs,
    }


def detect_cut_zone(analysis):
    """Detect if there's a tail image zone to cut."""
    paras = analysis['paragraphs']
    total = analysis['total_paras']

    text_paras = sum(1 for p in paras if p['text_len'] > 0 and not p['has_img'])
    img_paras = analysis['total_img_paras']

    if text_paras == 0 and img_paras > 0:
        return {'action': 'ocr', 'reason': '全文无文字段落，疑似扫描件合同'}

    if img_paras == 0:
        return None

    img_ratio = img_paras / max(1, text_paras + img_paras)
    if img_ratio > 0.5:
        return {'action': 'ocr', 'reason': f'图片段落占比{img_ratio:.0%}，疑似扫描件合同'}

    # Find tail image zones
    image_zones = []
    i = 0
    while i < total:
        p = paras[i]
        if p['is_appendix_heading']:
            zone_start = i
            zone_has_images = False
            zone_has_text_content = False
            j = i + 1

            while j < total:
                next_p = paras[j]
                if next_p['is_appendix_heading']:
                    break
                if next_p['has_img']:
                    zone_has_images = True
                if next_p['text_len'] > 20 and not next_p['has_img']:
                    zone_has_text_content = True
                j += 1

            image_zones.append({
                'start_idx': zone_start,
                'end_idx': j - 1,
                'heading': p['text'],
                'has_images': zone_has_images,
                'has_text': zone_has_text_content,
                'is_image_only': zone_has_images and not zone_has_text_content,
            })
            i = j
        else:
            i += 1

    # Find consecutive image-only appendices at the tail
    tail_cut_zones = []
    for zone in reversed(image_zones):
        if zone['is_image_only']:
            tail_cut_zones.insert(0, zone)
        else:
            break

    if not tail_cut_zones:
        return None

    cut_start = tail_cut_zones[0]['start_idx']
    cut_headings = [z['heading'] for z in tail_cut_zones]

    return {
        'action': 'cut',
        'cut_start_idx': cut_start,
        'cut_end_idx': total - 1,
        'cut_headings': cut_headings,
        'reason': f'末尾{len(tail_cut_zones)}个附件为纯图片：{", ".join(cut_headings)}',
    }


def preprocess_contract(docx_path, output_dir=None):
    """Main entry: analyze and optionally strip tail images."""
    if output_dir is None:
        output_dir = os.path.dirname(docx_path) or '.'

    basename = os.path.splitext(os.path.basename(docx_path))[0]

    analysis = analyze_contract(docx_path)
    cut_info = detect_cut_zone(analysis)

    print(f"\n=== 合同预处理分析 ===")
    print(f"文件: {os.path.basename(docx_path)}")
    print(f"段落数: {analysis['total_paras']}")
    print(f"文字字符: {analysis['total_text_chars']}")
    print(f"图片段落: {analysis['total_img_paras']}")
    print(f"媒体文件: {len(analysis['media_files'])} ({analysis['total_media_bytes']:,} bytes)")

    if cut_info is None:
        print(f"结论: 无需切割")
        return {'action': 'none', 'analysis': analysis}

    if cut_info['action'] == 'ocr':
        print(f"结论: {cut_info['reason']}，需OCR处理")
        return {'action': 'ocr', 'reason': cut_info['reason'], 'analysis': analysis}

    cut_start = cut_info['cut_start_idx']
    print(f"结论: 需切割 — {cut_info['reason']}")
    print(f"切割点: 段落 #{cut_start}")

    with zipfile.ZipFile(docx_path) as z:
        doc = etree.fromstring(z.read('word/document.xml'))
        all_files = {}
        for name in z.namelist():
            all_files[name] = z.read(name)

    body = doc.find(f'{{{W}}}body')
    paras = body.findall(f'{{{W}}}p')

    cut_paras_xml = []
    for i in range(cut_start, len(paras)):
        cut_paras_xml.append(etree.tostring(paras[i], encoding='unicode'))

    for i in range(len(paras) - 1, cut_start - 1, -1):
        body.remove(paras[i])

    cut_rids = set()
    for p_info in analysis['paragraphs'][cut_start:]:
        cut_rids.update(p_info['img_rids'])

    rels_xml = all_files.get('word/_rels/document.xml.rels', b'')
    if isinstance(rels_xml, bytes):
        rels_xml = rels_xml.decode()
    rid_to_media = {}
    for m in re.finditer(r'Id="(rId\d+)"[^/]*Target="(media/[^"]+)"', rels_xml):
        rid_to_media[m.group(1)] = f'word/{m.group(2)}'

    cut_media = {}
    for rid in cut_rids:
        media_path = rid_to_media.get(rid)
        if media_path and media_path in all_files:
            cut_media[media_path] = len(all_files[media_path])

    stripped_path = os.path.join(output_dir, f'{basename}_stripped.docx')
    all_files['word/document.xml'] = etree.tostring(doc, xml_declaration=True, encoding='UTF-8', standalone=True)

    with zipfile.ZipFile(stripped_path, 'w', zipfile.ZIP_DEFLATED) as zout:
        for name, data in all_files.items():
            zout.writestr(name, data)

    cutdata = {
        'original_file': os.path.basename(docx_path),
        'cut_start_idx': cut_start,
        'total_paras_original': len(paras) + len(cut_paras_xml),
        'cut_paragraphs_xml': cut_paras_xml,
        'cut_headings': cut_info['cut_headings'],
        'cut_media_files': list(cut_media.keys()),
        'reason': cut_info['reason'],
    }

    cutdata_path = os.path.join(output_dir, f'{basename}_cutdata.json')
    with open(cutdata_path, 'w', encoding='utf-8') as f:
        json.dump(cutdata, f, ensure_ascii=False, indent=2)

    stripped_size = os.path.getsize(stripped_path)
    original_size = os.path.getsize(docx_path)

    print(f"\n输出:")
    print(f"  stripped: {stripped_path} ({stripped_size:,} bytes)")
    print(f"  cutdata:  {cutdata_path}")
    print(f"  大小变化: {original_size:,} → {stripped_size:,} bytes ({stripped_size/original_size:.0%})")

    return {
        'action': 'cut',
        'stripped_path': stripped_path,
        'cutdata_path': cutdata_path,
        'cut_info': cut_info,
        'analysis': analysis,
    }


def restore_contract(reviewed_path, cutdata_path, output_path):
    """Restore cut content back into the reviewed file."""
    with open(cutdata_path, 'r', encoding='utf-8') as f:
        cutdata = json.load(f)

    with zipfile.ZipFile(reviewed_path) as z:
        doc = etree.fromstring(z.read('word/document.xml'))
        all_files = {}
        for name in z.namelist():
            all_files[name] = z.read(name)

    body = doc.find(f'{{{W}}}body')
    sect_pr = body.find(f'{{{W}}}sectPr')

    for para_xml in cutdata['cut_paragraphs_xml']:
        para_elem = etree.fromstring(para_xml)
        if sect_pr is not None:
            sect_pr.addprevious(para_elem)
        else:
            body.append(para_elem)

    original_dir = os.path.dirname(cutdata_path)
    original_name = cutdata['original_file']
    original_path = os.path.join(original_dir, original_name)

    if os.path.exists(original_path):
        with zipfile.ZipFile(original_path) as z_orig:
            for media_file in cutdata.get('cut_media_files', []):
                if media_file not in all_files and media_file in z_orig.namelist():
                    all_files[media_file] = z_orig.read(media_file)
                    print(f"  还原媒体文件: {media_file}")

    all_files['word/document.xml'] = etree.tostring(doc, xml_declaration=True, encoding='UTF-8', standalone=True)

    with zipfile.ZipFile(output_path, 'w', zipfile.ZIP_DEFLATED) as zout:
        for name, data in all_files.items():
            zout.writestr(name, data)

    restored_size = os.path.getsize(output_path)
    print(f"\n=== 合同还原完成 ===")
    print(f"还原文件: {output_path} ({restored_size:,} bytes)")
    print(f"还原段落: {len(cutdata['cut_paragraphs_xml'])} 个")
    print(f"还原附件: {', '.join(cutdata['cut_headings'])}")

    return output_path


if __name__ == '__main__':
    if len(sys.argv) < 2:
        print("Usage:")
        print("  预处理: python contract_preprocess.py preprocess <input.docx> [output_dir]")
        print("  还原:   python contract_preprocess.py restore <reviewed.docx> <cutdata.json> <output.docx>")
        sys.exit(1)

    action = sys.argv[1]

    if action == 'preprocess':
        docx_path = sys.argv[2]
        output_dir = sys.argv[3] if len(sys.argv) > 3 else None
        result = preprocess_contract(docx_path, output_dir)
        print(f"\nResult: {json.dumps({k: v for k, v in result.items() if k != 'analysis'}, ensure_ascii=False, indent=2)}")

    elif action == 'restore':
        reviewed_path = sys.argv[2]
        cutdata_path = sys.argv[3]
        output_path = sys.argv[4]
        restore_contract(reviewed_path, cutdata_path, output_path)

    else:
        print(f"Unknown action: {action}")
        sys.exit(1)