modules/exporter.py

"""
Export Module — NEW feature
Exports extracted entities to:
1. CSV  — for spreadsheet analysis
2. JSON — for downstream NLP pipelines
3. Annotated DOCX — entities highlighted in a Word document
"""

import csv
import json
import os
from datetime import datetime
from docx import Document
from docx.shared import RGBColor, Pt
from docx.oxml.ns import qn
from docx.oxml import OxmlElement


EXPORT_DIR = os.path.join(os.path.dirname(__file__), "..", "exports")
os.makedirs(EXPORT_DIR, exist_ok=True)


def _timestamp() -> str:
    return datetime.now().strftime("%Y%m%d_%H%M%S")


# ── DOCX highlight colors (RGB) ─────────────────────────────────────────────
DOCX_COLORS = {
    "Person":         RGBColor(0x4F, 0xC3, 0xF7),  # blue
    "Organization":   RGBColor(0x81, 0xC7, 0x84),  # green
    "Location":       RGBColor(0xFF, 0xB7, 0x4D),  # orange
    "Miscellaneous":  RGBColor(0xCE, 0x93, 0xD8),  # purple
}


def export_csv(entities: list[dict], source_name: str = "document") -> str:
    """Export entities to CSV. Returns file path."""
    fname = os.path.join(EXPORT_DIR, f"ner_{source_name}_{_timestamp()}.csv")
    with open(fname, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=["word", "label", "score", "start", "end"])
        writer.writeheader()
        writer.writerows(entities)
    return fname


def export_json(entities: list[dict], meta: dict, source_name: str = "document") -> str:
    """Export entities + metadata to JSON. Returns file path."""
    fname   = os.path.join(EXPORT_DIR, f"ner_{source_name}_{_timestamp()}.json")
    payload = {
        "source":   meta,
        "entities": entities,
        "count":    len(entities),
    }
    with open(fname, "w", encoding="utf-8") as f:
        json.dump(payload, f, indent=2, ensure_ascii=False)
    return fname


def export_annotated_docx(
    original_text: str,
    entities: list[dict],
    source_name: str = "document"
) -> str:
    """
    Create a Word document with entities highlighted in context.
    Returns file path.
    """
    fname = os.path.join(EXPORT_DIR, f"annotated_{source_name}_{_timestamp()}.docx")
    doc   = Document()

    # Title
    title = doc.add_heading("SmartDoc NER — Annotated Output", level=1)
    title.runs[0].font.color.rgb = RGBColor(0x1D, 0x6E, 0x6E)

    doc.add_paragraph(f"Source: {source_name}  |  Entities found: {len(entities)}")
    doc.add_paragraph("")

    # Build annotated paragraph by inserting runs
    para        = doc.add_paragraph()
    sorted_ents = sorted(entities, key=lambda e: e["start"])
    cursor      = 0

    for ent in sorted_ents:
        # Normal text before entity
        if ent["start"] > cursor:
            run = para.add_run(original_text[cursor:ent["start"]])
            run.font.size = Pt(11)

        # Highlighted entity run
        run = para.add_run(original_text[ent["start"]:ent["end"]])
        run.font.size  = Pt(11)
        run.font.bold  = True
        color = DOCX_COLORS.get(ent["label"], RGBColor(0xE0, 0xE0, 0xE0))
        run.font.color.rgb = color

        # Label superscript
        run2 = para.add_run(f"[{ent['label']}]")
        run2.font.size       = Pt(8)
        run2.font.color.rgb  = color

        cursor = ent["end"]

    # Remaining text
    if cursor < len(original_text):
        run = para.add_run(original_text[cursor:])
        run.font.size = Pt(11)

    # Entity summary table
    doc.add_heading("Entity Summary", level=2)
    table = doc.add_table(rows=1, cols=3)
    table.style = "Table Grid"
    hdr = table.rows[0].cells
    hdr[0].text, hdr[1].text, hdr[2].text = "Entity", "Type", "Confidence"

    for ent in sorted(entities, key=lambda e: e["score"], reverse=True):
        row = table.add_row().cells
        row[0].text = ent["word"]
        row[1].text = ent["label"]
        row[2].text = f'{ent["score"]}%'

    doc.save(fname)
    return fname