import gradio as gr
from collinear import Collinear
import os
collinear = Collinear(access_token=os.getenv('COLLINEAR_API_KEY'))

def update_inputs(input_style):
    if input_style == "Conv":
        return gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
    elif input_style == "NLI":
        return gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)
    elif input_style == "QA format":
        return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)


# Function to judge reliability based on the selected input format
async def judge_reliability(input_style, document, conversation, claim, question, answer):
    if input_style == "Conv":
        outputs= await collinear.judge.veritas.conversation(document,conversation[:-1],conversation[-1])
    elif input_style == "NLI":
        outputs = await collinear.judge.veritas.natural_language_inference(document,claim)
    elif input_style == "QA format":
        outputs = await collinear.judge.veritas.question_answer(document,question,answer)
    results = f"Reliability Judge Outputs: {outputs}"
    return results



# Create the interface using gr.Blocks
with gr.Blocks() as demo:
    with gr.Row():
        input_style_dropdown = gr.Dropdown(label="Input Style", choices=["Conv", "NLI", "QA format"], value="Conv", visible=True)

    with gr.Row():
        document_input = gr.Textbox(label="Document", lines=5, visible=True, value="Chris Voss is one of the best negotiators in the world. And he was born in Iowa, USA.")
        conversation_input = gr.Textbox(label="Conversation", lines=5, visible=True, value='[{"role": "user", "content": "Hi Chris Voss, Where are you born?"}, {"role": "assistant", "content": "I am born in Iowa"}]')
        claim_input = gr.Textbox(label="Claim", lines=5, visible=False, value="CV was born in Iowa")
        question_input = gr.Textbox(label="Question", lines=5, visible=False, value="Where is Chris Voss born?")
        answer_input = gr.Textbox(label="Answer", lines=5, visible=False, value="CV was born in Iowa")

    with gr.Row():
        result_output = gr.Textbox(label="Results")


    # Set the visibility of inputs based on the selected input style
    input_style_dropdown.change(
        fn=update_inputs, 
        inputs=[input_style_dropdown], 
        outputs=[document_input, conversation_input, claim_input, question_input, answer_input]
    )

    # Set the function to handle the reliability check
    gr.Button("Submit").click(
        fn=judge_reliability, 
        inputs=[input_style_dropdown, document_input, conversation_input, claim_input, question_input, answer_input], 
        outputs=result_output
    )

# Launch the demo
if __name__ == "__main__":
    demo.launch()