src/exporter.py

"""
exporter.py — Multi-format document export engine.
Features:
  - Searchable PDF with invisible text layer
  - Annotated PDF with colored bounding boxes
  - Structured JSON export with metadata
  - CSV / Excel tabular export
  - Plain text export
  - Redacted document export
"""

import io
import json
import csv
import datetime
import numpy as np
from typing import List, Dict, Any
from PIL import Image


def export_json(
    ocr_results: List[Dict],
    entities: Dict,
    metadata: Dict = None,
) -> bytes:
    """Export OCR results and entities to structured JSON bytes."""
    output = {
        "metadata": {
            "exported_at": datetime.datetime.utcnow().isoformat() + "Z",
            "total_words": len(ocr_results),
            "avg_confidence": float(np.mean([r["confidence"] for r in ocr_results])) if ocr_results else 0,
            **(metadata or {}),
        },
        "document_type": entities.get("document_type", "Unknown"),
        "entities": {
            "dates": entities.get("dates", []),
            "amounts": entities.get("amounts", []),
            "emails": entities.get("emails", []),
            "phones": entities.get("phones", []),
            "urls": entities.get("urls", []),
            "invoice_fields": entities.get("invoice_fields", {}),
            "pii_detected": {k: len(v) > 0 for k, v in entities.get("pii", {}).items()},
        },
        "words": [
            {
                "text": r["text"],
                "confidence": round(r["confidence"], 4),
                "x_min": r["x_min"],
                "y_min": r["y_min"],
                "x_max": r["x_max"],
                "y_max": r["y_max"],
                "center": list(r["center"]),
            }
            for r in ocr_results
        ],
    }
    return json.dumps(output, indent=2, ensure_ascii=False).encode("utf-8")


def export_csv(ocr_results: List[Dict], entities: Dict) -> bytes:
    """Export OCR results to CSV bytes."""
    output = io.StringIO()
    writer = csv.writer(output)

    # Header
    writer.writerow([
        "Index", "Text", "Confidence", "X_Min", "Y_Min", "X_Max", "Y_Max", "Width", "Height", "Center_X", "Center_Y"
    ])

    for i, r in enumerate(ocr_results):
        writer.writerow([
            i + 1,
            r["text"],
            f"{r['confidence']:.4f}",
            r["x_min"],
            r["y_min"],
            r["x_max"],
            r["y_max"],
            r["width"],
            r["height"],
            r["center"][0],
            r["center"][1],
        ])

    # Blank + entities section
    writer.writerow([])
    writer.writerow(["=== EXTRACTED ENTITIES ==="])
    writer.writerow(["Type", "Value"])
    for date in entities.get("dates", []):
        writer.writerow(["Date", date])
    for amt in entities.get("amounts", []):
        writer.writerow(["Amount", amt])
    for email in entities.get("emails", []):
        writer.writerow(["Email", email])
    for phone in entities.get("phones", []):
        writer.writerow(["Phone", phone])
    for field, value in entities.get("invoice_fields", {}).items():
        writer.writerow([f"Invoice: {field}", value])

    return output.getvalue().encode("utf-8")


def export_txt(ocr_results: List[Dict], redacted: bool = False, full_text: str = "") -> bytes:
    """Export clean plain text."""
    if full_text:
        return full_text.encode("utf-8")
    sorted_results = sorted(ocr_results, key=lambda r: (r["y_min"], r["x_min"]))
    text = " ".join(r["text"] for r in sorted_results)
    return text.encode("utf-8")


def export_searchable_pdf(
    pil_img: Image.Image,
    ocr_results: List[Dict],
    page_title: str = "OCR Document",
) -> bytes:
    """
    Generate a searchable PDF:
    - Original image as background
    - Invisible text layer placed at exact bbox coordinates
    Returns raw PDF bytes.
    """
    from reportlab.pdfgen import canvas as rl_canvas
    from reportlab.lib.pagesizes import A4
    from reportlab.lib.utils import ImageReader

    buffer = io.BytesIO()

    # Get image dimensions
    img_w, img_h = pil_img.size
    page_w, page_h = img_w, img_h

    c = rl_canvas.Canvas(buffer, pagesize=(page_w, page_h))

    # Draw background image
    img_buffer = io.BytesIO()
    pil_img.save(img_buffer, format="PNG")
    img_buffer.seek(0)
    c.drawImage(ImageReader(img_buffer), 0, 0, width=page_w, height=page_h)

    # Overlay invisible text at OCR positions
    c.setFillColorRGB(0, 0, 0, 0)  # Transparent text
    for result in ocr_results:
        text = result["text"]
        x = result["x_min"]
        # PDF coords: y=0 at bottom, image y=0 at top
        y = page_h - result["y_max"]
        font_size = max(4, result["height"])
        c.setFont("Helvetica", min(font_size, 72))
        c.drawString(x, y, text)

    c.save()
    buffer.seek(0)
    return buffer.read()


def export_annotated_pdf(
    annotated_pil: Image.Image,
    ocr_results: List[Dict],
    entities: Dict,
    title: str = "Annotated OCR Report",
) -> bytes:
    """
    Generate an annotated report PDF with:
    - Annotated image (bounding boxes)
    - Entity summary table
    - Full word-by-word confidence table
    """
    from reportlab.lib.pagesizes import A4
    from reportlab.lib.units import cm
    from reportlab.lib.styles import getSampleStyleSheet
    from reportlab.lib import colors
    from reportlab.platypus import SimpleDocTemplate, Image as RLImage, Table, TableStyle, Paragraph, Spacer
    from reportlab.lib.utils import ImageReader

    buffer = io.BytesIO()
    doc = SimpleDocTemplate(buffer, pagesize=A4, topMargin=1*cm, bottomMargin=1*cm)
    styles = getSampleStyleSheet()
    story = []

    # Title
    story.append(Paragraph(f"<b>{title}</b>", styles["Title"]))
    story.append(Paragraph(
        f"Generated: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')} | "
        f"Words: {len(ocr_results)} | Type: {entities.get('document_type', 'Unknown')}",
        styles["Normal"]
    ))
    story.append(Spacer(1, 0.3*cm))

    # Annotated image
    img_buffer = io.BytesIO()
    # Resize to fit A4 width
    max_w = 500
    iw, ih = annotated_pil.size
    scale = min(max_w / iw, 1.0)
    resized = annotated_pil.resize((int(iw * scale), int(ih * scale)), Image.LANCZOS)
    resized.save(img_buffer, format="PNG")
    img_buffer.seek(0)
    story.append(RLImage(img_buffer, width=int(iw * scale), height=int(ih * scale)))
    story.append(Spacer(1, 0.3*cm))

    # Entities section
    story.append(Paragraph("<b>Extracted Entities</b>", styles["Heading2"]))
    entity_data = [["Category", "Values"]]
    if entities.get("dates"):
        entity_data.append(["Dates", ", ".join(entities["dates"])])
    if entities.get("amounts"):
        entity_data.append(["Amounts", ", ".join(entities["amounts"])])
    if entities.get("emails"):
        entity_data.append(["Emails", ", ".join(entities["emails"])])
    if entities.get("phones"):
        entity_data.append(["Phones", ", ".join(entities["phones"])])
    if entities.get("invoice_fields"):
        for k, v in entities["invoice_fields"].items():
            entity_data.append([k.replace("_", " ").title(), str(v)])

    if len(entity_data) > 1:
        entity_table = Table(entity_data, colWidths=[120, 370])
        entity_table.setStyle(TableStyle([
            ("BACKGROUND", (0, 0), (-1, 0), colors.HexColor("#6366f1")),
            ("TEXTCOLOR", (0, 0), (-1, 0), colors.white),
            ("FONTNAME", (0, 0), (-1, 0), "Helvetica-Bold"),
            ("FONTSIZE", (0, 0), (-1, -1), 9),
            ("ROWBACKGROUNDS", (0, 1), (-1, -1), [colors.HexColor("#f8fafc"), colors.white]),
            ("GRID", (0, 0), (-1, -1), 0.5, colors.HexColor("#e2e8f0")),
            ("LEFTPADDING", (0, 0), (-1, -1), 6),
            ("RIGHTPADDING", (0, 0), (-1, -1), 6),
            ("TOPPADDING", (0, 0), (-1, -1), 4),
            ("BOTTOMPADDING", (0, 0), (-1, -1), 4),
        ]))
        story.append(entity_table)
        story.append(Spacer(1, 0.3*cm))

    # Word confidence table
    story.append(Paragraph("<b>Word-by-Word Confidence Table</b>", styles["Heading2"]))
    word_data = [["#", "Text", "Confidence", "Position"]]
    for i, r in enumerate(ocr_results[:100]):  # Limit to first 100
        conf_pct = f"{r['confidence']:.1%}"
        pos = f"({r['x_min']}, {r['y_min']})"
        word_data.append([str(i + 1), r["text"][:40], conf_pct, pos])

    word_table = Table(word_data, colWidths=[30, 260, 80, 120])
    word_table.setStyle(TableStyle([
        ("BACKGROUND", (0, 0), (-1, 0), colors.HexColor("#1e293b")),
        ("TEXTCOLOR", (0, 0), (-1, 0), colors.white),
        ("FONTNAME", (0, 0), (-1, 0), "Helvetica-Bold"),
        ("FONTSIZE", (0, 0), (-1, -1), 8),
        ("ROWBACKGROUNDS", (0, 1), (-1, -1), [colors.HexColor("#f1f5f9"), colors.white]),
        ("GRID", (0, 0), (-1, -1), 0.3, colors.HexColor("#cbd5e1")),
        ("LEFTPADDING", (0, 0), (-1, -1), 5),
        ("RIGHTPADDING", (0, 0), (-1, -1), 5),
        ("TOPPADDING", (0, 0), (-1, -1), 3),
        ("BOTTOMPADDING", (0, 0), (-1, -1), 3),
    ]))
    story.append(word_table)

    doc.build(story)
    buffer.seek(0)
    return buffer.read()


def export_bundle_zip(
    processed_img: Image.Image,
    ocr_results: List[Dict],
    entities: Dict,
    full_text: str,
    base_name: str = "document",
) -> bytes:
    """Generate a single ZIP archive containing all export formats."""
    import zipfile
    from src.visualizer import draw_bounding_boxes

    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
        # JSON
        zip_file.writestr(f"{base_name}_data.json", export_json(ocr_results, entities, {"filename": base_name}))
        # CSV
        zip_file.writestr(f"{base_name}_words.csv", export_csv(ocr_results, entities))
        # Plain text
        zip_file.writestr(f"{base_name}_text.txt", export_txt(ocr_results, full_text=full_text))
        # Searchable PDF
        try:
            searchable_pdf = export_searchable_pdf(processed_img, ocr_results, base_name)
            zip_file.writestr(f"{base_name}_searchable.pdf", searchable_pdf)
        except Exception:
            pass
        # Annotated PDF
        try:
            annotated_img = draw_bounding_boxes(processed_img, ocr_results, show_labels=True, show_confidence=True)
            report_pdf = export_annotated_pdf(annotated_img, ocr_results, entities, title=f"DocIQ Report - {base_name}")
            zip_file.writestr(f"{base_name}_report.pdf", report_pdf)
        except Exception:
            pass

    zip_buffer.seek(0)
    return zip_buffer.read()