import gradio as gr
from transformers import pipeline
import torch

MODEL_ID = "samandar1105/named_entity-recognition"

ENTITY_STYLES = {
    "PER":  {"color": "#FF6B6B", "emoji": "👤", "label": "Person"},
    "ORG":  {"color": "#4ECDC4", "emoji": "🏢", "label": "Organization"},
    "LOC":  {"color": "#45B7D1", "emoji": "📍", "label": "Location"},
    "MISC": {"color": "#FFA07A", "emoji": "🏷️", "label": "Miscellaneous"},
}

EXAMPLES = [
    ["Elon Musk, CEO of Tesla and SpaceX, met with Emmanuel Macron at the Élysée Palace in Paris."],
    ["The United Nations Security Council held an emergency meeting at its headquarters in New York City."],
    ["Apple announced a new partnership with Samsung to develop chips in South Korea and Taiwan."],
    ["Cristiano Ronaldo left Manchester United to join Al Nassr, a club based in Riyadh, Saudi Arabia."],
    ["NASA's Artemis program, led by Bill Nelson, plans to land astronauts on the Moon near the South Pole."],
]

print(f"Loading model: {MODEL_ID}")
ner_pipeline = pipeline(
    "token-classification",
    model=MODEL_ID,
    aggregation_strategy="simple",
    device=0 if torch.cuda.is_available() else -1,
)
print("Model loaded!")


def run_ner(text: str):
    if not text or text.strip() == "":
        return "<p style='color: gray;'>Please enter some text above.</p>", ""
    if len(text.strip()) < 5:
        return "<p style='color: orange;'>Text too short. Please enter a full sentence.</p>", ""

    entities = ner_pipeline(text)

    if not entities:
        return f"<p>{text}</p><p style='color: gray;'>No named entities detected.</p>", ""

    highlighted = ""
    last_end = 0
    for ent in entities:
        start, end, label, score = ent["start"], ent["end"], ent["entity_group"], ent["score"]
        style = ENTITY_STYLES.get(label, {"color": "#aaa", "emoji": "•", "label": label})
        highlighted += text[last_end:start]
        highlighted += (
            f'<mark style="background: {style["color"]}33; border: 2px solid {style["color"]}; '
            f'border-radius: 4px; padding: 2px 6px; margin: 0 2px; font-weight: bold;">'
            f'{text[start:end]}'
            f'<sup style="background: {style["color"]}; color: white; border-radius: 3px; '
            f'padding: 0 4px; margin-left: 3px; font-size: 10px;">'
            f'{style["emoji"]} {label}</sup></mark>'
        )
        last_end = end
    highlighted += text[last_end:]
    full_html = f'<div style="font-size:16px;line-height:2.2;padding:15px;background:#f9f9f9;border-radius:8px;">{highlighted}</div>'

    table_rows = ""
    for ent in entities:
        label = ent["entity_group"]
        style = ENTITY_STYLES.get(label, {"color": "#aaa", "emoji": "•", "label": label})
        table_rows += (
            f'<tr><td style="padding:6px 12px;font-weight:bold;">{ent["word"]}</td>'
            f'<td style="padding:6px 12px;"><span style="background:{style["color"]};color:white;'
            f'border-radius:4px;padding:2px 8px;font-size:12px;">{style["emoji"]} {label} — {style["label"]}</span></td>'
            f'<td style="padding:6px 12px;">{ent["score"]:.1%}</td></tr>'
        )
    table_html = f'<table style="width:100%;border-collapse:collapse;margin-top:10px;"><thead><tr style="background:#f0f0f0;"><th style="padding:8px 12px;text-align:left;">Entity</th><th style="padding:8px 12px;text-align:left;">Type</th><th style="padding:8px 12px;text-align:left;">Confidence</th></tr></thead><tbody>{table_rows}</tbody></table>'

    return full_html, table_html


with gr.Blocks(title="Named Entity Recognizer", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🏷️ Named Entity Recognizer
    Paste any English text and this AI will highlight named entities.
    **👤 PER** — People &nbsp; **🏢 ORG** — Organizations &nbsp; **📍 LOC** — Locations &nbsp; **🏷️ MISC** — Miscellaneous
    """)
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(label="📝 Enter text", placeholder="Paste a sentence here...", lines=5)
            with gr.Row():
                submit_btn = gr.Button("🔍 Find Entities", variant="primary", scale=3)
                clear_btn = gr.ClearButton([input_text], value="🗑️ Clear", scale=1)
    gr.Examples(examples=EXAMPLES, inputs=input_text, label="📌 Click an example to try it")
    highlighted_output = gr.HTML(label="📄 Highlighted Text")
    entity_table = gr.HTML(label="📊 Detected Entities")
    gr.Markdown("---\n**Model:** `samandar1105/named_entity-recognition` fine-tuned on CoNLL-2003 | F1: ~95%")
    submit_btn.click(fn=run_ner, inputs=[input_text], outputs=[highlighted_output, entity_table])
    input_text.submit(fn=run_ner, inputs=[input_text], outputs=[highlighted_output, entity_table])

if __name__ == "__main__":
    demo.launch(share=False, server_port=7860)