import gradio as gr
from typing import Union, Dict, Any, List

# Sample pre-calculated entities
sample_text = (
    "Recent studies on ocean currents from the Global Ocean Temperature Dataset "
    "(GOTD) indicate significant shifts in marine biodiversity."
)
sample_entities = [
    {"label": "named dataset", "text": "Global Ocean Temperature Dataset", "start": 29, "end": 62, "score": 0.99},
    {"label": "acronym", "text": "GOTD", "start": 64, "end": 68, "score": 0.98},
]
rels = [
    'acronym', 'author', 'data description',
    'data geography', 'data source', 'data type',
    'publication year', 'publisher', 'reference year', 'version'
]
MODELS = ["demo-model-1", "demo-model-2"]

# Annotate_query simulation

def annotate_query(
    query: str,
    labels: Union[str, List[str]],
    threshold: float = 0.3,
    nested_ner: bool = False,
    model_name: str = None
) -> Dict[str, Any]:
    # In a real app, you'd call parse_query/inference_pipeline here.
    # For simulation, reuse sample_entities.
    return {
        "text": query,
        "entities": [
            {"start": ent["start"], "end": ent["end"], "label": ent["label"]}
            for ent in sample_entities
        ]
    }

# Build Gradio UI
demo = gr.Blocks()
with demo:
    gr.Markdown(
        """
        ## Step: Annotate Query Simulation
        Enter text (prepopulated) and click **Annotate** to see how entities are highlighted.
        """
    )
    # Inputs
    query = gr.Textbox(lines=3, value=sample_text, label="Input Text")
    entities = gr.Textbox(value=", ".join(rels), label="Relations (unused in simulation)")
    threshold = gr.Slider(0, 1, value=0.3, step=0.01, label="Threshold")
    nested = gr.Checkbox(value=False, label="Nested NER")
    model = gr.Radio(choices=MODELS, value=MODELS[0], label="Model")

    # Outputs
    output_hl = gr.HighlightedText(label="Annotated Entities")

    # Button
    annotate_btn = gr.Button("Annotate")
    annotate_btn.click(
        fn=annotate_query,
        inputs=[query, entities, threshold, nested, model],
        outputs=[output_hl]
    )

demo.launch(debug=True)