import json
from typing import Any
import zlib

import zstandard as zstd


COMPRESSION_LIBS = {"zlib": zlib, "zstd": zstd}


def compress(lib: str or None, data: Any) -> Any:
    """Compress string data to decrease mongodb storage size.

    Args:
        lib: Compression lib name.
        data: Data to compress.

    Returns:
        Compressed data if compression is enabled else original value.

    Raises:
        NotImplementedError: Unknown compression lib name.

    """
    if lib:
        if lib in COMPRESSION_LIBS:
            return COMPRESSION_LIBS[lib].compress(json.dumps(data).encode())
        else:
            raise NotImplementedError()
    return data


def decompress(lib: str or None, data: Any) -> Any:
    """Decompress data from mongo using specific lib.

    Args:
        lib: Compression lib name.
        data: Data to decompress.

    Returns:
        Decompressed data if it was compressed before else data.

    Raises:
        NotImplementedError: Unknown compression lib name.

    """
    if lib:
        if lib in COMPRESSION_LIBS:
            return json.loads(COMPRESSION_LIBS[lib].decompress(data).decode())
        else:
            raise NotImplementedError()
    return data