import gradio as gr
from collinear import Collinear
import os
import json
from openai import AsyncOpenAI
from jinja2 import Template
collinear = Collinear(access_token=os.getenv('COLLINEAR_API_KEY'))
prompt = Template("""
iven the following QUESTION, DOCUMENT and ANSWER you must analyze the provided answer and determine whether it is faithful to the contents of the DOCUMENT. The ANSWER must not offer new information beyond the context provided in the DOCUMENT. The ANSWER also must not contradict information provided in the DOCUMENT. Output your final verdict by strictly following this format: "PASS" if the answer is faithful to the DOCUMENT and "FAIL" if the answer is not faithful to the DOCUMENT. Show your reasoning.
--
QUESTION (THIS DOES NOT COUNT AS BACKGROUND INFORMATION):
{{question}}

--
DOCUMENT:
{{context}}

--
ANSWER:
{{answer}}

--
""")

def update_inputs(input_style):
    if input_style == "Dialog":
        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)


async def lynx(input_style_dropdown,document_input,question_input,answer_input):
    if input_style_dropdown=='QA format':
        client = AsyncOpenAI(
        base_url="https://s6mipt5j797e6fql.us-east-1.aws.endpoints.huggingface.cloud/v1/", 
        api_key=os.getenv("HF_TOKEN") 
        )
        rendered_prompt = prompt.render(question=question_input,context=document_input,answer=answer_input)
        rendered_prompt +="""
        
Your output should be in JSON FORMAT with the keys "REASONING" and "SCORE":
{{"REASONING": <your reasoning as bullet points>, "SCORE": <your final score>}}
        """
        chat_completion = await client.chat.completions.create(
            model="tgi",
            messages=[
            {
                "role": "user",
                "content": rendered_prompt
            }
        ],
            top_p=None,
            temperature=None,
            max_tokens=150,
            stream=False,
            seed=None,
            frequency_penalty=None,
            presence_penalty=None
        )
        print(chat_completion)
        return chat_completion.choices.pop().message.content
    else:
        return 'NA'

# Function to judge reliability based on the selected input format
async def judge_reliability(input_style, document, conversation, claim, question, answer):
    if input_style == "Dialog":
        conversation = json.loads(conversation)
        print(conversation)
        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:
    gr.Markdown(
        """
        <p style='text-align: center;color:white'>
        Test Collinear Veritas and compare with Lynx 8B using the sample conversations below or type your own.
        Collinear Veritas can work with any input formats including NLI, QA, and dialog.
        </p>
        """
    )
    with gr.Row():
        input_style_dropdown = gr.Dropdown(label="Input Style", choices=["Dialog", "NLI", "QA format"], value="Dialog", visible=True)

    with gr.Row():
        document_input = gr.Textbox(label="Document", lines=5, visible=True, value="Alex is a good boy. He stays in California")
        conversation_input = gr.Textbox(label="Conversation", lines=5, visible=True, value='[{"role": "user", "content": "Where does Alex stay?"}, {"role": "assistant", "content": "Alex lives in California"}]')
        claim_input = gr.Textbox(label="Claim", lines=5, visible=False, value="Alex lives in California")
        question_input = gr.Textbox(label="Question", lines=5, visible=False, value="Where does Alex stay?")
        answer_input = gr.Textbox(label="Answer", lines=5, visible=False, value="Alex lives in California")

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

        lynx_output = gr.Textbox(label="Lynx Model")


    # 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
    ).then(
        fn=lynx,
        inputs=[input_style_dropdown,document_input,question_input,answer_input],
        outputs=lynx_output
    )

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