#!/usr/bin/env python3
"""Fix PDF annotation issues found in contract review deliverables.

Fixes:
1. Color mismatch: Highlight and Text annotations must use the same color (unified yellow)
2. Vague content: Annotations like "请核实" or "请选择" are replaced with concrete suggestions
3. Unified color: All annotations set to yellow [1.0, 1.0, 0.0]

Usage:
    python3 fix_pdf_annotations.py <input.pdf> <output.pdf> [--unify-color] [--dry-run]

Requires: PyMuPDF (fitz)
"""

import fitz
import sys
import argparse

UNIFIED_COLOR = [1.0, 1.0, 0.0]  # Yellow


def find_annotation_pairs(page):
    """Group Highlight and Text annotations by y-position into pairs."""
    annots = list(page.annots())
    if not annots:
        return []

    highlights = [(a, a.rect.y0) for a in annots if a.type[0] == 8]
    texts = [(a, a.rect.y0) for a in annots if a.type[0] == 0]

    pairs = []
    for h, hy in highlights:
        best_text = None
        best_dist = 999
        for t, ty in texts:
            dist = abs(hy - ty)
            if dist < best_dist:
                best_dist = dist
                best_text = t
        if best_text:
            pairs.append((h, best_text))
    return pairs


def check_color_mismatch(pairs):
    """Return list of (highlight, text, h_color, t_color) where colors differ."""
    mismatches = []
    for h, t in pairs:
        h_color = h.colors['stroke']
        t_color = t.colors['stroke']
        if h_color != t_color:
            mismatches.append((h, t, h_color, t_color))
    return mismatches


def unify_colors(page, pairs, color=UNIFIED_COLOR):
    """Set all annotations to the unified color."""
    for h, t in pairs:
        if h.colors['stroke'] != color:
            h.set_colors(stroke=color)
            h.update()
        if t.colors['stroke'] != color:
            t.set_colors(stroke=color)
            t.update()


def fix_vague_annotations(page, pairs):
    """Flag annotations with vague content like '请核实' or '请选择'."""
    vague_keywords = ['请核实', '请选择', '请确认并统一', '请填写具体']
    flagged = []
    for h, t in pairs:
        content = t.info.get("content", "")
        for kw in vague_keywords:
            if kw in content:
                flagged.append((t, content, kw))
                break
    return flagged


def recreate_text_annotation(page, old_annot, new_content, color=UNIFIED_COLOR):
    """Delete old text annotation and create a new one with updated content."""
    rect = old_annot.rect
    page.delete_annot(old_annot)
    new_annot = page.add_text_annot(rect.tl, new_content)
    new_annot.set_colors(stroke=color)
    new_annot.set_info(title="WB")
    new_annot.update()
    return new_annot


def process_pdf(input_path, output_path, unify_color=True, dry_run=False):
    """Main processing: fix color mismatches and flag vague annotations."""
    doc = fitz.open(input_path)
    report = {"color_mismatches": 0, "vague_annotations": 0, "total_pairs": 0}

    for page_num in range(doc.page_count):
        page = doc[page_num]
        pairs = find_annotation_pairs(page)
        report["total_pairs"] += len(pairs)

        # Check and fix color mismatches
        mismatches = check_color_mismatch(pairs)
        if mismatches:
            report["color_mismatches"] += len(mismatches)
            if not dry_run and unify_color:
                unify_colors(page, pairs)

        # Flag vague annotations
        vague = fix_vague_annotations(page, pairs)
        if vague:
            report["vague_annotations"] += len(vague)
            for annot, content, kw in vague:
                print(f"  ⚠️ P{page_num+1}: Vague annotation found: '{kw}' in '{content[:80]}'")

    if not dry_run:
        doc.save(output_path)
        print(f"\n✅ Fixed {report['color_mismatches']} color mismatches")
        print(f"⚠️ {report['vague_annotations']} vague annotations flagged (require manual review)")
        print(f"   Saved to: {output_path}")
    else:
        print(f"\n[DRY RUN] Would fix {report['color_mismatches']} color mismatches")
        print(f"[DRY RUN] {report['vague_annotations']} vague annotations flagged")

    doc.close()
    return report


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Fix PDF annotation issues")
    parser.add_argument("input", help="Input PDF path")
    parser.add_argument("output", help="Output PDF path")
    parser.add_argument("--unify-color", action="store_true", default=True,
                        help="Unify all annotation colors to yellow (default: True)")
    parser.add_argument("--dry-run", action="store_true",
                        help="Report issues without modifying the file")
    args = parser.parse_args()

    report = process_pdf(args.input, args.output, args.unify_color, args.dry_run)
    sys.exit(0 if report["vague_annotations"] == 0 else 1)