# SPDX-FileCopyrightText: 2023-present Massimiliano Pippi <mpippi@gmail.com>
#
# SPDX-License-Identifier: MIT
import pickle
from typing import Optional, Protocol, runtime_checkable


@runtime_checkable
class RenderCache(Protocol):
    """
    Interface for rendering cache.

    Prompts with the same template and same context are retrieved from the cache directly,
    instead of going through the rendering process.
    """

    def get(self, context: dict) -> Optional[str]: ...

    def set(self, context: dict, prompt: str) -> None: ...

    def clear(self) -> None: ...


class DefaultCache:
    """
    In-memory, default rendering cache.
    """

    def __init__(self) -> None:
        self._cache: dict[bytes, str] = {}

    def get(self, context: dict) -> Optional[str]:
        return self._cache.get(pickle.dumps(context, pickle.HIGHEST_PROTOCOL))

    def set(self, context: dict, prompt: str) -> None:
        self._cache[pickle.dumps(context, pickle.HIGHEST_PROTOCOL)] = prompt

    def clear(self) -> None:
        self._cache = {}
