import pickle

import pandas as pd


def save_to_lmdb(env, key, dataset):
    """
    Saves the given dataset to an LMDB environment with the given key.
    :param env: LMDB environment
    :param key: key to save the dataset to
    :param dataset: tuple of pandas dataframes
    """

    with env.begin(write=True) as txn:
        txn.put(key.encode('ascii'), pickle.dumps(dataset))


def load_from_lmdb(env, key):
    """
    Loads a dataset from an LMDB environment with the given key.
    :param env: LMDB environment
    :param key: key of the dataset to load
    :return: key and tuple of pandas dataframes (input, context, output)
    """
    with env.begin(write=False) as txn:
        try:
            data = pickle.loads(txn.get(key.encode('ascii')))
            return data
        except TypeError:
            raise KeyError(key)


def delete_from_lmdb(env, key):
    """
    Deletes a dataset from an LMDB environment with the given key.
    :param env: LMDB environment
    :param key: key of the dataset to delete
    """
    with env.begin(write=True) as txn:
        txn.delete(key.encode('ascii'))


def clear_lmdb(env):
    """
    Clears all datasets from an LMDB environment.
    :param env: LMDB environment
    """
    with env.begin(write=True) as txn:
        cursor = txn.cursor()
        for key, value in cursor:
            txn.delete(key)


def lmdb_dataset_generator(env):
    """
    Generator function to yield datasets from an LMDB environment.
    :param env: LMDB environment
    :return: generator
    """
    with env.begin(write=False) as txn:
        cursor = txn.cursor()
        for key, value in cursor:
            data = pickle.loads(value)
            yield data


def lmdb_contains(env, substring) -> bool:
    """
    Checks if the given substring is contained in any of the keys of the LMDB environment.
    :param env: LMDB environment
    :param substring: substring to search for
    :return: boolean
    """
    with env.begin(write=False) as txn:
        cursor = txn.cursor()
        for key, value in cursor:
            if substring in key.decode('ascii'):
                return True
    return False


def lmdb_substring_key_search(env, substring):
    """
    Searches for keys in the LMDB environment that contain the given substring.
    :param env: LMDB environment
    :param substring: substring to search for
    :return: list of keys
    """
    keys = []
    with env.begin(write=False) as txn:
        cursor = txn.cursor()
        for key, value in cursor:
            if substring in key.decode('ascii'):
                keys.append(key)
    return keys


def get_lmdb_keys(env, limit: int = None):
    """
    Get all keys in the LMDB environment.
    :param env: LMDB environment
    :param limit: maximum number of keys to return
    :return: list of keys
    """
    with env.begin(write=False) as txn:
        with txn.cursor() as cursor:
            keys = [key.decode("ascii") for key in cursor.iternext(keys=True, values=False)]

    return keys


def get_lmdb_keyspace_size(env):
    """
    Get the number of keys in the LMDB environment.
    :param env: LMDB environment
    :return: number of keys
    """
    with env.begin(write=False) as txn:
        return txn.stat()['entries']