# ╔════════════════════════════════════════════════════════════════════════════╗
# ║             PIPELINE v35: GROQ ONLY + DEBUG LOGS (FIX)                    ║
# ║             Usando apenas Groq API (sem Gemini)                           ║
# ╚════════════════════════════════════════════════════════════════════════════╝

import os
import json
import re
import time
from datetime import datetime
import gradio as gr
from groq import Groq

# ==================== 1. CONFIGURAÇÃO ====================
groq_key = os.getenv("GROQ_API_KEY", "SUA_GROQ_KEY_AQUUI")
groq_client = Groq(api_key=groq_key)

ARQUIVO_CONFIG = "protocolo.json"
ARQUIVO_HELP = "help.md"
DELAY_ENTRE_AGENTES = 5  # 3 minutos

print("🚀 App inicializada - GROQ ONLY")
print(f"  ✅ Groq: {'OK - API configurada' if groq_key != 'SUA_GROQ_KEY_AQUUI' else '⚠️ usando key placeholder'}")

# ==================== 2. UTILIDADES ====================
def carregar_protocolo():
    try:
        with open(ARQUIVO_CONFIG, "r", encoding="utf-8") as f:
            return f.read()
    except Exception as e:
        print(f"❌ Erro carregar_protocolo: {e}")
        return "[]"

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

def carregar_help():
    try:
        with open(ARQUIVO_HELP, "r", encoding="utf-8") as f:
            return f.read()
    except:
        return "# Help não encontrado\n\nCrie um arquivo help.md na raiz do projeto."

def ler_anexo(arquivo):
    if arquivo is None:
        return ""
    try:
        with open(arquivo.name, "r", encoding="utf-8") as f:
            conteudo = f.read()
            print(f"📎 Anexo lido: {os.path.basename(arquivo.name)} ({len(conteudo)} chars)")
            return f"\n\n[ANEXO: {os.path.basename(arquivo.name)}]\n{conteudo}\n[FIM ANEXO]\n"
    except Exception as e:
        print(f"❌ Erro ler_anexo {arquivo.name}: {e}")
        return ""

def verificar_stop(texto):
    if not texto:
        return False
    stop_detectado = bool(re.search(r'\bstop\b', str(texto), re.IGNORECASE))
    print(f"🛑 STOP detectado? {stop_detectado} em '{str(texto)[:100]}...'")
    return stop_detectado

# ==================== 3. ENGINE DE EXECUÇÃO (GROQ) ====================
def executar_no(timeline, config):
    print(f"\n🔥 === EXECUTANDO {config['nome']} ===")
    modelo = config.get('modelo', 'meta-llama/llama-4-maverick-17b-128e-instruct')
    print(f"   Modelo Groq: {modelo}")
    
    try:
        inicio = time.time()
        
        # Converte timeline para formato messages do Groq
        messages = []
        
        # System message com missão do agente
        messages.append({
            "role": "system",
            "content": f"AGENTE: {config['nome']}\nMISSÃO: {config['missao']}"
        })
        
        # Adiciona timeline como contexto
        for msg in timeline:
            role = msg.get('role')
            if role in ['user', 'assistant']:
                content = msg.get('content', '')
                # Se content for dict/list, serializa
                if isinstance(content, (dict, list)):
                    content = json.dumps(content, ensure_ascii=False)
                messages.append({
                    "role": role,
                    "content": str(content)
                })
        
        print(f"📤 Groq messages: {len(messages)} mensagens")
        print(f"   System: {messages[0]['content'][:200]}...")
        if len(messages) > 1:
            print(f"   Last user: {messages[-1]['content'][:200]}...")
        
        # Chama Groq API com streaming
        completion = groq_client.chat.completions.create(
            model=modelo,
            messages=messages,
            temperature=1,
            max_completion_tokens=4096,
            top_p=1,
            stream=True,
            stop=None
        )
        
        # Coleta resposta em streaming
        out_raw = ""
        for chunk in completion:
            content_chunk = chunk.choices[0].delta.content or ""
            out_raw += content_chunk
        
        tempo_exec = time.time() - inicio
        
        print(f"📥 OUTPUT GROQ ({len(out_raw)} chars, {tempo_exec:.2f}s):")
        print(out_raw[:1000])
        print("..." if len(out_raw) > 1000 else "")
        
        # Parse JSON se necessário
        content = out_raw
        if config.get('tipo_saida') == 'json':
            try:
                # Remove markdown code blocks se presentes
                cleaned = out_raw.strip()
                 
                content = json.loads(cleaned)
                print("✅ JSON parseado com sucesso")
            except Exception as parse_e:
                print(f"⚠️ Erro parse JSON: {parse_e}")
                print(f"   Primeiros 200 chars: {cleaned[:200]}")
                content = out_raw
        
        print(f"⏱️ Tempo total: {tempo_exec:.2f}s")
        
        return {
            "role": "assistant",
            "agent": config['nome'],
            "content": content,
            "raw": out_raw,
            "tempo": tempo_exec
        }, True
        
    except Exception as e:
        print(f"💥 ERRO GROQ: {str(e)}")
        import traceback
        traceback.print_exc()
        return {
            "role": "system",
            "error": str(e),
            "agent": config['nome']
        }, False

# ==================== 4. ORQUESTRADOR ====================
def orquestrador(texto, anexos_list, history, json_config, contexto_objetivo):
    print("\n" + "="*80)
    print("🎬 INICIANDO ORQUESTRADOR - NOVA EXECUÇÃO")
    print(f"📝 Input texto: '{texto[:200]}...' ({len(texto)} chars)")
    print(f"📎 Anexos: {len(anexos_list)} arquivos")
    
    if not texto.strip():
        print("⚠️ Texto vazio, abortando")
        yield history, []
        return
    
    history = history + [{"role": "user", "content": texto}]
    
    try:
        protocolo = json.loads(json_config)
        print(f"🔗 Protocolo: {len(protocolo)} agentes")
        for i, ag in enumerate(protocolo):
            print(f"   {i+1}. {ag['nome']} ({ag.get('modelo', 'default')})")
    except Exception as e:
        print(f"💥 Erro JSON config: {e}")
        history.append({"role": "assistant", "content": f"❌ Erro no JSON de Configuração: {str(e)}"})
        yield history, []
        return
    
    history.append({"role": "assistant", "content": ""})
    
    # Contexto inicial
    contexto_inicial = ""
    if contexto_objetivo and contexto_objetivo.strip():
        contexto_inicial += f"[OBJETIVO DO MODELO]\n{contexto_objetivo.strip()}\n[FIM OBJETIVO]\n\n"
        print(f"🎯 Objetivo adicionado: {len(contexto_objetivo)} chars")
    
    if anexos_list:
        for anexo in anexos_list:
            anexo_conteudo = ler_anexo(anexo)
            if anexo_conteudo:
                contexto_inicial += anexo_conteudo
    
    full_input = f"{contexto_inicial}{texto}".strip()
    timeline = [{"role": "user", "content": full_input}]
    print(f"🌐 Timeline inicial: {len(full_input)} chars")
    
    audit_data = []
    
    # Loop pelos agentes
    for idx, cfg in enumerate(protocolo):
        print(f"\n{'='*50}")
        print(f"🚀 FASE {idx+1}/{len(protocolo)}: {cfg['nome']}")
        
        history[-1]["content"] = f"⏳ Chamando agente: **{cfg['nome']}**... Aguarde."
        yield history, audit_data
        
        print(f"💤 Delay entre agentes: {DELAY_ENTRE_AGENTES}s")
        time.sleep(DELAY_ENTRE_AGENTES)
        
        history[-1]["content"] = f"⏳ Agente **{cfg['nome']}** processando..."
        yield history, audit_data
        
        res, sucesso = executar_no(timeline, cfg)
        timeline.append(res)
        
        audit_entry = {
            "step": idx + 1,
            "agent": cfg['nome'],
            "model": cfg.get('modelo', 'default'),
            "type": cfg.get('tipo_saida', 'texto'),
            "response_preview": res.get('content' , ''),
            "raw_len": len(res.get('raw', '')),
            "tempo": round(res.get('tempo', 0), 2),
            "sucesso": sucesso,
            "timestamp": datetime.now().strftime('%H:%M:%S')
        }
        audit_data.append(audit_entry)
        print(f"📋 Audit #{audit_entry['step']} salvo")
        
        conteudo_resposta = res.get('content', '')
        if verificar_stop(conteudo_resposta):
            print("🛑 STOP detectado - encerrando pipeline")
            history[-1]["content"] = str(conteudo_resposta) if isinstance(conteudo_resposta, str) else json.dumps(conteudo_resposta, ensure_ascii=False, indent=2)
            yield history, audit_data
            return
        
        # Se for último agente ou tipo texto, exibe resposta
        if idx == len(protocolo) - 1 or cfg.get('tipo_saida') == 'texto':
            texto_final = str(conteudo_resposta) if isinstance(conteudo_resposta, str) else json.dumps(conteudo_resposta, ensure_ascii=False, indent=2)
            print(f"✨ Exibindo resposta final: {len(texto_final)} chars")
            
            # Typewriter effect
            for i in range(0, len(texto_final), 5):
                history[-1]["content"] = texto_final[:i+5]
                yield history, audit_data
                time.sleep(0.05)
            
            history[-1]["content"] = texto_final
            yield history, audit_data
    
    print("🏁 Pipeline concluída com sucesso")
    print("="*80)

# ==================== 5. UI ====================
def ui_clean():
    config_init = carregar_protocolo()
    help_init = carregar_help()
    
    with gr.Blocks(title="AI Forensics - Groq") as app:
        anexos_state = gr.State([])
        
        with gr.Tabs():
            with gr.Tab("💬 Chat"):
                gr.Markdown("## Investigador AI (v35 - Groq Only)")
                chatbot = gr.Chatbot(label="Histórico", height=500)
                
                with gr.Row():
                    txt_in = gr.Textbox(show_label=False, placeholder="Digite sua mensagem...", lines=2, scale=9)
                    btn_send = gr.Button("📤 Enviar", variant="primary", scale=1)
            
            with gr.Tab("📎 Anexos & Contexto"):
                gr.Markdown("""
                ## Anexos e Contexto Factual
                Carregue arquivos que serão enviados **uma vez** antes da primeira operação.
                """)
                objetivo_text = gr.Textbox(
                    label="Objetivo do Modelo",
                    placeholder="Ex: Você é um analista forense imparcial...",
                    lines=5
                )
                gr.Markdown("### Anexos")
                anexos_upload = gr.File(
                    file_count="multiple",
                    file_types=[".txt", ".md", ".csv", ".json"]
                )
                anexos_display = gr.Textbox(label="Arquivos Carregados", interactive=False, lines=3)
                
                def atualizar_anexos(files):
                    if not files:
                        return [], "Nenhum arquivo carregado"
                    nomes = [os.path.basename(f.name) for f in files]
                    return files, f"📎 {len(files)} arquivo(s): {', '.join(nomes)}"
                
                anexos_upload.change(atualizar_anexos, anexos_upload, [anexos_state, anexos_display])
            
            with gr.Tab("⚙️ Protocolo"):
                gr.Markdown("""
                ## Edição do Protocolo de Agentes (Groq)
                
                **Modelos disponíveis:**
                - `meta-llama/llama-4-maverick-17b-128e-instruct` (padrão)
                - `meta-llama/llama-3.3-70b-versatile`
                - `deepseek-r1-distill-llama-70b`
                - `llama-3.1-70b-versatile`
                
                **Exemplo:**
                ```
                [
                  {
                    "nome": "Alien Escrivão",
                    "modelo": "meta-llama/llama-4-maverick-17b-128e-instruct",
                    "tipo_saida": "json",
                    "missao": "Analisar BO de forma imparcial..."
                  }
                ]
                ```
                """)
                
                with gr.Row():
                    btn_save_proto = gr.Button("💾 Salvar", variant="primary", size="sm")
                    btn_reload_proto = gr.Button("🔄 Recarregar", size="sm")
                proto_status = gr.Markdown("")
                code_json = gr.Code(value=config_init, language="json", lines=30)
                
                btn_save_proto.click(salvar_protocolo, code_json, proto_status)
                btn_reload_proto.click(lambda: carregar_protocolo(), outputs=code_json)
            
            with gr.Tab("🔍 Auditoria"):
                gr.Markdown("## Auditoria da Última Execução")
                audit_display = gr.JSON(label="Dados de Auditoria", value=[])
            
            with gr.Tab("❓ Ajuda"):
                help_content = gr.Markdown(help_init)
                btn_reload_help = gr.Button("🔄 Recarregar Help")
                btn_reload_help.click(lambda: carregar_help(), outputs=help_content)
        
        # Triggers
        btn_send.click(
            orquestrador,
            [txt_in, anexos_state, chatbot, code_json, objetivo_text],
            [chatbot, audit_display]
        ).then(lambda: "", outputs=txt_in)
        
        txt_in.submit(
            orquestrador,
            [txt_in, anexos_state, chatbot, code_json, objetivo_text],
            [chatbot, audit_display]
        ).then(lambda: "", outputs=txt_in)
    
    return app

if __name__ == "__main__":
    print("🎉 Lançando app Groq-only...")
    ui_clean().launch(css="footer{display:none!important;}")