import logging as lg
from concurrent.futures import ProcessPoolExecutor

from rdkit.Chem import (
    Kekulize,
    Mol,
    MolFromInchi,
    MolFromSmiles,
    MolToInchi,
    MolToInchiAndAuxInfo,
    MolToInchiKey,
    MolToSmiles,
    RemoveStereochemistry,
)
from rdkit.Chem.inchi import InchiToInchiKey

from fairfetched.utils._track import track

from ._optional import _papyrus_standardize

# ------ reading ------------


def _safe_smiles_to_mol(smiles: str) -> Mol | None:
    try:
        return MolFromSmiles(smiles)
    except Exception as e:
        lg.error(f"Failed on `MolFromSmiles` conversion: {smiles} with exception {e}")
        return None


def _safe_inchi_to_mol(inchi: str) -> Mol | None:
    try:
        return MolFromInchi(inchi)
    except Exception as e:
        lg.error(f"Failed on `MolFromInchi` conversion: {inchi} with exception {e}")
        return None


# ------ cleaning -----------


def _safe_remove_stereochemistry(mol: Mol | None) -> Mol | None:
    try:
        RemoveStereochemistry(mol)
        return mol
    except Exception as e:
        lg.error(f"Failed in RemoveStereochemistry(mol) with error {e}")
        return None


def _safe_validate_inchi(mol) -> Mol | None:
    try:
        inchi = MolToInchi(mol)
        mol = MolFromInchi(inchi)
        if not isinstance(mol, Mol):
            lg.error(f"Failed on InChI conversion: {inchi}")
            return None
        return mol
    except Exception:
        return None


def _safe_standardise_papyrus(mol: Mol | None, **kwargs) -> Mol | None:
    if not mol:
        return None
    try:
        mol, message = _papyrus_standardize(
            mol, return_type=True, raise_error=False, **kwargs
        )
        if not mol:
            lg.error(f"Papyrus-type standardisation did not return a mol: {message}")
            return None
        return mol
    except Exception as e:
        lg.error(f"failed papyrus standardisation with exception: {e}")
        return None


# ------ writing -----------


def mol_to_strings(mol: Mol | None) -> tuple[str, str, str, str] | None:
    """returns a tuple of kekulised smiles, inchi, inchiauxinfo, and inchikeys"""
    if not Mol:
        return None
    return MolToSmiles(mol), *MolToInchiAndAuxInfo(mol), MolToInchiKey(mol)


def _aeffect_standardisation_from_smiles(smiles: str) -> Mol | None:
    """Returns cleaned molecules from SMILES string or None upon failure or invalid molecules"""
    assert isinstance(smiles, str)
    if not smiles:
        lg.warning("empty smiles")
        return None
    mol = MolFromSmiles(smiles)
    if not isinstance(mol, Mol):
        lg.debug("Failed on initial `MolFromSmiles` conversion")
        return None
    try:
        RemoveStereochemistry(mol)
    except Exception as e:
        lg.error(f"Failed in RemoveStereochemistry(mol) with error {e}")
        return None
    # @TODO: check if stereochemistry removal should be before or after standardisation
    if not mol:
        lg.error(f"{smiles} could not be converted to a mol")
        return None
    try:
        mol, message = _papyrus_standardize(
            mol,
            filter_mixtures=True,  # @TODO: check if this is desired
            filter_inorganic=False,
            filter_non_small_molecule=False,
            return_type=True,
            raise_error=False,
        )
        if not mol:
            lg.warning(
                f"{smiles} could not be cleaned during Papyrus-type standardisation: {message}"
            )
            return None
        mol = MolFromInchi(MolToInchi(mol))
        if not mol:
            lg.warning(f"{smiles} did not give valid inchi: {message}")
        return mol
    except Exception as e:
        lg.error(
            f"{smiles} could not be cleaned in papyrus-style standardsation with exception: {e}"
        )
        return None
    return None


def inchis_to_mols(
    inchis: Iterable[str], cache: Path | None = None
) -> list[Mol | None]:
    """Convert InChIs to RDKit Mol objects"""
    return [
        _safe_inchi_to_mol(i) for i in track(inchis, desc="Converting InChI to mol")
    ]


def mols_to_kekulised_smiles(mols: Iterable[Mol]) -> list[str]:
    """Returns kekulised SMILES strings suitable for re-loading the molecule later.
    Removes stereochemistry.
    """
    return [
        _mol_to_kekulised_smiles(m)
        for m in track(mols, desc="converting mols to kekulised smiles")
    ]


def _mol_to_inchi_and_auxinfo(mol: Mol) -> tuple[str, str]:
    return MolToInchiAndAuxInfo(mol)


def _mol_to_inchi(mol: Mol) -> str:
    return MolToInchiAndAuxInfo(mol)[0]  # to maintain consistency


def mols_to_inchis(mols: Iterable[Mol]) -> list[str]:
    return [_mol_to_inchi(m) for m in track(mols, desc="converting mols to inchis")]


def mols_to_inchis_and_auxinfo(mols: Iterable[Mol]) -> list[tuple[str, str]]:
    return [
        _mol_to_inchi_and_auxinfo(m)
        for m in track(mols, desc="converting mols to inchis with auxinfo")
    ]


def inchis_to_inchikeys(inchis: Iterable[str]) -> list[str]:
    return [
        InchiToInchiKey(i) for i in track(inchis, desc="converting inchis to inchikeys")
    ]


def mols_to_inchikeys(mols: Iterable[Mol]) -> list[str]:
    return [
        InchiToInchiKey(MolToInchiAndAuxInfo(m)[0])
        for m in track(mols, desc="converting mols to inchikeys")
    ]


def mols_to_strings_zip(mols: Iterable[Mol]) -> zip[tuple[str, str, str, str]]:
    return zip(
        map(_mol_to_kekulised_smiles, mols),
        *(zip(*map(_mol_to_inchi_and_auxinfo, mols))),
        map(MolToInchiKey, mols),
    )


def mols_to_strings(mols: Iterable[Mol]) -> list[tuple[str, str, str, str]]:
    return list(mols_to_strings(mols))


def make_rdkit_fingerprint_fn(
    fp_type: str = "morgan", radius: int = 2, n_bits: int = 2048
):
    """Returns a rep_fn that produces a bit vector as a list[int]."""

    def fp_fn(mol) -> list[int] | None:
        from rdkit.Chem import rdMolDescriptors

        if fp_type == "morgan":
            bv = rdMolDescriptors.GetMorganFingerprintAsBitVect(
                mol, radius, nBits=n_bits
            )
        elif fp_type == "rdkit":
            from rdkit.Chem import RDKFingerprint

            bv = RDKFingerprint(mol, fpSize=n_bits)
        else:
            raise ValueError(f"Unknown fp_type: {fp_type}")
        return list(bv.ToList())

    return fp_fn
