import gradio as gr
from union.remote import UnionRemote
from flytekit.types.file import FlyteFile

remote = UnionRemote()

# Fetch recent executions
recent_executions = remote.recent_executions(limit=10)
executions = [
    e for e in recent_executions if e.spec.launch_plan.name == "bert-fine-tune.bert_ft"
]
recent_ex_id = executions[0].id.name
execution = remote.fetch_execution(name=recent_ex_id)

ft_llm = execution.outputs["o0"].remote_source
model_cache = execution.outputs["o1"].remote_source

predict_task = remote.fetch_task(name="bert-fine-tune.predict_sentiment")


def execute_flyte_task(text):
    inputs = {
        "text": text,
        "model_cache_dir": model_cache,
        "model": FlyteFile(ft_llm)
    }
    execution = remote.execute(predict_task, inputs=inputs, wait=True)

    # Extract the response
    response = execution.outputs["o0"]

    # Format the response
    sentiment = response['label']
    score = response['score']

    # Determine color based on sentiment
    color = "red" if sentiment == "NEGATIVE" else "green"

    # Format the output HTML
    output_html = f"""
    <div style="text-align: center;">
        <h2>Sentiment: <span style="color: {color};">{sentiment}</span></h2>
        <p>Confidence Score: {score:.2f}</p>
    </div>
    """

    return output_html


# Launch Gradio app
iface = gr.Interface(
    fn=execute_flyte_task,
    inputs=["text"],
    outputs=gr.HTML(),  # Change output to HTML for better formatting
    live=False,
)

iface.launch(debug=True)