import os
from openai import AzureOpenAI
import gradio as gr

client = AzureOpenAI(
    azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"), 
    api_key=os.getenv("AZURE_OPENAI_API_KEY"),  
    api_version="2024-02-01"
)

# Define three prompt options with user-friendly names
prompts = [
    "be a helpful assistant",
    "Crie um resumo de aproximadamente uma página da sentença abaixo em linguagem simples e acessível, sem jargões. O resumo deve conter: 1) Tipo de sentença; 2) fatos relevantes: 2.1) data do crime; 2.2) Local do crime; 2.3 vítimas; 2.4) descrição do fato apurado; 3) A decisão 4) A fundamentação.",
    """Abaixo consta uma transcrição de uma oitiva judicial. Faça uma citação indireta do que a depoente respondeu, separado por ponto e vírgula. O resultado deve constar tudo que foi dito pela testemunha, podendo ser omitido apenas o que for repetitivo.  O texto não deve ser interrompido até exaurir todo o conteúdo da oitiva.
    exemplo de reposta: Estava trabalhando durante os fatos; Trabalhava em um depósito de cerveja ao lado de um posto de gasolina; Quando soube dos fatos a vítima estava no hospital; Fechou o comércio e foi até o Hospital de Acopiara, todavia a vítima foi transferida; Ouviu dizer que Leandro chegou alterado e começou a espancar Vitória, tendo a mãe da adolescente a defendido; Neste momento o réu sacou um punhal e perfurou a ofendida; Não sabe de nenhuma desavença entre a vítima e o réu; Acredita que o réu estava "virado" de uma vaquejada; Conhece o réu apenas de vista; É irmão e sobrinho de Luciene."""
]

# Initialize a dictionary to store chat histories per user
chat_histories = {}

def predict(message, prompt_index):
    history = chat_histories.setdefault(prompt_index, [])

    selected_prompt = prompts[prompt_index]
    history_openai_format = [{"role": "system", "content": selected_prompt}]
    for human, assistant in history:
        history_openai_format.append({"role": "user", "content": human})
        history_openai_format.append({"role": "assistant", "content": assistant})
    history_openai_format.append({"role": "user", "content": message})

    response = client.chat.completions.create(
        model='gpt-4o-mini',
        messages=history_openai_format,
        temperature=1.0,
        stream=True
    )

    partial_message = ""
    for chunk in response:
        if chunk.choices and chunk.choices[0].delta.content is not None:
            partial_message += chunk.choices[0].delta.content
            yield partial_message

    chat_histories[prompt_index].append((message, partial_message))

def predict_wrapper(message, prompt_index):
    return list(predict(message, prompt_index))[-1]

def clear_history():
    global chat_histories
    chat_histories = {}

# Gradio Blocks interface
with gr.Blocks() as demo:
    with gr.Tab("Genérica"):
        message_input = gr.Textbox(label="Message")
        output = gr.Textbox()
        message_input.submit(predict_wrapper, inputs=[message_input, gr.State(0)], outputs=output)
    
    with gr.Tab("Resumir Sentença em linguagem simples"):
        message_input = gr.Textbox(label="Message")
        output = gr.Textbox()
        message_input.submit(predict_wrapper, inputs=[message_input, gr.State(1)], outputs=output)
    
    with gr.Tab("Transformar Transcrição de oitiva judicial em citação indireta"):
        message_input = gr.Textbox(label="Message")
        output = gr.Textbox()
        message_input.submit(predict_wrapper, inputs=[message_input, gr.State(2)], outputs=output)
    
    with gr.Tab("Clear Memory"):
        clear_button = gr.Button("Clear Memory")
        clear_button.click(fn=clear_history, inputs=[], outputs=[])

try:
    demo.launch(auth=([(os.getenv("username"), os.getenv("password"))]))
finally:
    clear_history()