#--- START OF FILE main (1).py ---

import os
import io
import tempfile
from dotenv import load_dotenv
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from langchain_core.prompts import ChatPromptTemplate
from langchain_groq import ChatGroq
from fastapi.responses import StreamingResponse

# RAG Imports
# MUDANÇA: Adicionado TextLoader para arquivos de texto
from langchain_community.document_loaders import PyPDFLoader, TextLoader
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
from langchain_core.output_parsers import StrOutputParser

load_dotenv()

# --- 1. Inicialização ---
app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# --- 2. Configurações de IA ---
HF_EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"

# Modelo Groq
model = ChatGroq(
    model="llama-3.3-70b-versatile",
    temperature=0.3 
)

embeddings = HuggingFaceEmbeddings(
    model_name=HF_EMBEDDING_MODEL,
    model_kwargs={'device': 'cpu'} 
) 

rag_system_prompt = (
    "Você é um assistente experiente e prestativo. "
    "Sua tarefa é fornecer respostas detalhadas e ricas em contexto com base nas informações fornecidas. "
    "Ao usar o contexto abaixo, sintetize os pontos principais e explique como eles se relacionam. "
    "Se o contexto for código (Python/MD), explique o funcionamento ou use como referência."
    "Se a resposta não estiver no contexto, diga honestamente que não sabe, não invente informações."
    "\n\nCONTEXTO DO DOCUMENTO:\n{context}"
)

rag_prompt = ChatPromptTemplate.from_messages(
    [("system", rag_system_prompt), ("human", "{input}")]
)

rag_chain = None
global_retriever = None

# --- 3. Modelo de Dados ---
class ChatRequest(BaseModel):
    content: str

# --- 4. Helpers ---
def format_docs(docs):
    return "\n\n---\n\n".join(
        f"📄 Conteúdo: {doc.page_content}\n(🔖 Fonte: {doc.metadata.get('source', 'Desconhecida')})" 
        for doc in docs
    )

# --- 5. Endpoints ---

@app.post("/upload-document")
async def upload_document(file: UploadFile = File(...)):
    global rag_chain, global_retriever

    # 1. Validação de extensão
    filename = file.filename.lower()
    allowed_extensions = [".pdf", ".txt", ".md", ".py"]
    
    if not any(filename.endswith(ext) for ext in allowed_extensions):
        raise HTTPException(status_code=400, detail=f"Extensão não suportada. Use: {allowed_extensions}")

    try:
        # Salva arquivo temporário com a extensão correta (importante para o Loader)
        file_ext = os.path.splitext(filename)[1]
        with tempfile.NamedTemporaryFile(delete=False, suffix=file_ext) as tmp_file:
            content = await file.read()
            tmp_file.write(content)
            temp_path = tmp_file.name
        
        # 2. Seleção do Loader baseado na extensão
        if filename.endswith(".pdf"):
            loader = PyPDFLoader(temp_path)
        else:
            # Para .txt, .md, .py usamos o TextLoader com UTF-8
            loader = TextLoader(temp_path, encoding="utf-8")

        docs = loader.load()

        # 3. Chunking
        # Se for código (.py), talvez chunks menores sejam melhores, mas manteremos o padrão por enquanto
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1200, chunk_overlap=200)
        splits = text_splitter.split_documents(docs)

        # 4. Vetorização
        vectorstore = FAISS.from_documents(documents=splits, embedding=embeddings)
        global_retriever = vectorstore.as_retriever(search_kwargs={"k": 6}) 
        
        rag_chain = rag_prompt | model | StrOutputParser()

        return {
            "message": "Processamento concluído!",
            "filename": file.filename,
            "total_chunks": len(splits),
            "type": file_ext
        }
    
    except Exception as e:
        print(f"Erro: {e}")
        raise HTTPException(status_code=500, detail=f"Falha ao processar: {e}")
    finally:
        if 'temp_path' in locals() and os.path.exists(temp_path):
            os.remove(temp_path)

@app.post("/chat")
async def chat(request: ChatRequest):
    current_chain = rag_chain
    context_str = ""

    if global_retriever:
        try:
            docs_source = global_retriever.invoke(request.content)
            context_str = format_docs(docs_source)
        except Exception as e:
            print(f"Erro na recuperação: {e}")
            context_str = "Erro ao recuperar contexto."
    else:
        current_chain = (
            ChatPromptTemplate.from_messages([("system", "Você é um assistente útil."), ("human", "{input}")])
            | model 
            | StrOutputParser()
        )

    async def stream_generator():
        try:
            input_data = {"input": request.content}
            if context_str:
                input_data["context"] = context_str

            async for chunk in current_chain.astream(input_data):
                if chunk:
                    yield chunk
            
            if context_str:
                debug_data = f"\n\n###__DEBUG__###\n**Auditoria de Contexto (RAG):**\n\n{context_str}"
                yield debug_data

        except Exception as e:
             print(f"Erro stream: {e}")
             yield f"Erro no serviço de IA: {e}"

    return StreamingResponse(stream_generator(), media_type="text/plain")

#--- END OF FILE main (1).py ---