import gradio as gr
from functools import partial
from transformers import pipeline, pipelines


######################
##### INFERENCE ######
######################
# Text Analysis
def cls_inference(input: list[str], pipe: pipeline) -> dict:
    results = pipe(input, top_k=None)
    return {x["label"]: x["score"] for x in results}


# POSP
def tagging(text: str, pipe: pipeline):
    output = pipe(text)
    return {"text": text, "entities": output}


# Text Analysis
def text_analysis(text, pipes: list[pipeline]):
    outputs = []
    for pipe in pipes:
        if isinstance(pipe, pipelines.token_classification.TokenClassificationPipeline):
            outputs.append(tagging(text, pipe))
        else:
            outputs.append(cls_inference(text, pipe))
    return outputs


######################
##### INTERFACE ######
######################
def text_interface(pipe: pipeline, examples: list[str], output_label: str, title: str, desc: str):
    return gr.Interface(
        fn=partial(cls_inference, pipe=pipe),
        inputs=[
            gr.Textbox(lines=5, label="Input Text"),
        ],
        title=title,
        description=desc,
        outputs=[gr.Label(label=output_label)],
        examples=examples,
        allow_flagging="never",
    )


def token_classification_interface(pipe: pipeline, examples: list[str], output_label: str, title: str, desc: str):
    return gr.Interface(
        fn=partial(tagging, pipe=pipe),
        inputs=[
            gr.Textbox(placeholder="Masukan kalimat di sini...", label="Input Text"),
        ],
        outputs=[gr.HighlightedText(label=output_label)],
        title=title,
        examples=examples,
        description=desc,
        allow_flagging="never",
    )


def text_analysis_interface(pipe: list, examples: list[str], output_label: str, title: str, desc: str):
    with gr.Blocks() as text_analysis_interface:
        gr.Markdown(title)
        gr.Markdown(desc)
        input_text = gr.Textbox(lines=5, label="Input Text")
        with gr.Row():
            outputs = [
                (
                    gr.HighlightedText(label=label)
                    if isinstance(p, pipelines.token_classification.TokenClassificationPipeline)
                    else gr.Label(label=label)
                )
                for label, p in zip(output_label, pipe)
            ]
        btn = gr.Button("Analyze")
        btn.click(
            fn=partial(text_analysis, pipes=pipe),
            inputs=[input_text],
            outputs=outputs,
        )
        gr.Examples(
            examples=examples,
            inputs=input_text,
            outputs=outputs,
        )
    return text_analysis_interface