import re
import uuid
from typing import List, Pattern, Union

from nexus_rag.domain.entities.molecule import Molecule
from nexus_rag.domain.entities.atom import Atom

class ConfigurableMarkdownChunker:
    """
    Divide o conteúdo de um Molecule em átomos menores com base em padrões configuráveis, 
    preservando a hierarquia estrutural do documento.
    """
    
    def __init__(self, patterns: List[Union[str, Pattern]], max_chunk_size: int = 3000):
        self.patterns = [re.compile(p) if isinstance(p, str) else p for p in patterns]
        self.max_chunk_size = max_chunk_size

    def chunk(self, molecule: Molecule) -> List[Atom]:
        atoms = []
        active_contexts = ["" for _ in self.patterns]
        text_buffer = []

        def get_context_prefix() -> str:
            parts = [f"[{ctx}]" for ctx in active_contexts if ctx]
            return " ".join(parts) + "\n" if parts else ""

        def save_atom():
            if not text_buffer:
                return
            # Preserva a hierarquia estrutural no prefixo do texto, mas mantém o conteúdo agrupado
            text = "\n".join(text_buffer).strip()
            if text:
                atoms.append(Atom(
                    id=f"atom-{uuid.uuid4()}",
                    molecule_id=molecule.id,
                    raw_text=f"{get_context_prefix()}{text}".strip()
                ))
            text_buffer.clear()

        # Divide o conteúdo em linhas e processa cada linha para detectar mudanças de contexto estrutural
        lines = molecule.content.splitlines()
        
        for line in lines:
            line_str = line.strip()
            if not line_str:
                text_buffer.append("") # preserva quebras de parágrafo
                continue

            matched_level = -1
            for level, pat in enumerate(self.patterns):
                # procura por correspondência total (match) ou parcial (search) para maior flexibilidade
                if pat.match(line_str) or pat.search(line_str):
                    matched_level = level
                    break
            
            if matched_level != -1:
                # Encontrou um novo ponto de corte estrutural, salva o átomo atual antes de mudar o contexto
                save_atom()
                
                # Atualiza o contexto ativo para o nível correspondente, limpando os níveis mais profundos
                clean_context = re.sub(r'^[#\-\s*]+', '', line_str).strip()
                active_contexts[matched_level] = clean_context
                
                # Reseta os contextos mais profundos, pois estamos em um novo bloco estrutural
                for i in range(matched_level + 1, len(self.patterns)):
                    active_contexts[i] = ""
                    
                # Importante: não adicionamos o cabeçalho como um átomo separado, pois ele é parte integrante do conteúdo que segue
                # Assim, o texto do cabeçalho e seu conteúdo relacionado permanecem agrupados, evitando a fragmentação de informações que pertencem juntas,
                # especialmente em casos onde o documento tem uma estrutura hierárquica clara (ex: um artigo seguido de parágrafos relacionados).
                text_buffer.append(line_str)
                continue
                
            text_buffer.append(line_str)
            
            # Caso o buffer de texto ultrapasse o tamanho máximo, salva o átomo atual.
            # Apenas quebra em um novo átomo se estivermos no início de uma nova linha (line_str == "") para evitar cortar sentenças no meio.
            current_size = sum(len(l) for l in text_buffer)
            if current_size > self.max_chunk_size and line_str == "":
                save_atom()
                
            # Adicionalmente, se o tamanho do buffer ultrapassar o limite mesmo sem encontrar uma nova linha, 
            # é melhor salvar o átomo para evitar buffers excessivamente grandes, mesmo que isso possa cortar um parágrafo no meio. 
            # Isso é um trade-off necessário para garantir que os átomos não fiquem muito grandes para processar pela LLM.
            elif current_size > self.max_chunk_size + 2000:
                save_atom()
                
        # Limpa o buffer restante após processar todas as linhas
        save_atom()
        return atoms
