import os
import gradio as gr
from huggingface_hub import InferenceClient

# Obtendo a API Key da variável de ambiente
api_key = os.getenv("HF_API_KEY")
if not api_key:
    raise EnvironmentError("A variável de ambiente 'HF_API_KEY' não está configurada.")

# Criação do cliente com a API Key
client = InferenceClient(api_key=api_key)

# Função para gerar a peça jurídica em português
def gerar_peca_juridica_criminal(informacoes_caso, tipo_peca):
    """
    Gera uma peça jurídica criminal em português com base no tipo e nas informações fornecidas.

    Parâmetros:
        informacoes_caso (str): Descrição detalhada do caso.
        tipo_peca (str): Tipo de peça jurídica a ser gerada (ex: habeas corpus, petição inicial).
    
    Retorna:
        str: Peça jurídica gerada.
    """
    mensagens = [
        {
            "role": "system",
            "content": (
                "Você é um especialista em direito criminal no Brasil e deve gerar peças jurídicas detalhadas. "
                "Todas as respostas devem estar em português, considerando o ordenamento jurídico brasileiro."
            )
        },
        {
            "role": "user",
            "content": f"Quero que você elabore um(a) {tipo_peca} com base nas seguintes informações do caso: {informacoes_caso}"
        }
    ]

    # Chamada ao modelo para gerar a resposta
    resposta = client.chat.completions.create(
        model="mistralai/Mistral-7B-Instruct-v0.3",
        messages=mensagens,
        max_tokens=1000
    )

    return resposta.choices[0].message["content"]

# Configuração do app Gradio
def gerar_interface_gradio():
    """
    Configura a interface do Gradio para gerar peças jurídicas.
    """
    with gr.Blocks() as interface:
        gr.Markdown("## Gerador de Peças Jurídicas Criminais")
        gr.Markdown(
            "Insira as informações detalhadas do caso e escolha o tipo de peça jurídica que deseja gerar."
        )
        
        with gr.Row():
            entrada_caso = gr.Textbox(
                label="Informações do Caso",
                placeholder="Descreva o caso detalhadamente aqui...",
                lines=5
            )
            entrada_tipo = gr.Textbox(
                label="Tipo de Peça Jurídica",
                placeholder="Exemplo: habeas corpus, petição inicial..."
            )
        
        botao_gerar = gr.Button("Gerar Peça Jurídica")
        saida_texto = gr.Textbox(label="Peça Jurídica Gerada", lines=15)

        botao_gerar.click(
            gerar_peca_juridica_criminal, 
            inputs=[entrada_caso, entrada_tipo], 
            outputs=saida_texto
        )
    
    return interface

# Executa o app
if __name__ == "__main__":
    app = gerar_interface_gradio()
    app.launch()