import re
import asyncio
from typing import List, Type
from enum import Enum
from pydantic import BaseModel, Field

from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.prompts import ChatPromptTemplate

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

DEFAULT_SYSTEM_PROMPT = (
    "You are an analytical assistant.\n"
    "We will provide a text fragment associated with the source document ID: {molecule_id}\n"
    "Important formatting rule: Target and Source Molecule IDs MUST be strictly formatted as {id_format_hint} in your output.\n\n"
    "Your mission is to extract the affirmative relationships (bonds) from this text.\n"
    "The MANDATORY relations MUST BE EXACTLY one of these values: {relation_types}.\n"
    "Fill the requested output structure in detail."
)

def create_dynamic_bond_model(relation_types: List[str], id_format_hint: str) -> Type[BaseModel]:
    BondTypeEnum = Enum("BondTypeEnum", {rt: rt for rt in relation_types})
    
    class DynamicBond(BaseModel):
        active_molecule_id: str = Field(..., description=f"Normalized ID of the source/active molecule. Must follow format: {id_format_hint}.")
        passive_molecule_id: str = Field(..., description=f"Normalized ID of the target/passive molecule. Must follow format: {id_format_hint} or 'DESCONHECIDO' if implicit.")
        bond_type: BondTypeEnum = Field(..., description="Relation type. MUST be strictly selected from the allowed enumeration.")
        target_section: str = Field("Seção não especificada", description="Target structural section (e.g., paragraph, article).")
        
    class DynamicBondExtraction(BaseModel):
        bonds: List[DynamicBond] = Field(default_factory=list, description="List of discovered relationships in the text.")
        
    return DynamicBondExtraction

class StructuredBondExtractor:
    """
    Extrator de relações estruturadas (Bonds) a partir de átomos textuais usando um modelo de linguagem.
    O extrator é configurado com tipos de relação permitidos, um padrão de gatilho para seleção de átomos candidatos e uma dica de formato para IDs de moléculas.
    """
    def __init__(
        self, 
        llm: BaseChatModel, 
        relation_types: List[str], 
        candidate_trigger_pattern: str,
        id_format_hint: str = "NUMERO/ANO",
        system_prompt_template: str = None
    ):
        self.llm = llm
        self.relation_types = relation_types
        self.id_format_hint = id_format_hint
        self.trigger_regex = re.compile(candidate_trigger_pattern, re.IGNORECASE)
        
        self.ExtractionModel = create_dynamic_bond_model(relation_types, id_format_hint)
        self.structured_llm = self.llm.with_structured_output(self.ExtractionModel)
        
        system_instruction = system_prompt_template or DEFAULT_SYSTEM_PROMPT
        system_instruction = system_instruction.replace("{relation_types}", ", ".join(self.relation_types))
        system_instruction = system_instruction.replace("{id_format_hint}", self.id_format_hint)
        
        self.prompt = ChatPromptTemplate.from_messages([
            ("system", system_instruction),
            ("human", "{raw_text}")
        ])
        
        self.chain = self.prompt | self.structured_llm

    def find_candidates(self, atoms: List[Atom]) -> List[Atom]:
        return [a for a in atoms if self.trigger_regex.search(a.raw_text)]

    async def _process_atom(self, atom: Atom, base_molecule_id: str) -> List[Bond]:
        try:
            result = await self.chain.ainvoke({
                "molecule_id": base_molecule_id, 
                "raw_text": atom.raw_text
            })
            bonds = []
            for d_bond in result.bonds:
                am = d_bond.active_molecule_id.strip()
                pm = d_bond.passive_molecule_id.strip()
                if not am:
                    continue
                    
                b_val = d_bond.bond_type.value if hasattr(d_bond.bond_type, 'value') else d_bond.bond_type
                bonds.append(Bond(
                    source_atom_id=atom.id,
                    active_molecule_id=am,
                    passive_molecule_id=pm,
                    bond_type=b_val,
                    target_section=d_bond.target_section or "Seção não especificada"
                ))
            return bonds
        except Exception as exc:
            print(f"[ERROR - Atom Extraction {atom.id}]: {exc}")
            return []

    async def extract_bonds_async(self, base_molecule_id: str, atoms: List[Atom]) -> List[Bond]:
        candidates = self.find_candidates(atoms)
        if not candidates:
            return []
            
        tasks = [self._process_atom(a, base_molecule_id) for a in candidates]
        results_matrix = await asyncio.gather(*tasks)
        
        all_bonds = [b for batch in results_matrix for b in batch]
        
        unique_bonds = []
        seen = set()
        for b in all_bonds:
            sig = (b.active_molecule_id, b.passive_molecule_id, b.target_section, b.bond_type)
            if sig not in seen:
                seen.add(sig)
                unique_bonds.append(b)
                
        return unique_bonds

    def extract_bonds(self, base_molecule_id: str, atoms: List[Atom]) -> List[Bond]:
        return asyncio.run(self.extract_bonds_async(base_molecule_id, atoms))
