import gradio as gr
from collinear import Collinear
conv_template = Template(
    """
    # Context:
    {{ document }}
    
    # Claim:
    {% for message in conversation %}
    {{ message.role }}: {{ message.content }}
    {% endfor %}
    """
)

qa_template = Template(
    """
    # Context:
    {{ document }}
    
    # Claim:
    user: {{ question }}
    assistant: {{ answer }}
    """
)

nli_template = Template(
    """
    # Context:
    {{ document }}
    
    # Claim:
    assistant: {{ claim }}
    """
)


# Function to dynamically update inputs based on the input style
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
def judge_reliability(input_style, document, conversation, claim, question, answer):
    # with torch.no_grad():
    #     if input_style == "Conv":
    #         conversation = json.loads(conversation)
    #         text = conv_template.render(document=document, conversation=conversation)
    #     elif input_style == "NLI":
    #         text = nli_template.render(document=document, claim=claim)
    #     elif input_style == "QA format":
    #         text = qa_template.render(document=document, question=question, answer=answer)

    #     print(text)
        
    #     outputs = model_pipeline(text)
    #     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()