from typing import List, Union, Pattern, Tuple
from langchain_core.language_models.chat_models import BaseChatModel

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

from nexus_rag.ingestion.parsers.pdf_parser import DoclingParser
from nexus_rag.ingestion.chunkers.markdown_chunker import ConfigurableMarkdownChunker
from nexus_rag.ingestion.extractors.bond_extractor import StructuredBondExtractor

class NexusPipelineRunner:
    """
    Orquestrador final da injeção de documentos.
    Encapsula o parser (Docling), chunker e extractor (LLM) no fluxo completo.
    """
    def __init__(
        self, 
        llm: BaseChatModel, 
        relation_types: List[str], 
        chunking_rules: List[Union[str, Pattern]],
        candidate_trigger_pattern: str,
        id_format_hint: str = "NUMERO/ANO",
        max_chunk_size: int = 3000,
        system_prompt_template: str = None
    ):
        self.parser = DoclingParser()
        self.chunker = ConfigurableMarkdownChunker(
            patterns=chunking_rules, 
            max_chunk_size=max_chunk_size
        )
        self.extractor = StructuredBondExtractor(
            llm=llm, 
            relation_types=relation_types, 
            candidate_trigger_pattern=candidate_trigger_pattern,
            id_format_hint=id_format_hint,
            system_prompt_template=system_prompt_template
        )

    def process_from_pdf(self, pdf_path: str, molecule_id: str, title: str, aliases: List[str] = None) -> Tuple[Molecule, List[Atom], List[Bond]]:
        """
        Ingere um arquivo PDF usando Docling, empacota como Molecule, 
        gera os atoms e extrai os relacionamentos (Bonds) síncronamente.
        """
        markdown_content = self.parser.parse_pdf(pdf_path)
        
        mol = Molecule(
            id=molecule_id,
            title=title,
            content=markdown_content,
            aliases=aliases or []
        )
        
        atoms = self.chunker.chunk(mol)
        bonds = self.extractor.extract_bonds(base_molecule_id=mol.id, atoms=atoms)
        return mol, atoms, bonds

    async def process_from_pdf_async(self, pdf_path: str, molecule_id: str, title: str, aliases: List[str] = None) -> Tuple[Molecule, List[Atom], List[Bond]]:
        """Versão assíncrona do processamento desde o PDF suportando concorrência no LLM."""
        markdown_content = self.parser.parse_pdf(pdf_path)
        
        mol = Molecule(
            id=molecule_id,
            title=title,
            content=markdown_content,
            aliases=aliases or []
        )
        
        atoms = self.chunker.chunk(mol)
        bonds = await self.extractor.extract_bonds_async(base_molecule_id=mol.id, atoms=atoms)
        return mol, atoms, bonds
