from __future__ import annotations

from rdkit import Chem


def splice_molecule(scaffold: Chem.Mol, fragment: Chem.Mol, position: int) -> Chem.Mol:
    combined = Chem.CombineMols(scaffold, fragment)
    rwmol = Chem.RWMol(combined)

    scaffold_atom_count = scaffold.GetNumAtoms()
    total_atoms = rwmol.GetNumAtoms()

    scaffold_dummy_idx = -1
    for atom_idx in range(scaffold_atom_count):
        atom = rwmol.GetAtomWithIdx(atom_idx)
        if atom.GetAtomicNum() == 0 and atom.GetIsotope() == position:
            scaffold_dummy_idx = atom_idx
            break
    if scaffold_dummy_idx == -1:
        raise ValueError(f"Scaffold dummy atom with isotope {position} not found")

    fragment_dummy_idx = -1
    for atom_idx in range(scaffold_atom_count, total_atoms):
        atom = rwmol.GetAtomWithIdx(atom_idx)
        if atom.GetAtomicNum() == 0:
            fragment_dummy_idx = atom_idx
            break
    if fragment_dummy_idx == -1:
        raise ValueError("Fragment does not contain a dummy atom")

    scaffold_dummy = rwmol.GetAtomWithIdx(scaffold_dummy_idx)
    fragment_dummy = rwmol.GetAtomWithIdx(fragment_dummy_idx)
    if scaffold_dummy.GetDegree() != 1:
        raise ValueError("Scaffold dummy atom must have exactly one neighbor")
    if fragment_dummy.GetDegree() != 1:
        raise ValueError("Fragment dummy atom must have exactly one neighbor")

    scaffold_anchor_idx = scaffold_dummy.GetNeighbors()[0].GetIdx()
    fragment_anchor_idx = fragment_dummy.GetNeighbors()[0].GetIdx()
    fragment_bond = rwmol.GetBondBetweenAtoms(fragment_dummy_idx, fragment_anchor_idx)
    if fragment_bond is None:
        raise ValueError("Fragment dummy atom is not bonded to an anchor atom")

    rwmol.AddBond(scaffold_anchor_idx, fragment_anchor_idx, fragment_bond.GetBondType())

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

    product = rwmol.GetMol()
    Chem.SanitizeMol(product)
    return product
