import gradio as gr
import pandas as pd
from src.utils import get_rag_chain


rag = get_rag_chain()


# Write a function to process the RAG results
def query_fc(query):
    # query = "Is Africa the youngest continent in the world?"
    result = rag.invoke(query)
    docs = [doc.metadata for doc in result['source_documents']]
    df = pd.DataFrame(docs)

    df.url = df.apply(lambda x: "<a href='{}'>{}</a>".format(x.url, x.title),
                      axis=1)
    df['publisher'] = df.apply(lambda x: "<a href='https://{}'>{}</a>".
                               format(x.publisher_site, x.publisher_name), axis=1)
    df.drop(columns=['language_code', 'title', 'claim_date', 'review_date',
                     'publisher_site', 'publisher_name'], inplace=True)
    df.rename(columns={'url': 'FC article', 'claim': 'Claim', 'publisher': 'FC Publisher',
                       'claimant': 'Claimant', 'textual_rating': 'FC Rating'},
              inplace=True)

    # Reorder the columns in the DataFrame
    column_order = ['Claim', 'FC Rating', 'FC article', 'FC Publisher', 'Claimant']
    df = df.reindex(columns=column_order)

    return (result['result'],
            "<div style='max-width:100%; max-height:360px; overflow:auto'>"
            + df.to_html(index=False, escape=False) + "</div>")


app = gr.Interface(
    fn=query_fc,
    inputs=gr.Textbox(placeholder="Enter your query here...", label='Query'),
    outputs=[
        gr.Textbox(label="Fact-check"),
        gr.HTML(label="Source Documents")],  # FIXME: the label is not showing
    examples=[
        ["Is Joe Biden offering motel stays to undocumented immigrants?"],
        ["Did Justin Trudeau sit in protest in support of the protesting Indian farmers?"],
    ])

if __name__ == "__main__":
    app.launch()