modules/rag_normas.py

"""
RAG Normativo - Modulo de Consulta a NT-01/2025 CBMGO
Usa sentence-transformers + FAISS para busca semantica
"""
import json
import os
import numpy as np

try:
    from sentence_transformers import SentenceTransformer
    import faiss
    HAS_FAISS = True
except ImportError:
    HAS_FAISS = False

from datasets import load_dataset


class RAGNormas:
    """
    Modulo RAG para consulta semantica a NT-01/2025 CBMGO.
    
    Uso:
        rag = RAGNormas()
        rag.build_from_jsonl("data/chunks.jsonl")
        resultado = rag.responder("Quais extintores sao obrigatorios para comercio?")
    """
    
    def __init__(self, model_name: str = "sentence-transformers/all-mpnet-base-v2"):
        self.model_name = model_name
        self.model = None
        self.index = None
        self.chunks = []
        self._initialized = False
    
    def _load_model(self):
        if not HAS_FAISS:
            raise ImportError("Instale: pip install sentence-transformers faiss-cpu")
        if self.model is None:
            print(f"Carregando modelo: {self.model_name}")
            self.model = SentenceTransformer(self.model_name)
    
    def build_from_jsonl(self, chunks_path: str):
        """Constroi indice FAISS a partir de arquivo .jsonl de chunks."""
        self._load_model()
        
        if not os.path.exists(chunks_path):
            raise FileNotFoundError(f"Arquivo nao encontrado: {chunks_path}")
        
        self.chunks = []
        with open(chunks_path, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line:
                    self.chunks.append(json.loads(line))
        
        print(f"Carregando {len(self.chunks)} chunks...")
        textos = [c["texto"] for c in self.chunks]
        embeddings = self.model.encode(textos, show_progress_bar=True,
                                        batch_size=32, normalize_embeddings=True)
        
        dim = embeddings.shape[1]
        self.index = faiss.IndexFlatIP(dim)  # Inner Product = cosine para vetores normalizados
        self.index.add(embeddings.astype(np.float32))
        self._initialized = True
        print(f"Indice FAISS construido: {self.index.ntotal} vetores, dim={dim}")
    
    def load_from_files(self, index_path: str, chunks_path: str):
        """Carrega indice FAISS pre-construido."""
        self._load_model()
        self.index = faiss.read_index(index_path)
        self.chunks = []
        with open(chunks_path, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line:
                    self.chunks.append(json.loads(line))
        self._initialized = True
        print(f"Indice carregado: {self.index.ntotal} vetores, {len(self.chunks)} chunks")
    
    def save_index(self, index_path: str):
        """Salva indice FAISS em disco."""
        if self.index is None:
            raise RuntimeError("Indice nao construido. Execute build_from_jsonl() primeiro.")
        faiss.write_index(self.index, index_path)
        print(f"Indice salvo em: {index_path}")
    
    def buscar(self, query: str, top_k: int = 5) -> list:
        """Busca os top_k chunks mais relevantes para a query."""
        if not self._initialized:
            raise RuntimeError("RAG nao inicializado. Execute build_from_jsonl() ou load_from_files().")
        
        emb = self.model.encode([query], normalize_embeddings=True)
        scores, indices = self.index.search(emb.astype(np.float32), top_k)
        
        resultados = []
        for score, idx in zip(scores[0], indices[0]):
            if idx < len(self.chunks):
                chunk = self.chunks[idx].copy()
                chunk["score"] = float(score)
                resultados.append(chunk)
        return resultados
    
    def responder(self, query: str, top_k: int = 5) -> str:
        """Retorna contexto formatado com referencias normativas."""
        resultados = self.buscar(query, top_k)
        if not resultados:
            return "Nenhum resultado encontrado para a consulta."
        
        linhas = [f"Consulta: {query}", "=" * 60]
        for i, r in enumerate(resultados, 1):
            ref = r.get("ref", "NT-01/2025")
            artigo = r.get("artigo", "")
            texto = r.get("texto", "")
            score = r.get("score", 0)
            linhas.append(f"[{i}] {ref}{' - ' + artigo if artigo else ''} (relevancia: {score:.3f})")
            linhas.append(texto[:500] + ("..." if len(texto) > 500 else ""))
            linhas.append("-" * 40)
        
        return "\n".join(linhas)
    
    def responder_fallback(self, query: str) -> str:
        """Fallback sem FAISS: busca por palavras-chave nos chunks carregados."""
        if not self.chunks:
            return self._responder_estatico(query)
        
        query_lower = query.lower()
        keywords = query_lower.split()
        
        scored = []
        for chunk in self.chunks:
            texto = chunk.get("texto", "").lower()
            score = sum(1 for kw in keywords if kw in texto)
            if score > 0:
                scored.append((score, chunk))
        
        scored.sort(key=lambda x: x[0], reverse=True)
        top = scored[:3]
        
        if not top:
            return self._responder_estatico(query)
        
        linhas = [f"Resultado para: {query}", "=" * 50]
        for score, chunk in top:
            ref = chunk.get("ref", "NT-01/2025")
            artigo = chunk.get("artigo", "")
            texto = chunk.get("texto", "")
            linhas.append(f"[{ref}{' - ' + artigo if artigo else ''}]")
            linhas.append(texto[:400])
            linhas.append("-" * 30)
        return "\n".join(linhas)
    
    def _responder_estatico(self, query: str) -> str:
        """Respostas estaticas para perguntas comuns sobre NT-01/2025."""
        q = query.lower()
        
        respostas = {
            "extintor": """NT-01/2025 CBMGO - Extintores Portateis (Anexo B)
Todas as edificacoes devem possuir extintores portateis.
- Grupo A (residencial): 1 extintor a cada 500m2, cap. 2-A:10-B:C
- Grupo D (comercial): 1 extintor a cada 400m2, cap. 2-A:20-B:C
- Grupo F (industrial): 1 extintor a cada 300m2, cap. 4-A:20-B:C
Distancia maxima de qualquer ponto ao extintor: 15 metros.
Ref: NT-01/2025 Tabela B-1 | ABNT NBR 12693""",
            
            "hidrante": """NT-01/2025 CBMGO - Sistema de Hidrantes (Art. 18-22)
Obrigatorio para edificacoes com area > 750m2 ou altura > 12m.
- Altura < 12m: Mangueira / Coluna Seca, vazao 300 L/min
- Altura 12-30m: Hidrante com Reservatorio, vazao 450 L/min
- Altura > 30m: Chuveiros Automaticos + Hidrante, vazao 600 L/min
Reserva tecnica de incendio minima: 18m3 (coluna seca) a 36m3 (alta pressao)
Ref: NT-01/2025 Art. 18 | ABNT NBR 13714""",
            
            "spda": """NT-01/2025 CBMGO - SPDA (Art. 25)
Sistema de Protecao contra Descargas Atmosfericas.
Obrigatorio para edificacoes com area total > 1000m2.
Tambem obrigatorio para: hospitais, escolas, locais de reuniao,
depositos de inflamaveis e edificacoes com altura > 20m.
Projeto conforme ABNT NBR 5419.
Ref: NT-01/2025 Art. 25""",
            
            "iluminacao": """NT-01/2025 CBMGO - Iluminacao de Emergencia (Art. 20)
Obrigatoria para edificacoes com altura > 12m.
Tambem obrigatoria em: saidas de emergencia, escadas, corredores,
locais de reuniao com lotacao > 50 pessoas.
Autonomia minima: 1 hora (ABNT NBR 10898).
Nivel minimo de iluminamento: 3 lux nas rotas de fuga.
Ref: NT-01/2025 Art. 20 | ABNT NBR 10898""",
            
            "sinalizacao": """NT-01/2025 CBMGO - Sinalizacao de Emergencia (Art. 19)
Obrigatoria em todas as edificacoes.
Itens obrigatorios:
- Placas de saida de emergencia (verde)
- Indicacao de extintores e hidrantes
- Planta de localizacao em locais de reuniao
- Instrucoes de emergencia
Ref: NT-01/2025 Art. 19 | ABNT NBR 13434""",
            
            "ocupacao": """NT-01/2025 CBMGO - Classificacao de Ocupacoes (Tabela 1)
Grupos de ocupacao:
A - Residencial (unifamiliar/multifamiliar)
B - Servicos de hospedagem (hoteis, pousadas)
C - Reuniao de publico (teatros, cinemas, igrejas)
D - Servico profissional, pessoal e tecnico (comercio, escritorios)
E - Educacional e cultura fisica
F - Locais de saude e institucional
G - Servico automotivo, transporte e pecas
H - Depositos, armazenamento geral
I - Industrial
Ref: NT-01/2025 Tabela 1""",
        }
        
        for chave, resp in respostas.items():
            if chave in q:
                return resp
        
        return (
            "Consulta: " + query + "\n"
            "\nBase de dados normativa NT-01/2025 CBMGO disponivel.\n"
            "Para consulta completa, construa o indice FAISS com build_from_jsonl().\n"
            "\nTopicos disponveis: extintores, hidrantes, SPDA, iluminacao, sinalizacao, ocupacao."
        )


# Instancia global (lazy loading)
_rag_instance = None

def get_rag() -> RAGNormas:
    global _rag_instance
    if _rag_instance is None:
        _rag_instance = RAGNormas()
        # Tentar carregar chunks se existirem
        if os.path.exists("data/chunks.jsonl"):
            try:
                if os.path.exists("data/embeds.faiss") and HAS_FAISS:
                    _rag_instance.load_from_files("data/embeds.faiss", "data/chunks.jsonl")
                else:
                    # Carregar apenas chunks para fallback por palavras-chave
                    with open("data/chunks.jsonl", "r", encoding="utf-8") as f:
                        _rag_instance.chunks = [json.loads(l) for l in f if l.strip()]
            except Exception as e:
                print(f"Aviso RAG: {e}")
    return _rag_instance


if __name__ == "__main__":
    rag = RAGNormas()
    print(rag._responder_estatico("extintores portateis"))
    print("\n---\n")
    print(rag._responder_estatico("sistema de hidrantes"))