import os
from io import StringIO, BytesIO
from rdkit import Chem
from rdkit.Chem import Descriptors
from typing import List, Optional
from types import ModuleType
from importlib.resources import files
import random


class BaseMolDataset:
    """
    Helper functions for loading the data.
    I had considered making it a "module factory" but gets quickly unintelligible.
    """

    dataset_package: ModuleType = '__main__'

    @classmethod
    def _package_files(cls):
        """Return a Traversable for the dataset package."""
        return files(cls.dataset_package)

    def __init__(self):
        raise NotImplementedError('Virtual method')

    def raise_not_implemented(self):
        raise NotImplementedError('Method not implemented for this dataset')

    @classmethod
    def file2mol_func(cls, filename: str) -> str:
        return os.path.splitext(filename)[0].strip()

    @classmethod
    def mol2file_func(cls, hit_name: str) -> str:
        return f'{hit_name.strip()}.mol'

    @classmethod
    def get_molblock(cls, hit_name: Optional[str] = None) -> str:
        """
        returns the mol block for the given hit. Chosen randomly if unspecified.
        See ``get_mpro_hit_list()`` for options.
        """
        if hit_name is None:
            hit_name = random.choice(cls.get_hit_list())
        return cls.get_text(cls.mol2file_func(hit_name))
    
    @classmethod
    def get_hit_list(cls) -> List[str]:
        return [cls.file2mol_func(r.name) for r in cls._package_files().iterdir() if '.mol' in r.name]

    @classmethod
    def get_text(cls, filename: str) -> str:
        return cls._package_files().joinpath(filename).read_text(encoding='utf-8')

    @classmethod
    def get_bytes(cls, filename: str) -> bytes:
        return cls._package_files().joinpath(filename).read_bytes()

    @classmethod
    def exists(cls, filename: str) -> bool:
        return cls._package_files().joinpath(filename).is_file()

    @classmethod
    def get_mol(cls, hit_name: Optional[str] = None) -> Chem.Mol:
        """
        returns the Chem.Mol instance for the given hit. Chosen randomly if unspecified.
        See ``get_mpro_hit_list()`` for options.
        """
        mol = Chem.MolFromMolBlock(cls.get_molblock(hit_name))
        mol.SetProp('_Name', hit_name)
        return mol

    @classmethod
    def get_sdf_mols(cls, filename: str) -> List[Chem.Mol]:
        """
        Get the Chem.Mols in the specified sdf file.

        :param filename:
        :return:
        """
        bytes_io = BytesIO(cls.get_bytes(filename))
        with Chem.ForwardSDMolSupplier(bytes_io) as reader:
            return list(reader)

    @classmethod
    def get_filtered_mol(cls, mw_cutoff: float = float('inf')) -> Chem.Mol:
        """
        Return a random Chem.Mol from the MPro dataset,
        with one of the following filtering criteria:

        * ``mw_cutoff``: molWt cutoff
        """
        # For now only MolWt
        choices = cls.get_hit_list()
        random.shuffle(choices)
        for hit_name in choices:
            mol = cls.get_mol(hit_name)
            if mw_cutoff >= Descriptors.MolWt(mol):
                return mol
        else:
            raise ValueError('No hit matches the specified values')

    @classmethod
    def get_n_filtered_mols(cls, amount: int, **cutoffs) -> List[Chem.Mol]:
        """Get ``amount`` of the mols (Chem.Mol) randomly
         that match a cutoff criterion.
        As listed in ``get_filtered_mol``"""
        mols = []
        mol_names = set()
        assert amount > 0, 'A zero amount does nothing.'
        for i in range(len(cls.get_hit_list())):
            mol = cls.get_filtered_mol(**cutoffs)
            mol_name = mol.GetProp('_Name')
            if mol_name in mol_names:
                continue
            mol_names.add(mol_name)
            mols.append(mol)
            if len(mols) == int(amount):
                return mols
        else:
            raise ValueError(f'There are only {len(mols)} molecules ({mol_names}) that match the criteria {cutoffs}')

    @classmethod
    def get_template(cls) -> str:
        return cls.get_text('template.pdb')

# not for the classes to use
def extend_doc(cls):
    """
    This adds usage notes dynamically to the modules.
    The spacing will be off.
    """
    extra = f"""
    The module {cls.dataset_package.__name__} contains the data, but
    not functions to retrieve them without using ``importlib.resources``.
    These are added by the class {cls.__name__} as class methods.
    Not all of the methods are implemented for all datasets.
    """
    cls.dataset_package.__doc__ += extra
