app.py

import os
import json
import time
import hashlib
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor, as_completed

import gradio as gr
import google.generativeai as genai

# Dependências para PDF
try:
    import PyPDF2
    PDF_SUPPORT = True
except ImportError:
    PDF_SUPPORT = False
    print("⚠️  PyPDF2 não instalado. Instale com: pip install PyPDF2")

# ==================== 1. CONFIGURAÇÃO ====================

api_key = os.getenv("GOOGLE_API_KEY", "SUA_API_KEY_AQUI")
if api_key and api_key != "SUA_API_KEY_AQUI":
    genai.configure(api_key=api_key)

# Modelos do Gemini
model_flash = genai.GenerativeModel("gemini-flash-latest")
model_pro = genai.GenerativeModel("gemini-flash-latest")

ARQUIVO_CONFIG = "protocolo.json"
PASTA_TRANSCRICOES = "transcricoes"
PAGES_PER_CHUNK = 10
MAX_WORKERS = 5  # Limite de chamadas paralelas

os.makedirs(PASTA_TRANSCRICOES, exist_ok=True)

# ==================== 2. UTILIDADES ====================

def carregar_protocolo():
    """ Carrega o protocolo. Se não existir, cria um com exemplo de STOP. """
    try:
        with open(ARQUIVO_CONFIG, "r", encoding="utf-8") as f:
            return f.read()
    except FileNotFoundError:
        # Protocolo padrão que inclui um agente com a lógica de pergunta ao usuário
        protocolo_padrao = [
            {"nome": "Leitor Inicial", "modelo": "flash", "missao": "Leia o contexto e resuma os fatos principais em 3 a 5 pontos."},
            {
                "nome": "Advogado de Acusação",
                "modelo": "pro",
                "missao": "Com base nos fatos, formule uma pergunta crucial para o usuário para fortalecer um caso. Sua resposta DEVE ser APENAS um JSON no formato: {\"tipo\": \"pergunta_usuario\", \"pergunta\": \"Sua pergunta aqui\"}"
            },
            {"nome": "Analista Final", "modelo": "pro", "missao": "Considere a resposta do usuário e os fatos iniciais para dar um parecer final sobre o caso."}
        ]
        return json.dumps(protocolo_padrao, indent=2)

def salvar_protocolo(conteudo):
    try:
        json.loads(conteudo)
        with open(ARQUIVO_CONFIG, "w", encoding="utf-8") as f:
            f.write(conteudo)
        return "✅ Protocolo salvo com sucesso!"
    except Exception as e:
        return f"❌ Erro ao salvar: {str(e)}"

def limpar_nome_arquivo(nome):
    nome_base = os.path.basename(nome)
    nome_limpo = "".join([c for c in nome_base if c.isalnum() or c in (' ', '.', '_', '-')]).strip()
    return nome_limpo + ".json"

def extrair_texto_pdf(caminho_pdf):
    # (Implementação existente, sem alterações)
    try:
        with open(caminho_pdf, 'rb') as f:
            reader = PyPDF2.PdfReader(f)
            paginas = []
            for i, page in enumerate(reader.pages):
                texto = page.extract_text()
                paginas.append({
                    "numero": i + 1,
                    "texto": texto,
                    "metadata": str(page)[:200]
                })
            return paginas, None
    except Exception as e:
        return None, str(e)


def fragmentar_pdf(paginas, tamanho_chunk=PAGES_PER_CHUNK):
    # (Implementação existente, sem alterações)
    chunks = []
    for i in range(0, len(paginas), tamanho_chunk):
        chunk = paginas[i:i + tamanho_chunk]
        num_inicio = chunk[0]["numero"]
        num_fim = chunk[-1]["numero"]
        
        texto_consolidado = "\n---QUEBRA DE PÁGINA---\n".join(
            [f"[PÁGINA {p['numero']}]\n{p['texto']}" for p in chunk]
        )
        
        chunks.append({
            "id": f"chunk_{num_inicio}_{num_fim}",
            "paginas": f"{num_inicio}-{num_fim}",
            "num_paginas": len(chunk),
            "texto": texto_consolidado,
            "metadata": [p["metadata"] for p in chunk]
        })
    return chunks

def processar_pdf_completo(arquivo_pdf):
    # (Implementação existente, sem alterações)
    if not PDF_SUPPORT:
        return None, "❌ PyPDF2 não disponível"
    
    try:
        paginas, erro = extrair_texto_pdf(arquivo_pdf.name if hasattr(arquivo_pdf, 'name') else arquivo_pdf)
        if erro:
            return None, f"❌ Erro ao ler PDF: {erro}"
        
        chunks = fragmentar_pdf(paginas)
        nome_arquivo = os.path.basename(arquivo_pdf.name if hasattr(arquivo_pdf, 'name') else arquivo_pdf)
        
        return {
            "arquivo": nome_arquivo,
            "total_paginas": len(paginas),
            "total_chunks": len(chunks),
            "chunks": chunks,
            "tipo": "pdf"
        }, None
    except Exception as e:
        return None, f"❌ Erro no processamento: {str(e)}"

def ler_arquivo_texto(arquivo):
    # (Implementação existente, sem alterações)
    if arquivo is None: return None
    try:
        with open(arquivo.name, "r", encoding="utf-8") as f:
            conteudo = f.read()
        return {
            "arquivo": os.path.basename(arquivo.name),
            "conteudo": conteudo,
            "tipo": "texto"
        }
    except: return None

# ==================== 3. PIPELINE DE IA ====================

def transcrever_chunk(chunk_data, config_agentes):
    # (Implementação existente, sem alterações)
    modelo = model_flash
    try:
        if config_agentes and isinstance(config_agentes, list):
            if config_agentes[0].get("modelo") == "pro":
                modelo = model_pro
    except:
        pass
    
    prompt = f"""
ANÁLISE DE DOCUMENTO (OCR/LEITURA):
Transcreva e estruture o conteúdo das páginas {chunk_data['paginas']}.
Texto extraído:
{chunk_data['texto']}

Retorne JSON: {{ "transcricao": "...", "objetos": ["..."], "resumo": "..." }}
"""
    try:
        for tentativa in range(3):
            try:
                resposta = modelo.generate_content(prompt)
                texto_resp = resposta.text.replace("```json", "").replace("```", "")
                return json.loads(texto_resp.strip()), None
            except Exception as inner_e:
                if "429" in str(inner_e):
                    time.sleep(2 * (tentativa + 1))
                    continue
                raise inner_e
    except Exception as e:
        return None, str(e)

# ==================== 4. GERENCIADOR DE ARQUIVOS ====================

class GerenciadorArquivos:
    # (Implementação existente, sem alterações)
    def __init__(self):
        self.arquivos = {} 

    def adicionar(self, arquivo, arquivo_id):
        self.arquivos[arquivo_id] = {
            "arquivo": arquivo,
            "nome": os.path.basename(arquivo.name),
            "status": "adicionado",
            "processado": None,
            "transcricao": None
        }

    def gerar_prompt_com_transcricoes(self, texto_usuario):
        prompt = texto_usuario + "\n\n--- CONTEXTO DOS ARQUIVOS ---\n"
        count = 0
        for _, item in self.arquivos.items():
            if item["status"] == "processado" and item["transcricao"]:
                count += 1
                trans = item["transcricao"]
                nome = item["nome"]
                prompt += f"\n[ARQUIVO: {nome}]\n"
                
                if isinstance(trans, dict) and "chunks_processados" in trans:
                    for chunk in trans["chunks_processados"]:
                        if chunk.get("status") == "OK":
                            resumo = chunk.get('resumo', '')
                            resumo = str(resumo) if resumo else ""
                            prompt += f"Páginas {chunk['paginas']}: {resumo}\n"
                            
                            texto_full = chunk.get('transcricao', '')
                            if texto_full:
                                texto_seguro = str(texto_full)
                                prompt += f"Trecho: {texto_seguro[:400]}...\n"
                            else:
                                prompt += "Trecho: (vazio)\n"

                elif isinstance(trans, dict) and "conteudo" in trans:
                    conteudo = str(trans['conteudo'])
                    prompt += f"Conteúdo: {conteudo[:1000]}...\n"
        
        if count == 0:
            prompt += "(Nenhum arquivo processado ainda)"
        return prompt

# Instância Global
gerenciador = GerenciadorArquivos()

# ==================== 5. FUNÇÕES DE ORQUESTRAÇÃO ====================

def automacao_upload_processamento(files, history, config_json):
    # (Implementação existente, sem alterações)
    if not files:
        return history

    try:
        config_agentes = json.loads(config_json)
    except:
        config_agentes = []

    if history is None:
        history = []

    history.append([None, f"📂 **SISTEMA:** Recebi {len(files)} arquivo(s). Verificando cache e processando..."])
    yield history
    
    ids_para_processar = []
    
    for f in files:
        arquivo_id = f"arq_{int(time.time()*1000)}_{f.name}"
        gerenciador.adicionar(f, arquivo_id)
        ids_para_processar.append(arquivo_id)
    
    for arq_id in ids_para_processar:
        item = gerenciador.arquivos[arq_id]
        nome = item["nome"]
        
        nome_cache = limpar_nome_arquivo(nome)
        caminho_cache = os.path.join(PASTA_TRANSCRICOES, nome_cache)
        
        if os.path.exists(caminho_cache):
            try:
                with open(caminho_cache, "r", encoding="utf-8") as cache_file:
                    dados_cache = json.load(cache_file)
                item["transcricao"] = dados_cache
                item["status"] = "processado"
                if nome.lower().endswith('.pdf') and "chunks_processados" in dados_cache:
                     item["processado"] = {"tipo": "pdf", "chunks": []}
                history.append([None, f"♻️ **Cache Encontrado:** `{nome}` já foi processado. Carregando..."])
                yield history
                continue
            except Exception as e:
                history.append([None, f"⚠️ Erro cache `{nome}`: {e}. Reprocessando..."])

        history.append([None, f"⚙️ **Processando:** `{nome}`..."])
        yield history

        if nome.lower().endswith('.pdf'):
            if not PDF_SUPPORT:
                history.append([None, f"❌ Erro em `{nome}`: Biblioteca PDF ausente."])
                yield history
                continue

            pdf_proc, erro = processar_pdf_completo(item["arquivo"])
            if erro:
                history.append([None, f"❌ Erro em `{nome}`: {erro}"])
                yield history
                continue
            
            item["processado"] = pdf_proc
            chunks = pdf_proc["chunks"]
            total_chunks = len(chunks)
            
            chunks_ordenados = [None] * total_chunks
            
            history.append([None, f"📄 `{nome}` fragmentado em {total_chunks} partes. Iniciando IA (Paralelo: {MAX_WORKERS} threads)..."])
            yield history

            with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
                futures_map = {}
                for i, chunk in enumerate(chunks):
                    future = executor.submit(transcrever_chunk, chunk, config_agentes)
                    futures_map[future] = i
                
                concluidos = 0
                for future in as_completed(futures_map):
                    index_original = futures_map[future]
                    res, err = future.result()
                    
                    if err:
                        chunks_ordenados[index_original] = {"status": "ERRO", "paginas": chunks[index_original]["paginas"]}
                    else:
                        chunks_ordenados[index_original] = {
                            "status": "OK",
                            "paginas": chunks[index_original]["paginas"],
                            "transcricao": res.get("transcricao"),
                            "resumo": res.get("resumo")
                        }
                    
                    concluidos += 1
                    if concluidos % 2 == 0 or concluidos == total_chunks:
                        msg_base = f"📄 `{nome}`: Processando partes... ({concluidos}/{total_chunks})"
                        history[-1][1] = msg_base
                        yield history

            dados_finais = {
                "arquivo": nome,
                "data_processamento": str(datetime.now()),
                "chunks_processados": chunks_ordenados
            }
            
            item["transcricao"] = dados_finais
            item["status"] = "processado"
            
            try:
                with open(caminho_cache, "w", encoding="utf-8") as f_out:
                    json.dump(dados_finais, f_out, indent=2, ensure_ascii=False)
                history.append([None, f"💾 `{nome}` processado e salvo no cache."])
            except Exception as e:
                 history.append([None, f"⚠️ Erro ao salvar cache: {e}"])
            
            yield history

        else:
            res = ler_arquivo_texto(item["arquivo"])
            if res:
                item["processado"] = res
                dados_finais = {"conteudo": res["conteudo"], "data_processamento": str(datetime.now())}
                item["transcricao"] = dados_finais
                item["status"] = "processado"
                
                with open(caminho_cache, "w", encoding="utf-8") as f_out:
                    json.dump(dados_finais, f_out, indent=2, ensure_ascii=False)
                
                history.append([None, f"✅ `{nome}` (Texto) lido e salvo."])
            else:
                history.append([None, f"❌ Falha ao ler `{nome}`."])
            yield history
            
    history.append([None, "🏁 **Processamento de lote finalizado.** Os arquivos estão prontos para análise."])
    yield history


def chat_orquestrador(message, history, config_json, pipeline_state):
    """
    Orquestra a conversa. Pode iniciar uma nova pipeline ou continuar uma que foi pausada.
    """
    
    # --- LÓGICA DE CONTINUAÇÃO ---
    if pipeline_state.get("is_paused"):
        history.append([message, None])
        
        # Recupera o estado
        timeline_execucao = pipeline_state["timeline"]
        agentes_restantes = pipeline_state["remaining_agents"]
        
        # Adiciona a resposta do usuário à trilha de auditoria
        timeline_execucao.append({
            "passo": len(timeline_execucao) + 1,
            "tipo": "resposta_usuario",
            "conteudo": message
        })
        
        # Reseta o estado para evitar loops
        pipeline_state["is_paused"] = False
        
        # Continua a execução do ponto onde parou
        yield from executar_pipeline(history, timeline_execucao, agentes_restantes, pipeline_state)
        return

    # --- LÓGICA DE INÍCIO DE UMA NOVA CONVERSA ---
    try:
        prompt_contexto = gerenciador.gerar_prompt_com_transcricoes(message)
    except Exception as e:
        history.append([message, f"❌ Erro ao gerar contexto: {str(e)}"])
        yield history, [], pipeline_state
        return
    
    try:
        protocolo = json.loads(config_json)
    except:
        history.append([message, "❌ Erro no JSON de Configuração do Protocolo."])
        yield history, [], pipeline_state
        return

    history.append([message, None])
    
    # Inicia uma nova trilha de auditoria
    timeline_execucao = [{"passo": 1, "tipo": "prompt_usuario", "conteudo": prompt_contexto}]
    yield history, timeline_execucao, pipeline_state
    
    # Inicia a execução com todos os agentes do protocolo
    yield from executar_pipeline(history, timeline_execucao, protocolo, pipeline_state)


def executar_pipeline(history, timeline_execucao, agentes_a_executar, pipeline_state):
    """
    Função core que executa a lista de agentes em sequência.
    Pode ser pausada se um agente pedir input do usuário.
    """
    
    passo_atual = len(timeline_execucao) + 1
    
    for i, cfg in enumerate(agentes_a_executar):
        time.sleep(3)
        nome_agente = cfg.get("nome", "Agente")
        modelo_agente = model_pro if cfg.get("modelo") == "pro" else model_flash
        
        msg_atual = history[-1][1] or ""
        history[-1][1] = msg_atual + f"⏳ **{nome_agente}** está analisando...\n"
        yield history, timeline_execucao, pipeline_state
        
        prompt_agente = f"""
--- HISTÓRICO DA CONVERSA ATÉ AGORA ---
{json.dumps(timeline_execucao, ensure_ascii=False, indent=2)}
-----------------
Sua Identidade: {nome_agente}
Sua Missão Específica Agora: {cfg['missao']}
Responda de forma concisa e direta, focando apenas na sua missão.
"""
        try:
            inicio = time.time()
            resp = modelo_agente.generate_content(prompt_agente)
            texto_resp = resp.text
            duracao = time.time() - inicio

            # --- LÓGICA DE PAUSA (STOP) ---
            try:
                # Tenta interpretar a resposta como JSON para verificar se é uma pergunta
                resposta_json = json.loads(texto_resp)
                if resposta_json.get("tipo") == "pergunta_usuario":
                    pergunta = resposta_json.get("pergunta", "Não foi possível extrair a pergunta.")
                    
                    # Salva o estado atual da pipeline
                    pipeline_state["is_paused"] = True
                    pipeline_state["timeline"] = timeline_execucao
                    # Salva os agentes que AINDA NÃO rodaram
                    pipeline_state["remaining_agents"] = agentes_a_executar[i+1:]
                    
                    # Adiciona a pergunta à auditoria e ao chat
                    timeline_execucao.append({"passo": passo_atual, "tipo": "pergunta_agente", "agente": nome_agente, "pergunta": pergunta})
                    msg_atual = history[-1][1].replace(f"⏳ **{nome_agente}** está analisando...\n", "")
                    history[-1][1] = msg_atual + f"**{nome_agente}** precisa de mais informações:\n\n> *{pergunta}*\n\nAguardando sua resposta na caixa de texto abaixo..."
                    
                    # Encerra a execução atual e aguarda o usuário
                    yield history, timeline_execucao, pipeline_state
                    return # Sai da função
            except (json.JSONDecodeError, TypeError):
                # Se não for um JSON de pergunta, é uma resposta normal
                pass
            # ---------------------------

            timeline_execucao.append({"passo": passo_atual, "tipo": "resposta_agente", "agente": nome_agente, "resposta": texto_resp})
            
            msg_atual = history[-1][1].replace(f"⏳ **{nome_agente}** está analisando...\n", "")
            
            # Não exibe o conteúdo da resposta do modelo no chat, apenas a confirmação
            novo_trecho = f"✅ **[{nome_agente}]** concluiu sua análise em ({duracao:.1f}s).\n"
            history[-1][1] = msg_atual + novo_trecho
            yield history, timeline_execucao, pipeline_state
            
        except Exception as e:
            timeline_execucao.append({"passo": passo_atual, "tipo": "erro_agente", "agente": nome_agente, "erro": str(e)})
            msg_atual = history[-1][1]
            history[-1][1] = msg_atual.replace(f"⏳ **{nome_agente}** está analisando...\n", "") + f"\n❌ Erro em {nome_agente}: {str(e)}\n"
            yield history, timeline_execucao, pipeline_state
        
        passo_atual += 1

# ==================== 6. UI (Gradio) ====================

def ui_v29_stop_logic():
    css = """
    footer {display: none !important;}
    .contain {border: none !important;}
    """
    
    config_inicial = carregar_protocolo()

    with gr.Blocks(title="AI Forensics Auto", css=css, theme=gr.themes.Soft()) as app:
        
        # Estado da configuração dos agentes
        state_config = gr.State(config_inicial)
        # NOVO: Estado para controlar a pausa/continuação da pipeline
        pipeline_state = gr.State({"is_paused": False, "timeline": [], "remaining_agents": []})

        with gr.Tabs():
            with gr.Tab("💬 Investigação"):
                
                chatbot = gr.Chatbot(
                    height=400, 
                    show_label=False, 
                    show_copy_button=True,
                    render_markdown=True,
                    label="Chat de Investigação"
                )
                
                with gr.Row():
                    txt_input = gr.Textbox(
                        scale=8, 
                        show_label=False, 
                        placeholder="Digite sua instrução ou responda à pergunta do agente...", 
                        lines=1
                    )
                    btn_enviar = gr.Button("Enviar 📨", variant="primary", scale=1)

                with gr.Accordion("📂 Adicionar Arquivos para Análise", open=False):
                    gr.Markdown("Selecione arquivos (PDF, TXT). A transcrição iniciará **automaticamente**.")
                    file_uploader = gr.File(
                        file_count="multiple",
                        file_types=[".pdf", ".txt", ".json", ".md"],
                        label="Arraste arquivos aqui ou clique para selecionar"
                    )

            with gr.Tab("🕵️ Auditoria"):
                gr.Markdown("### Trilha de Auditoria\nExibe o histórico completo de prompts e respostas de cada agente na última execução.")
                json_audit = gr.JSON(label="Timeline da Execução da Última Mensagem")

            with gr.Tab("⚙️ Contexto & Config"):
                gr.Markdown("### Protocolo dos Agentes\nDefina a sequência e as missões dos agentes de IA. Para pausar e pedir input, use a missão de exemplo para o agente retornar um JSON específico.")
                with gr.Row():
                    btn_save_cfg = gr.Button("💾 Salvar Alterações")
                    lbl_cfg_status = gr.Label(show_label=False)
                
                code_config = gr.Code(value=config_inicial, language="json", label="protocolo.json")
                
                btn_save_cfg.click(salvar_protocolo, inputs=[code_config], outputs=[lbl_cfg_status])
                # Atualiza o state_config em memória após salvar
                btn_save_cfg.click(lambda x: x, inputs=[code_config], outputs=[state_config])

        # Ação de clique agora passa o pipeline_state para o orquestrador
        btn_enviar.click(
            chat_orquestrador,
            inputs=[txt_input, chatbot, state_config, pipeline_state],
            outputs=[chatbot, json_audit, pipeline_state] # Atualiza o estado da pipeline
        ).then(
            lambda: "", outputs=[txt_input]
        )

        file_uploader.upload(
            automacao_upload_processamento,
            inputs=[file_uploader, chatbot, state_config],
            outputs=[chatbot]
        )

    return app

if __name__ == "__main__":
    ui_v29_stop_logic().launch()