import pytest
from rdkit import Chem

from macro_lactone_toolkit.splicing.fragment_prep import activate_fragment


def test_activate_smart_ethanol():
    mol = activate_fragment("CCO", strategy="smart")

    assert mol is not None
    assert mol.GetNumAtoms() == 4
    dummy_atom = next(atom for atom in mol.GetAtoms() if atom.GetAtomicNum() == 0)
    assert dummy_atom.GetNeighbors()[0].GetSymbol() == "O"
    assert "*" in Chem.MolToSmiles(mol)


def test_activate_smart_amine():
    mol = activate_fragment("CCN", strategy="smart")

    dummy_atom = next(atom for atom in mol.GetAtoms() if atom.GetAtomicNum() == 0)
    assert dummy_atom.GetNeighbors()[0].GetSymbol() == "N"


def test_activate_random_pentane():
    mol = activate_fragment("CCCCC", strategy="random")

    assert mol.GetNumAtoms() == 6
    dummy_atom = next(atom for atom in mol.GetAtoms() if atom.GetAtomicNum() == 0)
    assert dummy_atom.GetNeighbors()[0].GetSymbol() == "C"


def test_activate_smart_fallback():
    mol = activate_fragment("CCC", strategy="smart")

    dummy_atom = next(atom for atom in mol.GetAtoms() if atom.GetAtomicNum() == 0)
    assert dummy_atom.GetNeighbors()[0].GetSymbol() == "C"
    assert Chem.SanitizeMol(mol) == Chem.SanitizeFlags.SANITIZE_NONE


def test_invalid_smiles():
    with pytest.raises(ValueError):
        activate_fragment("NotASmiles", strategy="smart")