from __future__ import annotations

from rdkit import Chem
from rdkit.Chem import Descriptors

from ._core import (
    build_fragment_smiles,
    build_numbering_result,
    collect_side_chain_atoms,
    ensure_mol,
    find_macrolactone_candidates,
    is_intrinsic_lactone_neighbor,
)
from .analyzer import MacroLactoneAnalyzer
from .errors import AmbiguousMacrolactoneError, FragmentationError, MacrolactoneDetectionError
from .models import FragmentationResult, RingNumberingResult, SideChainFragment


class MacrolactoneFragmenter:
    """Number and fragment 12-20 membered macrolactones."""

    def __init__(self, ring_size: int | None = None):
        if ring_size is not None and not 12 <= ring_size <= 20:
            raise ValueError("ring_size must be between 12 and 20")
        self.ring_size = ring_size
        self.analyzer = MacroLactoneAnalyzer()

    def number_molecule(self, mol_input: str | Chem.Mol) -> RingNumberingResult:
        mol, _ = ensure_mol(mol_input)
        candidate = self._select_candidate(mol)
        return build_numbering_result(mol, candidate)

    def fragment_molecule(
        self,
        mol_input: str | Chem.Mol,
        parent_id: str | None = None,
    ) -> FragmentationResult:
        mol, parent_smiles = ensure_mol(mol_input)
        candidate = self._select_candidate(mol)
        numbering = build_numbering_result(mol, candidate)
        parent_id = parent_id or "molecule_0"

        ring_atom_set = set(candidate.ring_atoms)
        fragments: list[SideChainFragment] = []

        for position, ring_atom_idx in numbering.position_to_atom.items():
            ring_atom = mol.GetAtomWithIdx(ring_atom_idx)
            for neighbor in ring_atom.GetNeighbors():
                neighbor_idx = neighbor.GetIdx()
                if neighbor_idx in ring_atom_set:
                    continue
                if is_intrinsic_lactone_neighbor(mol, candidate, ring_atom_idx, neighbor_idx):
                    continue

                side_chain_atoms = collect_side_chain_atoms(mol, neighbor_idx, ring_atom_set)
                if not side_chain_atoms:
                    continue

                try:
                    fragment_smiles_labeled = build_fragment_smiles(
                        mol=mol,
                        side_chain_atoms=side_chain_atoms,
                        side_chain_start_idx=neighbor_idx,
                        ring_atom_idx=ring_atom_idx,
                        dummy_isotope=position,
                    )
                    fragment_smiles_plain = build_fragment_smiles(
                        mol=mol,
                        side_chain_atoms=side_chain_atoms,
                        side_chain_start_idx=neighbor_idx,
                        ring_atom_idx=ring_atom_idx,
                        dummy_isotope=0,
                    )
                except Exception as exc:  # pragma: no cover - sanitized by tests
                    raise FragmentationError(str(exc)) from exc

                plain_fragment = Chem.MolFromSmiles(fragment_smiles_plain)
                if plain_fragment is None:
                    raise FragmentationError("Generated fragment_smiles_plain could not be read by RDKit.")

                fragment_id = f"{parent_id}_frag_{len(fragments)}"
                fragments.append(
                    SideChainFragment(
                        parent_id=parent_id,
                        fragment_id=fragment_id,
                        cleavage_position=position,
                        attachment_atom_idx=ring_atom_idx,
                        fragment_smiles_labeled=fragment_smiles_labeled,
                        fragment_smiles_plain=fragment_smiles_plain,
                        atom_count=sum(1 for atom in plain_fragment.GetAtoms() if atom.GetAtomicNum() != 0),
                        molecular_weight=round(Descriptors.MolWt(plain_fragment), 4),
                    )
                )

        return FragmentationResult(
            parent_id=parent_id,
            parent_smiles=parent_smiles,
            ring_size=numbering.ring_size,
            numbering=numbering,
            fragments=fragments,
        )

    def _select_candidate(self, mol: Chem.Mol):
        candidates = find_macrolactone_candidates(mol, ring_size=self.ring_size)
        if not candidates:
            requested = f"{self.ring_size}-membered " if self.ring_size is not None else ""
            raise MacrolactoneDetectionError(f"No valid {requested}macrolactone was detected.")

        valid_ring_sizes = sorted({candidate.ring_size for candidate in candidates})
        if len(candidates) > 1 or len(valid_ring_sizes) > 1:
            raise AmbiguousMacrolactoneError(
                "Multiple valid macrolactone candidates were detected. Pass ring_size explicitly."
            )
        return candidates[0]
