from __future__ import annotations

from typing import Iterable

from rdkit import Chem

from .._core import collect_side_chain_atoms, ensure_mol, find_macrolactone_candidates, is_intrinsic_lactone_neighbor
from ..fragmenter import MacrolactoneFragmenter


def prepare_macrolactone_scaffold(
    mol_input: str | Chem.Mol,
    positions: Iterable[int],
    ring_size: int | None = None,
) -> tuple[Chem.Mol, dict[int, int]]:
    positions = list(positions)
    mol, _ = ensure_mol(mol_input)
    fragmenter = MacrolactoneFragmenter(ring_size=ring_size)
    numbering = fragmenter.number_molecule(mol)
    candidate = find_macrolactone_candidates(mol, ring_size=numbering.ring_size)[0]
    ring_atom_set = set(numbering.ordered_atoms)

    atoms_to_remove: set[int] = set()
    dummy_specs: list[tuple[int, int, Chem.BondType]] = []

    for position in positions:
        if position not in numbering.position_to_atom:
            raise ValueError(f"Position {position} not found in ring numbering.")
        ring_atom_idx = numbering.position_to_atom[position]
        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)
            atoms_to_remove.update(side_chain_atoms)
            bond = mol.GetBondBetweenAtoms(ring_atom_idx, neighbor_idx)
            if bond is not None:
                dummy_specs.append((ring_atom_idx, position, bond.GetBondType()))

        if not any(spec_position == position for _, spec_position, _ in dummy_specs):
            dummy_specs.append((ring_atom_idx, position, Chem.BondType.SINGLE))

    rwmol = Chem.RWMol(mol)
    for ring_atom_idx, position, bond_type in dummy_specs:
        dummy_atom = Chem.Atom(0)
        dummy_atom.SetIsotope(position)
        dummy_idx = rwmol.AddAtom(dummy_atom)
        rwmol.AddBond(ring_atom_idx, dummy_idx, bond_type)

    for atom_idx in sorted(atoms_to_remove, reverse=True):
        rwmol.RemoveAtom(atom_idx)

    scaffold = rwmol.GetMol()
    Chem.SanitizeMol(scaffold)
    dummy_map = {
        atom.GetIsotope(): atom.GetIdx()
        for atom in scaffold.GetAtoms()
        if atom.GetAtomicNum() == 0 and atom.GetIsotope() in positions
    }
    return scaffold, dummy_map
