import abc
import inspect
from typing import Any, Iterable, Mapping

from typing_extensions import Self

from .instrumentation import CACHE_SIZE


class Annotation(abc.ABC):
    _cache: dict[tuple[type[Self], inspect.Signature], Mapping[str, Self]] = {}

    @classmethod
    def annotated_parameters(cls, signature: inspect.Signature) -> Mapping[str, Self]:
        key = (cls, signature)
        if key in cls._cache:
            CACHE_SIZE.set(len(cls._cache), {"cache": "annotation"})
            return cls._cache[key]

        annotated: dict[str, Self] = {}

        for param_name, param in signature.parameters.items():
            if param.annotation == inspect.Parameter.empty:
                continue

            try:
                metadata: Iterable[Any] = param.annotation.__metadata__
            except AttributeError:
                continue

            for arg_type in metadata:
                if isinstance(arg_type, cls):
                    annotated[param_name] = arg_type
                elif isinstance(arg_type, type) and issubclass(arg_type, cls):
                    annotated[param_name] = arg_type()

        cls._cache[key] = annotated
        CACHE_SIZE.set(len(cls._cache), {"cache": "annotation"})
        return annotated


class Logged(Annotation):
    """Instructs docket to include arguments to this parameter in the log.

    If `length_only` is `True`, only the length of the argument will be included in
    the log.

    Example:

    ```python
    @task
    def setup_new_customer(
        customer_id: Annotated[int, Logged],
        addresses: Annotated[list[Address], Logged(length_only=True)],
        password: str,
    ) -> None:
        ...
    ```

    In the logs, you's see the task referenced as:

    ```
    setup_new_customer(customer_id=123, addresses[len 2], password=...)
    ```
    """

    length_only: bool = False

    def __init__(self, length_only: bool = False) -> None:
        self.length_only = length_only

    def format(self, argument: Any) -> str:
        if self.length_only:
            if isinstance(argument, (dict, set)):
                return f"{{len {len(argument)}}}"
            elif isinstance(argument, tuple):
                return f"(len {len(argument)})"
            elif hasattr(argument, "__len__"):
                return f"[len {len(argument)}]"

        return repr(argument)