#--- START OF FILE app.py ---

import chainlit as cl
import httpx
import os 

API_URL = "http://localhost:8000"
UPLOAD_URL = f"{API_URL}/upload-document"
CHAT_URL = f"{API_URL}/chat"

RAG_INITIALIZED = False

@cl.on_chat_start
async def start():
    global RAG_INITIALIZED
    RAG_INITIALIZED = False

    await cl.Message(
        content="👋 **Bem-vindo ao RAG Multi-Formato!**\n\nPara começar, faça o upload de um arquivo. Suporto: **PDF, TXT, Markdown (.md) e Python (.py)**.",
        author="Sistema"
    ).send()

    # MUDANÇA: Lista de extensões aceitas atualizada
    files = await cl.AskFileMessage(
        content="📂 Por favor, envie seu arquivo:",
        accept=[".pdf", ".txt", ".md", ".py"],
        max_size_mb=20,
        timeout=180,
    ).send()

    if files:
        file = files[0]
        msg = cl.Message(content=f"⏳ Processando '{file.name}'...", author="Sistema")
        await msg.send()
        
        try:
            # Pega o MIME type ou define genericamente como octet-stream se falhar, 
            # o backend agora confia mais na extensão do arquivo.
            mime_type = getattr(file, 'type', 'application/octet-stream')
            
            with open(file.path, "rb") as f:
                file_content = f.read()

            files_payload = {
                "file": (file.name, file_content, mime_type)
            }

            async with httpx.AsyncClient(timeout=300) as client:
                response = await client.post(UPLOAD_URL, files=files_payload)
                
                if response.status_code == 200:
                    RAG_INITIALIZED = True
                    data = response.json()
                    chunks = data.get("total_chunks", "?")
                    # Removemos 'total_pages' do destaque principal pois txt/py não tem página
                    
                    msg.content = (
                        f"✅ **Processamento Concluído!**\n\n"
                        f"📄 Arquivo: `{file.name}`\n"
                        f"✂️ Segmentos gerados (Chunks): {chunks}\n\n"
                        f"Base de conhecimento atualizada! Pode perguntar."
                    )
                    await msg.update()
                else:
                    RAG_INITIALIZED = False
                    error_detail = response.json().get('detail', 'Erro desconhecido')
                    msg.content = f"❌ Erro ao processar. Código: {response.status_code}. Detalhe: {error_detail}"
                    await msg.update()

        except Exception as e:
            msg.content = f"❌ Ocorreu um erro crítico: {e}"
            await msg.update()
        
@cl.on_message
async def main(message: cl.Message):
    global RAG_INITIALIZED
    
    if not RAG_INITIALIZED:
        await cl.Message(
            content="⚠️ Nenhum documento carregado. Por favor, recarregue a página para enviar um arquivo.",
            author="Sistema"
        ).send()
        return

    msg = cl.Message(content="")
    await msg.send()

    audit_step = cl.Step(name="🔎 Processando & Auditoria", type="process")
    audit_step.input = message.content
    await audit_step.send()

    debug_separator = "###__DEBUG__###"
    is_debug_mode = False
    debug_content = ""

    try:
        async with httpx.AsyncClient(timeout=60) as client:
            async with client.stream("POST", CHAT_URL, json={"content": message.content}) as response:
                if response.status_code != 200:
                    msg.content = f"❌ Erro na IA: {response.status_code}"
                    await msg.update()
                    audit_step.status = "failed"
                    await audit_step.update()
                    return

                async for chunk in response.aiter_text():
                    if is_debug_mode:
                        debug_content += chunk
                        continue

                    if debug_separator in chunk:
                        parts = chunk.split(debug_separator)
                        if parts[0]:
                            await msg.stream_token(parts[0])
                        is_debug_mode = True
                        if len(parts) > 1:
                            debug_content += parts[1]
                    else:
                        await msg.stream_token(chunk)
        
        await msg.update()
        
        if debug_content:
            audit_step.output = debug_content
        else:
            audit_step.output = "Nenhum contexto adicional retornado."
            
        await audit_step.update()

    except Exception as e:
        msg.content = f"❌ Erro de conexão: {e}"
        await msg.update()
        audit_step.status = "failed"
        await audit_step.update()

#--- END OF FILE app.py ---