lazy import ast
lazy from collections.abc import Callable
lazy from functools import wraps
lazy from inspect import getsource, signature, unwrap
lazy from pathlib import Path
lazy from sys import version_info
lazy from textwrap import dedent
lazy from types import FunctionType
lazy from uuid import uuid4

lazy import dill
lazy from blake3 import blake3
lazy from loguru import logger

BASE_FOLDER = Path(
    f"./.cache/{version_info.major}.{version_info.minor}.{version_info.micro}/"
)


def _function_hash[**P, R](fn: Callable[P, R]) -> str:
    fn = unwrap(fn)

    if not isinstance(fn, FunctionType):
        raise ValueError(f"fn {fn} is not an instance of types.FunctionType")

    source = getsource(fn)
    parsed = ast.parse(dedent(source))
    code = ast.dump(parsed)
    sig = signature(fn)
    name: str = fn.__qualname__

    hasher = blake3()
    hasher.update(name.encode("utf-8"))
    hasher.update(sig.format().encode("utf-8"))
    hasher.update(code.encode("utf-8"))

    return hasher.hexdigest()


def _parameter_hash[**P, R](
    fn: Callable[P, R], *args: P.args, **kwargs: P.kwargs
) -> str:
    fn = unwrap(fn)
    sig = signature(fn)

    bound = sig.bind(*args, **kwargs)
    bound.apply_defaults()

    serialized = dill.dumps(bound.arguments)
    hasher = blake3(serialized)

    return hasher.hexdigest()


def persistent[**P, R](func: Callable[P, R]) -> Callable[P, R]:
    fn_hash = _function_hash(func)

    folder = BASE_FOLDER / fn_hash
    folder.mkdir(parents=True, exist_ok=True)

    @wraps(func)
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
        param_hash = _parameter_hash(func, *args, **kwargs)

        path = folder / f"{param_hash}.dill"

        if not path.exists():
            logger.opt(colors=True).debug(
                f"<white>[<cyan>CACHE</cyan>] [<red>MISS</red>] {func.__qualname__}({args}, {kwargs})</white>"
            )
            result = func(*args, **kwargs)
            payload = dill.dumps(result)

            temp_path = path.with_suffix(f".tmp.{uuid4().hex}")
            temp_path.write_bytes(payload)
            temp_path.replace(path)
        else:
            logger.opt(colors=True).debug(
                f"<white>[<cyan>CACHE</cyan>] [<green>HIT</green>] {func.__qualname__}({args}, {kwargs})</white>"
            )
            payload = path.read_bytes()
            result = dill.loads(payload)

        return result

    return wrapper