from multiprocessing import Pool
import orjson
import gzip
from hugedict.mrsw_rocksdb import SecondarySyncedRocksDBDict
from hugedict.rocksdb import RocksDBDict
from hugedict.types import Fn
from loguru import logger


class ParallelFnWrapper:
    def __init__(self, fn: Fn, ignore_error=False):
        self.fn = fn
        self.ignore_error = ignore_error

    def run(self, args):
        idx, r = args
        try:
            r = self.fn(r)
            return idx, r
        except:
            logger.error(f"[ParallelMap] Error while process item {idx}")
            if self.ignore_error:
                return idx, None
            raise


class SecondaryRocksDBCacheFn:
    def __init__(
        self,
        db: SecondarySyncedRocksDBDict[bytes, bytes],
        fn: Fn,
        namespace: str = "",
    ):
        self.db = db
        self.namespace = namespace

        self.fn = fn
        self.fn_name = fn.__name__

    def run(self, *args, **kwargs):
        key = orjson.dumps((self.fn_name, args, kwargs))

        if key not in self.db:
            output = self.fn(*args, **kwargs)
            self.db[key] = output
            return output

        return self.db[key]
