app.py

import gradio as gr
import os
import json
import logging
from dotenv import load_dotenv
import google.generativeai as genai
from fpdf import FPDF
import re
from datetime import datetime
import PyPDF2
import matplotlib.pyplot as plt
import io

# Configuração de Logs
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Carregar API Key
load_dotenv()
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")

if GEMINI_API_KEY:
    genai.configure(api_key=GEMINI_API_KEY)

# --- CONFIGURAÇÃO IA (USE UM MODELO ESTÁVEL) ---
MODEL_NAME = "gemini-3-flash-preview" 

GENERATION_CONFIG = {
    "temperature": 0.5,
    "max_output_tokens": 15000,
    "response_mime_type": "application/json"
}

# --- LISTAS DE PEÇAS ---
PECAS_TRABALHO = [
    "Aleatória (Simulado Real)", "Reclamação Trabalhista", "Contestação", 
    "Recurso Ordinário", "Contrarrazões", "Inquérito p/ Falta Grave", 
    "Consignação em Pagamento", "Agravo de Petição", "Mandado de Segurança"
]

PECAS_CONSTITUCIONAL = [
    "Aleatória (Simulado Real)", "ADI", "ADC", "ADPF", "ADO", 
    "Mandado de Segurança", "Ação Popular", "Ação Civil Pública", 
    "Recurso Ordinário (ROC)", "Recurso Extraordinário (RE)", "Reclamação"
]

PECAS_PENAL = [
    "Aleatória (Simulado Real)", "Relaxamento de Prisão em Flagrante", "Liberdade Provisória",
    "Queixa-Crime", "Resposta à Acusação", "Memoriais", "Apelação",
    "Recurso em Sentido Estrito (RESE)", "Agravo em Execução", "Revisão Criminal"
]

# --- FUNÇÕES UTILITÁRIAS ---

def limpar_json(texto_bruto):
    if not texto_bruto: return {}
    texto = re.sub(r'^```json', '', texto_bruto, flags=re.MULTILINE)
    texto = re.sub(r'^```', '', texto, flags=re.MULTILINE)
    try:
        return json.loads(texto)
    except:
        match = re.search(r'\{.*\}', texto, re.DOTALL)
        if match: return json.loads(match.group(0))
        return {}

def contar_palavras(texto):
    if not texto: return "0 palavras"
    return f"{len(texto.split())} palavras"

# --- PDF ENGINE ---

class JusTutorPDF(FPDF):
    def header(self):
        self.set_font('Arial', 'B', 14)
        self.set_text_color(50, 50, 50)
        self.cell(0, 10, 'JusTutor Pro - Simulado OAB 2a Fase', 0, 1, 'C')
        self.ln(5)

    def footer(self):
        self.set_y(-15)
        self.set_font('Arial', 'I', 8)
        self.cell(0, 10, f'Pagina {self.page_no()}', 0, 0, 'C')

def safe_text(text):
    if not text: return ""
    if not isinstance(text, str): text = str(text)
    replacements = {'–': '-', '—': '-', '“': '"', '”': '"', '’': "'", '‘': "'", '…': '...'}
    for k, v in replacements.items(): text = text.replace(k, v)
    return text.encode('latin-1', 'replace').decode('latin-1')

def draw_table_row(pdf, item, valor, nota, motivo):
    pdf.set_font("Arial", "", 10)
    line_height = 5
    x_start = pdf.get_x()
    y_start = pdf.get_y()
    
    pdf.multi_cell(90, line_height, safe_text(item), split_only=True)
    h_item = len(pdf.multi_cell(90, line_height, safe_text(item), split_only=True)) * line_height
    pdf.multi_cell(60, line_height, safe_text(motivo), split_only=True)
    h_motivo = len(pdf.multi_cell(60, line_height, safe_text(motivo), split_only=True)) * line_height
    row_height = max(h_item, h_motivo, 10)
    
    if y_start + row_height > 270:
        pdf.add_page()
        y_start = pdf.get_y()
    
    pdf.set_xy(x_start, y_start)
    pdf.multi_cell(90, line_height, safe_text(item), border=1)
    pdf.set_xy(x_start + 90, y_start)
    pdf.cell(20, row_height, safe_text(str(valor)), border=1, align='C')
    pdf.set_xy(x_start + 110, y_start)
    pdf.cell(20, row_height, safe_text(str(nota)), border=1, align='C')
    pdf.set_xy(x_start + 130, y_start)
    pdf.set_font("Arial", "I", 9)
    pdf.multi_cell(60, line_height, safe_text(motivo), border=1)
    pdf.set_y(y_start + row_height)

def gerar_pdf_caderno(area, enc_peca, q1, q2, q3, q4):
    pdf = JusTutorPDF()
    pdf.set_auto_page_break(auto=True, margin=15)
    pdf.add_page()
    pdf.set_font("Arial", "B", 12)
    pdf.cell(0, 10, safe_text(f"CADERNO DE QUESTOES - {area}"), 0, 1, 'C')
    pdf.cell(0, 10, safe_text(f"Gerado em: {datetime.now().strftime('%d/%m/%Y')}"), 0, 1, 'C')
    pdf.ln(10)
    
    def add_q(title, body):
        if not body: return
        pdf.set_font("Arial", "B", 12)
        pdf.set_fill_color(230, 230, 230)
        pdf.cell(0, 10, safe_text(title), 0, 1, 'L', fill=True)
        pdf.ln(2)
        pdf.set_font("Arial", "", 11)
        pdf.multi_cell(0, 6, safe_text(body))
        pdf.ln(8)

    add_q("PECA PRATICO-PROFISSIONAL", enc_peca)
    add_q("QUESTAO 1", q1); add_q("QUESTAO 2", q2); add_q("QUESTAO 3", q3); add_q("QUESTAO 4", q4)
    filename = f"Caderno_Prova_{datetime.now().strftime('%H%M%S')}.pdf"
    pdf.output(filename)
    return filename

def gerar_pdf_completo(area, enc_peca, q1, q2, q3, q4, r_peca, r_q1, r_q2, r_q3, r_q4, dados_correcao, modelo, mentor):
    pdf = JusTutorPDF()
    pdf.set_auto_page_break(auto=True, margin=15)
    pdf.add_page()
    
    # Respostas
    pdf.set_font("Arial", "B", 14)
    pdf.cell(0, 10, safe_text("FOLHA DE RESPOSTAS"), 0, 1, 'C')
    pdf.ln(5)
    
    def add_resp(title, content):
        if not content: content = "[Em branco]"
        pdf.set_font("Arial", "B", 11)
        pdf.cell(0, 8, safe_text(title), 0, 1)
        pdf.set_font("Courier", "", 10)
        pdf.multi_cell(0, 5, safe_text(content), border=1)
        pdf.ln(5)
        
    add_resp("Sua Peca", r_peca); add_resp("Q1", r_q1); add_resp("Q2", r_q2); add_resp("Q3", r_q3); add_resp("Q4", r_q4)
    
    # Tabela
    pdf.add_page()
    pdf.set_font("Arial", "B", 14)
    pdf.cell(0, 10, safe_text("ESPELHO DE CORRECAO"), 0, 1, 'C')
    pdf.ln(5)
    pdf.set_fill_color(200, 200, 200)
    pdf.set_font("Arial", "B", 10)
    pdf.cell(90, 8, "ITEM AVALIADO", 1, 0, 'C', True)
    pdf.cell(20, 8, "VALOR", 1, 0, 'C', True)
    pdf.cell(20, 8, "NOTA", 1, 0, 'C', True)
    pdf.cell(60, 8, "JUSTIFICATIVA", 1, 1, 'C', True)
    
    try:
        if isinstance(dados_correcao, str):
             pdf.set_font("Arial", "", 10); pdf.multi_cell(0, 5, safe_text(dados_correcao))
        else:
            tabela = dados_correcao.get("tabela_notas", [])
            nota_final = dados_correcao.get("nota_final", "0.0")
            for linha in tabela:
                draw_table_row(pdf, linha.get("item", ""), linha.get("valor", ""), linha.get("nota_aluno", ""), linha.get("motivo", ""))
            pdf.ln(5)
            pdf.set_font("Arial", "B", 12)
            pdf.cell(0, 10, safe_text(f"NOTA FINAL: {nota_final} / 10.0"), 0, 1, 'R')
    except Exception as e:
        pdf.multi_cell(0, 5, safe_text(f"Erro tabela: {str(e)}"))

    if modelo:
        pdf.add_page(); pdf.set_font("Arial", "B", 14); pdf.cell(0, 10, safe_text("GABARITO MODELO"), 0, 1, 'C')
        pdf.set_font("Arial", "", 10); pdf.multi_cell(0, 5, safe_text(str(modelo).replace("**", "")))

    if mentor:
        pdf.add_page(); pdf.set_font("Arial", "B", 14); pdf.cell(0, 10, safe_text("MENTOR VADE MECUM"), 0, 1, 'C')
        pdf.set_font("Arial", "", 10); pdf.multi_cell(0, 5, safe_text(str(mentor).replace("**", "")))

    filename = f"Relatorio_Final_{datetime.now().strftime('%H%M%S')}.pdf"
    pdf.output(filename)
    return filename

# --- PROMPTS ---

def get_system_instruction(area, modo="geracao"):
    base = f"Você é a BANCA EXAMINADORA SÊNIOR DA FGV (OAB 2ª FASE - {area}). "
    
    # REGRAS ESPECÍFICAS DE PENAL
    if area == "Direito Penal":
        base += """
        CRITÉRIOS DE PENAL:
        1. Identificação da Peça: Focar no MOMENTO PROCESSUAL (Inquérito, Denúncia, Sentença, Trânsito em Julgado, Execução).
        2. Estrutura: Endereçamento, Preliminares (Nulidade, Extinção da Punibilidade), Mérito (Absolvição, Desclassificação), Subsidiárias (Dosimetria, Regime, Substituição de Pena).
        3. Datar no último dia do prazo (Súmula 710 STF para Processual Penal).
        """

    if modo == "geracao":
        return base + """
        OBJETIVO: Criar uma prova de NÍVEL DIFÍCIL, com múltiplas teses e pegadinhas processuais.
        ESTRUTURA JSON DE SAÍDA:
        {
            "peca_enunciado": "Texto longo (min. 400 palavras)...",
            "q1_enunciado": "...", "q2_enunciado": "...", "q3_enunciado": "...", "q4_enunciado": "...",
            "gabarito_oculto": { ... }
        }
        """
    else: # MODO CORREÇÃO
        return base + """
      OBJETIVO: Agir como o MELHOR PROFESSOR DE CURSINHO PREPARATÓRIO DO BRASIL.
        
        DIRETRIZES PARA O GABARITO MODELO (FUNDAMENTAL):
        1. NÃO RESUMA. Escreva a peça INTEIRA, linha por linha.
        2. Use linguagem jurídica culta e técnica (Ex: "Data venia", "Impende destacar", "Fulcro no art...").
        3. Estruture em Tópicos: I - Preliminares; II - Prejudiciais; III - Mérito; IV - Pedidos.
        4. Na fundamentação, use o SILOGISMO:
           - Premissa Maior (A Lei/Súmula diz X);
           - Premissa Menor (O caso narra Y);
           - Conclusão (Logo, deve-se aplicar Z).
        
        DIRETRIZES PARA CORREÇÃO:
        1. Seja rigoroso na pontuação (0.00 a 10.00).
        2. Se o aluno não citou o artigo exato, desconte ponto.
        
        ESTRUTURA JSON OBRIGATÓRIA:
        {
            "tabela_notas": [
                {"item": "Endereçamento", "valor": "0.10", "nota_aluno": "0.10", "motivo": "Correto"},
                {"item": "Mérito - Tese X", "valor": "0.50", "nota_aluno": "0.00", "motivo": "Não citou Art. Y"}
            ],
            "nota_final": "4.50",
            "gabarito_modelo": "TEXTO COMPLETO DA PEÇA IDEAL (Como se fosse para protocolar no tribunal) + RESPOSTAS COMPLETAS DAS QUESTÕES.",
            "mentor_vade_mecum": "..."
        }
        """

def get_prompt_gerador(area, peca_alvo):
    instrucao = f"Peça obrigatória: {peca_alvo}" if "Aleatória" not in peca_alvo else "Escolha peça aleatória."
    extra_penal = "Para Penal: Defina crime, pena em abstrato e o exato momento processual." if area == "Direito Penal" else ""
    return f"Gere simulado de {area}. {instrucao}. Inclua datas, valores e pegadinhas. {extra_penal}"

def get_prompt_correcao(area, gabarito, respostas):
    return f"""
    Corrija a prova comparando com o Gabarito.
    GABARITO: {json.dumps(gabarito, ensure_ascii=False)}
    RESPOSTAS ALUNO: {json.dumps(respostas, ensure_ascii=False)}
    """

# --- BACKEND ---

def gerar_simulado(area, peca_treino):
    if not GEMINI_API_KEY: return [None]*11 + ["❌ Erro API Key"]
    
    model = genai.GenerativeModel(model_name=MODEL_NAME, system_instruction=get_system_instruction(area, "geracao"), generation_config=GENERATION_CONFIG)
    
    try:
        response = model.generate_content(get_prompt_gerador(area, peca_treino))
        dados = limpar_json(response.text)
        pdf_caderno = gerar_pdf_caderno(area, dados.get("peca_enunciado"), dados.get("q1_enunciado"), dados.get("q2_enunciado"), dados.get("q3_enunciado"), dados.get("q4_enunciado"))
        
        return (
            dados.get("gabarito_oculto"), 
            dados.get("peca_enunciado"), dados.get("peca_enunciado"),  
            dados.get("q1_enunciado"), dados.get("q1_enunciado"),
            dados.get("q2_enunciado"), dados.get("q2_enunciado"),
            dados.get("q3_enunciado"), dados.get("q3_enunciado"),
            dados.get("q4_enunciado"), dados.get("q4_enunciado"),
            gr.update(visible=True), gr.update(interactive=True),
            pdf_caderno, gr.update(visible=True),
            "✅ Prova Gerada! Baixe o caderno."
        )
    except Exception as e:
        logger.error(f"Erro: {e}")
        return [None]*15 + [f"❌ Erro: {str(e)}"]

def corrigir_simulado(gabarito, area, r_peca, r_q1, r_q2, r_q3, r_q4, enc_peca, enc_q1, enc_q2, enc_q3, enc_q4):
    if not gabarito: return [None]*4 + ["⚠️ Gere a prova antes."]
    respostas = {"peca": r_peca, "q1": r_q1, "q2": r_q2, "q3": r_q3, "q4": r_q4}
    model = genai.GenerativeModel(model_name=MODEL_NAME, system_instruction=get_system_instruction(area, "correcao"), generation_config=GENERATION_CONFIG)
    
    try:
        response = model.generate_content(get_prompt_correcao(area, gabarito, respostas))
        dados = limpar_json(response.text)
        
        # Garante o calculo da nota final se a IA falhar
        nota_final = dados.get("nota_final", "0.0")
        if nota_final == "0.0" or nota_final == "":
            try:
                soma = sum([float(x.get('nota_aluno', 0)) for x in dados.get('tabela_notas', [])])
                nota_final = f"{soma:.2f}"
                dados['nota_final'] = nota_final
            except: pass

        # Constrói o Markdown para a tela com a nota em destaque
        tabela_md = f"# 🏆 NOTA FINAL: {nota_final} / 10.0\n\n"
        tabela_md += "| Item Avaliado | Valor Máx | Sua Nota | Comentário |\n|---|---|---|---|\n"
        for l in dados.get("tabela_notas", []):
            tabela_md += f"| {l.get('item','')} | {l.get('valor','')} | {l.get('nota_aluno','')} | {l.get('motivo','')} |\n"
        
        pdf_final = gerar_pdf_completo(area, enc_peca, enc_q1, enc_q2, enc_q3, enc_q4, r_peca, r_q1, r_q2, r_q3, r_q4, dados, dados.get("gabarito_modelo"), dados.get("mentor_vade_mecum"))
        return (tabela_md, dados.get("gabarito_modelo", "Erro"), dados.get("mentor_vade_mecum", "Erro"), pdf_final, "✅ Correção Realizada!")
    except Exception as e:
        return [f"Erro: {e}"]*3 + [None, "❌ Falha"]

# --- ORÁCULO ---
def analisar_pdf_historico(file_obj, area):
    if not file_obj: return None, "⚠️ Envie o PDF."
    keywords = {
        "Direito do Trabalho": {"Reclamação": ["reclamação"], "Contestação": ["contestação"], "RO": ["recurso ordinário"]},
        "Direito Constitucional": {"ADI": ["ação direta"], "MS": ["mandado de segurança"], "Ação Popular": ["ação popular"]},
        "Direito Penal": {"Apelação": ["apelação", "sentença condenatória"], "RESE": ["recurso em sentido estrito", "pronúncia"], "Memoriais": ["memoriais", "audiência de instrução"], "Agravo": ["agravo em execução"]}
    }
    termos = keywords.get(area, {})
    contagem = {k: 0 for k in termos.keys()}
    try:
        reader = PyPDF2.PdfReader(file_obj)
        texto = "".join([p.extract_text().lower() + "\n" for p in reader.pages])
        for peca, lista in termos.items():
            for t in lista: contagem[peca] += texto.count(t)
        fig, ax = plt.subplots(figsize=(8,4))
        ax.barh(list(contagem.keys()), list(contagem.values()), color="#2c3e50")
        plt.tight_layout()
        top = max(contagem, key=contagem.get) if sum(contagem.values()) > 0 else "?"
        return fig, f"### 🔮 Tendência: {top}"
    except Exception as e: return None, f"Erro: {e}"

# --- INTERFACE ---

def interface():
    with gr.Blocks(theme=gr.themes.Soft(), title="JusTutor Ultimate") as demo:
        gabarito_state = gr.State()
        gr.Markdown("# ⚖️ JusTutor Pro: Trabalho, Constitucional e Penal")
        
        with gr.Tabs():
            with gr.TabItem("📝 Sala de Prova"):
                with gr.Row(variant="panel"):
                    area_sel = gr.Dropdown(["Direito do Trabalho", "Direito Constitucional", "Direito Penal"], value="Direito Penal", label="Área")
                    peca_sel = gr.Dropdown(PECAS_PENAL, value="Aleatória (Simulado Real)", label="Modo")
                    btn_gerar = gr.Button("🚀 Iniciar Prova", variant="primary")
                    status = gr.Textbox(label="Status", value="Pronto.", interactive=False)
                
                with gr.Group(visible=False) as area_prova:
                    gr.Markdown("---")
                    with gr.Row():
                        btn_down_caderno = gr.DownloadButton("📥 Baixar Caderno de Questões (PDF)", visible=False)
                    
                    with gr.Row():
                        with gr.Column(scale=1): 
                            gr.Markdown("### 📘 Caderno de Questões")
                            with gr.Accordion("PEÇA (5.0)", open=True):
                                enc_peca = gr.Markdown("...")
                                enc_peca_st = gr.State()
                            with gr.Accordion("Q1 (1.25)", open=False): 
                                enc_q1 = gr.Markdown(); enc_q1_st = gr.State()
                            with gr.Accordion("Q2 (1.25)", open=False): 
                                enc_q2 = gr.Markdown(); enc_q2_st = gr.State()
                            with gr.Accordion("Q3 (1.25)", open=False): 
                                enc_q3 = gr.Markdown(); enc_q3_st = gr.State()
                            with gr.Accordion("Q4 (1.25)", open=False): 
                                enc_q4 = gr.Markdown(); enc_q4_st = gr.State()
                        
                        with gr.Column(scale=1):
                            gr.Markdown("### ✏️ Folha de Respostas")
                            with gr.Tabs():
                                with gr.TabItem("Peça"):
                                    resp_peca = gr.Textbox(label="Sua Peça", lines=25)
                                    count_lbl = gr.Markdown("0 palavras")
                                    resp_peca.change(contar_palavras, resp_peca, count_lbl)
                                with gr.TabItem("Questões"):
                                    resp_q1 = gr.Textbox(label="Q1", lines=5)
                                    resp_q2 = gr.Textbox(label="Q2", lines=5)
                                    resp_q3 = gr.Textbox(label="Q3", lines=5)
                                    resp_q4 = gr.Textbox(label="Q4", lines=5)
                            
                            btn_corrigir = gr.Button("✅ Entregar e Corrigir", variant="stop", interactive=False)

                    with gr.Group(visible=True) as area_resultados:
                        gr.Markdown("---")
                        with gr.Tabs():
                            with gr.TabItem("📊 Espelho FGV (Tabela)"): res_correcao = gr.Markdown()
                            with gr.TabItem("📚 Gabarito Modelo"): res_modelo = gr.Markdown()
                            with gr.TabItem("🧭 Mentor Vade Mecum"): res_mentor = gr.Markdown()
                        
                        btn_pdf_final = gr.DownloadButton("📄 Baixar Relatório Final", label="Baixar PDF Final")

            with gr.TabItem("🔮 Oráculo"):
                file_oracle = gr.File(label="PDF Provas Anteriores")
                btn_oracle = gr.Button("Analisar")
                plot_oracle = gr.Plot()
                txt_oracle = gr.Markdown()

        def atualiza_lista(area):
            if area == "Direito do Trabalho": return gr.update(choices=PECAS_TRABALHO)
            elif area == "Direito Constitucional": return gr.update(choices=PECAS_CONSTITUCIONAL)
            else: return gr.update(choices=PECAS_PENAL)
        
        area_sel.change(atualiza_lista, area_sel, peca_sel)

        btn_gerar.click(
            fn=lambda: "⏳ Gerando...", outputs=status
        ).then(
            fn=gerar_simulado,
            inputs=[area_sel, peca_sel],
            outputs=[
                gabarito_state, 
                enc_peca_st, enc_peca, enc_q1_st, enc_q1, enc_q2_st, enc_q2, enc_q3_st, enc_q3, enc_q4_st, enc_q4,
                area_prova, btn_corrigir, btn_down_caderno, btn_down_caderno, status
            ]
        )

        btn_corrigir.click(
            fn=lambda: "⏳ Corrigindo...", outputs=status
        ).then(
            fn=corrigir_simulado,
            inputs=[gabarito_state, area_sel, resp_peca, resp_q1, resp_q2, resp_q3, resp_q4, enc_peca_st, enc_q1_st, enc_q2_st, enc_q3_st, enc_q4_st],
            outputs=[res_correcao, res_modelo, res_mentor, btn_pdf_final, status]
        )

        btn_oracle.click(analisar_pdf_historico, [file_oracle, area_sel], [plot_oracle, txt_oracle])

    return demo

if __name__ == "__main__":
    app = interface()
    app.launch()