import os
import gradio as gr
import openai

from dotenv import load_dotenv
load_dotenv()

llm_api_options = ["OpenAI API","Azure OpenAI API","Google PaLM API", "Llama 2"]
TEST_MESSAGE = "My favorite TV shows are The Mentalist, The Blacklist, Designated Survivor, and Unforgettable. What are ten series that I should watch next?"


def test_handler(optionSelection, prompt: str = "Write an introductory paragraph to explain Generative AI to the reader of this content."):        
    if optionSelection not in llm_api_options:
        raise ValueError("Invalid choice!")

    match optionSelection:
        case  "OpenAI API":
            #model = "gpt-35-turbo"
            model = "gpt-4"
            system_prompt: str = "Explain in detail to help student understand the concept.", 
            assistant_prompt: str = None, 
            
            messages = [
                {"role": "user", "content": f"{prompt}"},
                {"role": "system", "content": f"{system_prompt}"},
                {"role": "assistant", "content": f"{assistant_prompt}"}
            ]
            
            openai.api_key = os.getenv("OPENAI_API_KEY")
            openai.api_version = '2020-11-07'
            
            completion = openai.ChatCompletion.create(
                model = model, 
                messages = messages,
                temperature = 0.7
            )
            
            print(completion)
            response = completion["choices"][0]["message"].content
            print(response)    
            return message, response
        case  "Azure OpenAI API":
            return "", ""
        case  "Google PaLM API":
            return "", ""
        case  "Llama 2":
            return "", ""


with gr.Blocks() as LLMDemoTabbedScreen:
    with gr.Tab("Text-to-Text (Text Completion)"):
        llm_selection = gr.Radio(llm_api_options, label="Select one", info="Which service do you want to use?", value="OpenAI API")
        test_string = gr.Textbox(label="Try String", value=TEST_MESSAGE, lines=2)                   
        test_string_response = gr.Textbox(label="Response")
        test_string_output_info = gr.Label(value="Output Info", label="Info")
        test_button = gr.Button("Try it")

    test_button.click(
        fn=test_handler,
        inputs=[llm_api_options, test_string],
        outputs=[test_string_output_info, test_string_response]
    )
	
	
if __name__ == "__main__":
    LLMDemoTabbedScreen.launch()