import re
import gradio as gr

# Simulated GLiNER output
SIMULATED_ENTITIES = {
    "Household Survey Data": "named dataset",
    "Ethiopia":             "data geography",
    "2020":                 "reference year",
    "World Bank":           "publisher"
}

def highlight_text(text):
    """
    Return a dict with the original text and a list of start/end/label spans.
    """
    entities = []
    for phrase, label in SIMULATED_ENTITIES.items():
        # find all occurrences of phrase
        for m in re.finditer(re.escape(phrase), text):
            entities.append({
                "start": m.start(),
                "end":   m.end(),
                "label": label
            })
    return {"text": text, "entities": entities}

with gr.Blocks() as demo:
    gr.Markdown("## Simulated GLiNER Entity Highlighter\n"
                "Type a sentence (or use the default) then hit **Highlight Entities** or press Enter.")
    
    text_input = gr.Textbox(
        label="Input Text",
        lines=3,
        value="The Household Survey Data for Ethiopia in 2020 was published by the World Bank."
    )
    highlight_btn = gr.Button("Highlight Entities")
    highlighted   = gr.HighlightedText(label="Highlighted Entities")
    
    # Wire up both the button and pressing Enter to our highlighter
    highlight_btn.click(fn=highlight_text, inputs=text_input, outputs=highlighted)
    text_input.submit(fn=highlight_text, inputs=text_input, outputs=highlighted)
    # Also do an initial run on load so you immediately see highlights
    demo.load(fn=highlight_text, inputs=text_input, outputs=highlighted)
    
    gr.Markdown("""
**Entity Legend**  
- **named dataset**: Household Survey Data  
- **data geography**: Ethiopia  
- **reference year**: 2020  
- **publisher**: World Bank
""")

if __name__ == "__main__":
    demo.launch(inline=True)