import gradio as gr
from datetime import datetime
from huggingface_hub import InferenceClient
import os

class DefenseGenerator:
    def __init__(self):
        self.api_key = os.getenv("HF_API_KEY")
        if not self.api_key:
            raise EnvironmentError("HF_API_KEY não configurada")
        self.client = InferenceClient(api_key=self.api_key)
        
        self.document_types = {
            "resposta_acusacao": "Resposta à Acusação",
            "habeas_corpus": "Habeas Corpus",
            "revogacao_preventiva": "Revogação de Prisão Preventiva",
            "liberdade_provisoria": "Liberdade Provisória",
            "relaxamento_prisao": "Relaxamento de Prisão",
            "apelacao": "Recurso de Apelação",
            "embargos": "Embargos de Declaração"
        }

    def generate_content(self, prompt: str) -> str:
        try:
            response = self.client.chat.completions.create(
                model="mistralai/Mistral-7B-Instruct-v0.3",
                messages=[
                    {"role": "system", "content": "Você é um advogado criminalista experiente. Gere peças sem usar asteriscos ou marcadores. Use numeração romana para seções principais (I, II, III...) e números arábicos para subitens (1, 2, 3...). Mantenha formatação consistente e profissional."},
                    {"role": "user", "content": prompt}
                ],
                max_tokens=2000,
                temperature=0.7
            )
            return response.choices[0].message["content"]
        except Exception as e:
            return f"Erro na geração: {str(e)}"

    def gerar_documento(self, tipo: str, fatos: str, dados_basicos: dict) -> str:
        prompt = f"""Gere uma peça de {self.document_types[tipo]} completa, formatada e sem asteriscos, com base nos seguintes fatos:

{fatos}

A peça deve seguir este formato:

EXCELENTÍSSIMO(A) SENHOR(A) DOUTOR(A) [AUTORIDADE]

[Título da Peça em Maiúsculas]

I - QUALIFICAÇÃO DAS PARTES
[Nome e dados completos]

II - FUNDAMENTOS JURÍDICOS
1. [Primeiro fundamento]
2. [Segundo fundamento]
[...]

III - JURISPRUDÊNCIA APLICÁVEL
1. [Primeira jurisprudência]
2. [Segunda jurisprudência]
[...]

IV - DOUTRINA APLICÁVEL
[Citações doutrinárias relevantes]

V - PRELIMINARES
1. [Primeira preliminar]
2. [Segunda preliminar]
[...]

VI - MÉRITO
[Desenvolvimento dos argumentos]

VII - PEDIDOS
1. [Primeiro pedido]
2. [Segundo pedido]
[...]

Local, Data.

[Advogado]
OAB/[Estado] [Número]

Dados para formatação:
Comarca: {dados_basicos['comarca']}
Vara: {dados_basicos['vara']}
Cliente: {dados_basicos['cliente']}
OAB: {dados_basicos['oab']}"""

        return self.generate_content(prompt)

def create_interface():
    with gr.Blocks(theme=gr.themes.Soft()) as demo:
        gr.Markdown("# Sistema Criminal.AI", elem_classes="text-center text-2xl font-bold mb-4")
        
        with gr.Group():
            tipo = gr.Dropdown(
                choices=list(DefenseGenerator().document_types.keys()),
                value="resposta_acusacao",
                label="Tipo de Peça",
                container=False,
                scale=1
            )
        
        with gr.Row():
            with gr.Column(scale=1):
                comarca = gr.Textbox(
                    label="Comarca",
                    placeholder="Ex: São Paulo",
                    container=False
                )
                vara = gr.Textbox(
                    label="Vara",
                    placeholder="Ex: 1ª",
                    container=False
                )
            
            with gr.Column(scale=1):
                cliente = gr.Textbox(
                    label="Nome do Cliente",
                    placeholder="Ex: João da Silva",
                    container=False
                )
                oab = gr.Textbox(
                    label="OAB",
                    placeholder="Ex: SP123456",
                    container=False
                )
        
        fatos = gr.Textbox(
            label="Descreva os fatos",
            placeholder="Descreva a situação do cliente e o que precisa ser argumentado...",
            lines=5,
            container=False
        )
        
        with gr.Row():
            gerar_btn = gr.Button("Gerar Peça", variant="primary", scale=2)
            limpar_btn = gr.Button("Limpar", scale=1)
            
        output = gr.TextArea(
            label="Documento Gerado",
            lines=20,
            show_copy_button=True,
            container=False
        )

        def process(tipo, fatos, comarca, vara, cliente, oab):
            gerador = DefenseGenerator()
            dados = {
                "comarca": comarca,
                "vara": vara,
                "cliente": cliente,
                "oab": oab
            }
            return gerador.gerar_documento(tipo, fatos, dados)
            
        def clear():
            return {
                tipo: "resposta_acusacao",
                comarca: "",
                vara: "",
                cliente: "",
                oab: "",
                fatos: "",
                output: ""
            }
            
        gerar_btn.click(
            fn=process,
            inputs=[tipo, fatos, comarca, vara, cliente, oab],
            outputs=output,
            api_name="gerar_peca"
        )
        
        limpar_btn.click(
            fn=clear,
            outputs=[tipo, comarca, vara, cliente, oab, fatos, output],
            api_name="limpar"
        )

        exemplos = [
            ["resposta_acusacao", "Cliente preso em flagrante por furto de celular iPhone 13 da vítima Maria em 10/01/2024. É primário, tem residência fixa e trabalho. A vítima não reconheceu com certeza.", "São Paulo", "1ª", "João Silva", "SP123456"],
            ["habeas_corpus", "Cliente preso preventivamente há 90 dias por tráfico. Apreendidas 10g de maconha, é primário, tem residência e trabalho fixos. Não foram apreendidos outros materiais típicos de tráfico.", "São Paulo", "2ª", "Pedro Santos", "SP654321"]
        ]
        
        gr.Examples(
            examples=exemplos,
            inputs=[tipo, fatos, comarca, vara, cliente, oab],
            label="Exemplos de Uso"
        )

    return demo

if __name__ == "__main__":
    demo = create_interface()
    demo.launch()