"""
SmartDoc NER Pro — Upgraded Named Entity Recognition System
Author: Nafees Ahmad | PAF-IAST

Modules:
  1. Text NER          — direct text input with inline highlighting
  2. Document NER      — upload PDF / DOCX / TXT
  3. Batch Processing  — multiple documents at once
  4. Entity Analytics  — frequency, co-occurrence, confidence stats
  5. Search & Filter   — search within extracted entities
  6. Export            — CSV, JSON, annotated DOCX

Model: dslim/bert-large-NER  (F1 ~92.8% on CoNLL-2003)
       vs spacy en_core_web_sm (F1 ~85%)  — significant accuracy gain
"""

import gradio as gr
import pandas as pd
import sys
import os

sys.path.insert(0, os.path.dirname(__file__))

from modules.ner_engine      import run_ner, get_highlighted_html, LABEL_COLORS
from modules.doc_reader      import read_document, chunk_text
from modules.analytics       import (build_summary_table, label_counts,
                                     top_entities, co_occurrence,
                                     confidence_stats, deduplicate_entities)
from modules.exporter        import export_csv, export_json, export_annotated_docx
from modules.entity_search   import search_entities, filter_by_label, filter_by_confidence
from modules.batch_processor import process_single, process_batch, aggregate_entities

# ── Global state (simple in-memory store for current session) ────────────────
_state = {
    "entities":    [],
    "text":        "",
    "meta":        {},
    "batch_results": [],
}


# ════════════════════════════════════════════════════════════════════════════
#  TAB 1 — Text NER
# ════════════════════════════════════════════════════════════════════════════
def process_text(text: str, min_confidence: float):
    if not text.strip():
        return "<p style='color:gray'>Please enter some text.</p>", [], "", ""

    entities = run_ner(text)
    entities = filter_by_confidence(entities, min_confidence)

    _state["entities"] = entities
    _state["text"]     = text
    _state["meta"]     = {"type": "Text input", "word_count": len(text.split()), "filename": "text_input"}

    html    = get_highlighted_html(text, entities)
    rows    = build_summary_table(entities)
    stats   = confidence_stats(entities)
    counts  = label_counts(entities)

    stats_text = (
        f"Total entities: {stats['total']}  |  "
        f"Avg confidence: {stats['avg']}%  |  "
        f"Min: {stats['min']}%  |  Max: {stats['max']}%\n"
        f"Breakdown: {counts}"
    )

    headers = ["Entity", "Type", "Confidence (%)"]
    df      = pd.DataFrame(rows, columns=headers) if rows else pd.DataFrame(columns=headers)

    return html, df, stats_text, f"{len(entities)} entities extracted"


# ════════════════════════════════════════════════════════════════════════════
#  TAB 2 — Document NER
# ════════════════════════════════════════════════════════════════════════════
def process_document(file, min_confidence: float):
    if file is None:
        return "<p style='color:gray'>Please upload a file.</p>", [], ""

    try:
        text, meta = read_document(file.name)
    except ValueError as e:
        return f"<p style='color:red'>{e}</p>", [], ""

    chunks   = chunk_text(text, max_chars=400)
    entities = []
    offset   = 0

    for chunk in chunks:
        ents = run_ner(chunk)
        for e in ents:
            e["start"] += offset
            e["end"]   += offset
        entities.extend(ents)
        offset += len(chunk) + 1

    entities = filter_by_confidence(entities, min_confidence)
    _state["entities"] = entities
    _state["text"]     = text
    _state["meta"]     = meta

    html   = get_highlighted_html(text[:3000], entities)   # display first 3000 chars
    rows   = build_summary_table(entities)
    counts = label_counts(entities)

    headers = ["Entity", "Type", "Confidence (%)"]
    df      = pd.DataFrame(rows, columns=headers) if rows else pd.DataFrame(columns=headers)

    info = (
        f"File: {meta['filename']}  |  Type: {meta['type']}  |  "
        f"Pages: {meta['pages']}  |  Words: {meta['word_count']}  |  "
        f"Entities found: {len(entities)}  |  Breakdown: {counts}"
    )
    return html, df, info


# ════════════════════════════════════════════════════════════════════════════
#  TAB 3 — Export
# ════════════════════════════════════════════════════════════════════════════
def do_export(export_type: str):
    entities = _state["entities"]
    text     = _state["text"]
    meta     = _state["meta"]
    name     = meta.get("filename", "document").replace(".", "_")

    if not entities:
        return None, "No entities to export. Run NER first."

    if export_type == "CSV":
        path = export_csv(entities, name)
    elif export_type == "JSON":
        path = export_json(entities, meta, name)
    elif export_type == "Annotated DOCX":
        path = export_annotated_docx(text, entities, name)
    else:
        return None, "Unknown export type."

    return path, f"Exported {len(entities)} entities to {os.path.basename(path)}"


# ════════════════════════════════════════════════════════════════════════════
#  TAB 4 — Analytics
# ════════════════════════════════════════════════════════════════════════════
def show_analytics():
    entities = _state["entities"]
    if not entities:
        return "Run NER on a document first.", "", ""

    deduped  = deduplicate_entities(entities)
    top      = top_entities(entities, top_n=5)
    pairs    = co_occurrence(entities)
    stats    = confidence_stats(entities)
    counts   = label_counts(entities)

    top_text = "TOP ENTITIES PER TYPE:\n" + "-"*40 + "\n"
    for label, items in top.items():
        top_text += f"\n{label}:\n"
        for word, freq in items:
            top_text += f"  • {word} (×{freq})\n"

    pairs_text = ""
    if pairs:
        pairs_text = "\nPERSON — ORGANIZATION CO-OCCURRENCES:\n" + "-"*40 + "\n"
        for person, org in pairs[:10]:
            pairs_text += f"  {person}  ↔  {org}\n"

    stats_text = (
        f"\nCONFIDENCE STATISTICS:\n" + "-"*40 + "\n"
        f"  Total entities:   {stats['total']}\n"
        f"  Unique entities:  {len(deduped)}\n"
        f"  Avg confidence:   {stats['avg']}%\n"
        f"  Max confidence:   {stats['max']}%\n"
        f"  Min confidence:   {stats['min']}%\n\n"
        f"LABEL DISTRIBUTION:\n" + "-"*40 + "\n"
    )
    for label, count in counts.items():
        bar = "█" * int(count / max(counts.values()) * 20)
        stats_text += f"  {label:<20} {bar} {count}\n"

    return top_text, pairs_text, stats_text


# ════════════════════════════════════════════════════════════════════════════
#  TAB 5 — Search & Filter
# ════════════════════════════════════════════════════════════════════════════
def search_and_filter(query: str, label_filter: list, min_conf: float):
    entities = _state["entities"]
    if not entities:
        return [], "Run NER first."

    filtered = search_entities(entities, query)
    filtered = filter_by_label(filtered, label_filter) if label_filter else filtered
    filtered = filter_by_confidence(filtered, min_conf)

    rows    = build_summary_table(filtered)
    headers = ["Entity", "Type", "Confidence (%)"]
    df      = pd.DataFrame(rows, columns=headers) if rows else pd.DataFrame(columns=headers)
    info    = f"{len(filtered)} entities match your filters (from {len(entities)} total)"
    return df, info


# ════════════════════════════════════════════════════════════════════════════
#  TAB 6 — Batch Processing
# ════════════════════════════════════════════════════════════════════════════
def process_batch_files(files, min_confidence: float):
    if not files:
        return [], "Upload files first."

    paths   = [f.name for f in files]
    results = process_batch(paths)

    summary_rows = []
    for res in results:
        if "error" in res:
            summary_rows.append([res["path"], "ERROR", res["error"], 0, 0, 0, 0])
        else:
            c = res["counts"]
            summary_rows.append([
                res["meta"]["filename"],
                res["meta"]["type"],
                res["meta"]["word_count"],
                c.get("Person", 0),
                c.get("Organization", 0),
                c.get("Location", 0),
                sum(c.values()),
            ])

    _state["batch_results"] = results
    all_ents = aggregate_entities(results)
    _state["entities"] = all_ents
    _state["text"]     = " ".join(r.get("text", "") for r in results)
    _state["meta"]     = {"filename": "batch", "type": "Batch", "pages": "N/A",
                          "word_count": sum(r.get("meta", {}).get("word_count", 0) for r in results)}

    headers = ["File", "Type", "Words", "Persons", "Orgs", "Locations", "Total"]
    df      = pd.DataFrame(summary_rows, columns=headers)
    info    = f"Processed {len(results)} files. Total entities: {len(all_ents)}. Results available in Analytics & Export tabs."
    return df, info


# ════════════════════════════════════════════════════════════════════════════
#  BUILD GRADIO UI
# ════════════════════════════════════════════════════════════════════════════
with gr.Blocks(
    title="SmartDoc NER Pro",
    theme=gr.themes.Soft(primary_hue="teal", secondary_hue="blue"),
) as app:

    gr.Markdown("""
    # 🔍 SmartDoc NER Pro
    **Advanced Named Entity Recognition** — Persons · Organizations · Locations · Miscellaneous

    **Model:** `dslim/bert-large-NER` (F1 ≈ 92.8% on CoNLL-2003)
    **Supports:** Text · PDF · DOCX · TXT · Batch Processing · Export (CSV / JSON / DOCX)

    *Built by Nafees Ahmad | PAF-IAST Pakistan*
    """)

    with gr.Tabs():

        # ── Tab 1: Text Input ──────────────────────────────────────────────
        with gr.Tab("📝 Text NER"):
            gr.Markdown("### Enter text directly to extract named entities")
            with gr.Row():
                with gr.Column(scale=2):
                    txt_input   = gr.Textbox(
                        label="Input Text",
                        placeholder="Paste any text here... e.g. 'Nafees Ahmad studied at PAF-IAST in Haripur, Pakistan.'",
                        lines=8
                    )
                    txt_min_conf = gr.Slider(50, 100, value=70, step=5, label="Min Confidence (%)")
                    txt_btn     = gr.Button("Extract Entities", variant="primary")
                with gr.Column(scale=1):
                    txt_status  = gr.Textbox(label="Status", interactive=False)
                    txt_stats   = gr.Textbox(label="Statistics", lines=3, interactive=False)

            txt_html   = gr.HTML(label="Highlighted Text")
            txt_table  = gr.DataFrame(label="Extracted Entities", interactive=False)

            txt_btn.click(
                process_text,
                inputs=[txt_input, txt_min_conf],
                outputs=[txt_html, txt_table, txt_stats, txt_status]
            )

            # Quick example texts
            gr.Examples(
                examples=[
                    ["Nafees Ahmad is a Software Engineering student at PAF-IAST in Haripur Hazara, Pakistan. He won the Pak Angels Generative AI Hackathon in 2024 and completed internships at TIERS Limited and Advanced Telecom Services.", 70],
                    ["Elon Musk, CEO of Tesla and SpaceX, announced a new factory in Austin, Texas. Amazon's Jeff Bezos also unveiled plans for a facility in Berlin, Germany.", 70],
                    ["The World Health Organization (WHO) and UNICEF signed a new agreement in Geneva, Switzerland, to support healthcare initiatives in South Asia and sub-Saharan Africa.", 70],
                ],
                inputs=[txt_input, txt_min_conf]
            )

        # ── Tab 2: Document Upload ─────────────────────────────────────────
        with gr.Tab("📄 Document NER"):
            gr.Markdown("### Upload a PDF, DOCX, or TXT file")
            with gr.Row():
                doc_file     = gr.File(label="Upload Document", file_types=[".pdf", ".docx", ".txt"])
                doc_min_conf = gr.Slider(50, 100, value=70, step=5, label="Min Confidence (%)")
                doc_btn      = gr.Button("Process Document", variant="primary")
            doc_info    = gr.Textbox(label="Document Info", interactive=False)
            doc_html    = gr.HTML(label="Highlighted Text (first 3000 chars)")
            doc_table   = gr.DataFrame(label="Extracted Entities", interactive=False)

            doc_btn.click(
                process_document,
                inputs=[doc_file, doc_min_conf],
                outputs=[doc_html, doc_table, doc_info]
            )

        # ── Tab 3: Batch Processing ────────────────────────────────────────
        with gr.Tab("📦 Batch Processing"):
            gr.Markdown("### Process multiple documents at once")
            batch_files    = gr.File(label="Upload Multiple Files", file_count="multiple",
                                     file_types=[".pdf", ".docx", ".txt"])
            batch_min_conf = gr.Slider(50, 100, value=70, step=5, label="Min Confidence (%)")
            batch_btn      = gr.Button("Process All", variant="primary")
            batch_info     = gr.Textbox(label="Batch Status", interactive=False)
            batch_table    = gr.DataFrame(label="Per-Document Summary", interactive=False)

            batch_btn.click(
                process_batch_files,
                inputs=[batch_files, batch_min_conf],
                outputs=[batch_table, batch_info]
            )

        # ── Tab 4: Search & Filter ─────────────────────────────────────────
        with gr.Tab("🔎 Search & Filter"):
            gr.Markdown("### Search and filter entities from the last processed document")
            with gr.Row():
                sf_query  = gr.Textbox(label="Search keyword", placeholder="e.g. Ahmad, Microsoft, Pakistan")
                sf_labels = gr.CheckboxGroup(
                    ["Person", "Organization", "Location", "Miscellaneous"],
                    label="Filter by type", value=[]
                )
                sf_conf   = gr.Slider(0, 100, value=0, step=5, label="Min Confidence (%)")
            sf_btn   = gr.Button("Search", variant="primary")
            sf_info  = gr.Textbox(label="Results", interactive=False)
            sf_table = gr.DataFrame(label="Filtered Entities", interactive=False)

            sf_btn.click(
                search_and_filter,
                inputs=[sf_query, sf_labels, sf_conf],
                outputs=[sf_table, sf_info]
            )

        # ── Tab 5: Analytics ───────────────────────────────────────────────
        with gr.Tab("📊 Analytics"):
            gr.Markdown("### Entity analytics from the last processed document or batch")
            an_btn = gr.Button("Run Analytics", variant="primary")
            with gr.Row():
                an_top   = gr.Textbox(label="Top Entities per Type", lines=15, interactive=False)
                an_pairs = gr.Textbox(label="Person–Organization Co-occurrences", lines=15, interactive=False)
            an_stats = gr.Textbox(label="Confidence & Distribution Statistics", lines=12, interactive=False)

            an_btn.click(
                show_analytics,
                inputs=[],
                outputs=[an_top, an_pairs, an_stats]
            )

        # ── Tab 6: Export ─────────────────────────────────────────────────
        with gr.Tab("💾 Export"):
            gr.Markdown("### Export extracted entities (from last processed document)")
            with gr.Row():
                ex_type = gr.Radio(["CSV", "JSON", "Annotated DOCX"], label="Export format", value="CSV")
                ex_btn  = gr.Button("Export", variant="primary")
            ex_info = gr.Textbox(label="Export status", interactive=False)
            ex_file = gr.File(label="Download exported file")

            ex_btn.click(
                do_export,
                inputs=[ex_type],
                outputs=[ex_file, ex_info]
            )

    gr.Markdown("""
    ---
    **Legend:**
    <span style='background:#4FC3F7;padding:2px 8px;border-radius:4px'>Person</span>
    <span style='background:#81C784;padding:2px 8px;border-radius:4px;margin-left:6px'>Organization</span>
    <span style='background:#FFB74D;padding:2px 8px;border-radius:4px;margin-left:6px'>Location</span>
    <span style='background:#CE93D8;padding:2px 8px;border-radius:4px;margin-left:6px'>Miscellaneous</span>
    """)


if __name__ == "__main__":
    app.launch()