from __future__ import annotations

from contextvars import ContextVar, Token
from threading import RLock
from typing import TYPE_CHECKING, ClassVar, Protocol

if TYPE_CHECKING:
    from collections.abc import Callable, Sequence

from kele.control.callback import HookMixin


class SupportsDescription(Protocol):
    """Protocol for objects that expose a description field."""

    description: str


class DescriptionJoiner:
    """Utilities for merging description fragments."""

    @staticmethod
    def merge(parts: Sequence[str], *, separator: str = " | ") -> str:
        """Combine description fragments, preserving order and removing duplicates/empties."""
        seen: set[str] = set()
        merged: list[str] = []
        for part in parts:
            # Flatten already-joined fragments so nested combinations do not duplicate tokens.
            for fragment in part.split(separator):
                stripped = fragment.strip()
                if not stripped or stripped in seen:
                    continue
                seen.add(stripped)
                merged.append(stripped)
        return separator.join(merged)


class DescriptionRegistryMixin(HookMixin):
    """Optional registry for looking up description-bearing objects by key."""

    _kv_registry_enabled: ClassVar[bool] = False
    _kv_registry: ClassVar[dict[str, dict[str, dict[str, object]]]] = {}
    _kv_registry_namespace: ClassVar[str] = "default"
    _kv_registry_default_namespace: ClassVar[str] = "default"
    _kv_registry_active_namespace: ClassVar[ContextVar[str | None]] = ContextVar(
        "description_registry_active_namespace",
        default=None,
    )
    _kv_registry_lock: ClassVar[RLock] = RLock()

    @classmethod
    def _registry_bucket(cls, namespace: str) -> dict[str, object]:
        namespace_bucket = cls._kv_registry.setdefault(namespace, {})
        return namespace_bucket.setdefault(cls.__name__, {})

    @classmethod
    def _resolve_namespace(cls, namespace: str | None) -> str:
        if namespace is not None:
            return namespace
        active_namespace = cls._kv_registry_active_namespace.get()
        if active_namespace is not None:
            return active_namespace
        return cls._kv_registry_namespace

    @classmethod
    def activate_kv_registry_namespace(cls, namespace: str) -> Token[str | None]:
        """Activate namespace for the current context."""
        return cls._kv_registry_active_namespace.set(namespace)

    @classmethod
    def reset_kv_registry_namespace(cls, token: Token[str | None]) -> None:
        """Restore previous namespace context."""
        cls._kv_registry_active_namespace.reset(token)

    @classmethod
    def _enable_kv_registry_namespace(cls, *, enabled: bool = True, namespace: str) -> None:
        """Enable or disable key->value recording for an explicit namespace (engine-internal use)."""
        if not namespace:
            raise ValueError("namespace must be a non-empty string.")
        with cls._kv_registry_lock:
            cls._kv_registry_enabled = enabled
            cls._kv_registry_namespace = namespace
            cls._registry_bucket(namespace)

    @classmethod
    def enable_kv_registry(cls, *, enabled: bool = True, namespace: str | None = None) -> None:
        """Enable or disable key->value recording for the user-facing default namespace only."""
        default_namespace = cls._kv_registry_default_namespace
        if namespace not in {None, default_namespace}:
            raise ValueError(
                f"Only namespace={default_namespace!r} is allowed for enable_kv_registry(); "
                "custom namespaces are reserved for engine-internal scoping."
            )
        cls._enable_kv_registry_namespace(enabled=enabled, namespace=default_namespace)

    @classmethod
    def record_kv(cls, key: object, value: object) -> None:
        """Record a key->value pair if the registry is enabled."""
        if not cls._kv_registry_enabled:
            return
        target_namespace = cls._resolve_namespace(None)
        with cls._kv_registry_lock:
            bucket = cls._registry_bucket(target_namespace)
            bucket[str(key)] = value

    @classmethod
    def get_kv(cls, key: object, default: object | None = None, *, namespace: str | None = None) -> object | None:
        """Retrieve a value by key if the registry is enabled."""
        if not cls._kv_registry_enabled:
            return default
        target_namespace = cls._resolve_namespace(namespace)
        with cls._kv_registry_lock:
            namespace_bucket = cls._kv_registry.get(target_namespace, {})
            bucket = namespace_bucket.get(cls.__name__, {})
            return bucket.get(str(key), default)

    @classmethod
    def get_kv_items(cls, *, namespace: str | None = None) -> list[tuple[str, object]]:
        """Return a snapshot of key/value entries for one namespace."""
        if not cls._kv_registry_enabled:
            return []
        target_namespace = cls._resolve_namespace(namespace)
        with cls._kv_registry_lock:
            namespace_bucket = cls._kv_registry.get(target_namespace, {})
            bucket = namespace_bucket.get(cls.__name__, {})
            return list(bucket.items())

    @classmethod
    def clear_kv_registry(cls, *, namespace: str | None = None) -> None:
        """Clear registry data for one namespace."""
        if namespace is None:
            raise ValueError("namespace must be provided. Use clear_all_kv_registry() for global clear.")
        with cls._kv_registry_lock:
            namespace_bucket = cls._kv_registry.get(namespace)
            if namespace_bucket is None:
                return
            namespace_bucket.pop(cls.__name__, None)
            if not namespace_bucket:
                cls._kv_registry.pop(namespace, None)

    @classmethod
    def clear_all_kv_registry(cls) -> None:
        """Clear registry data for all namespaces explicitly."""
        with cls._kv_registry_lock:
            cls._kv_registry.clear()


class DescriptionHandlerMixin(HookMixin):
    """Mixin for optional description handlers."""

    description_handler: ClassVar[Callable[[SupportsDescription], str] | None] = None

    @classmethod
    def set_description_handler(cls, handler: Callable[[SupportsDescription], str] | None) -> None:
        """Set a custom description handler (None to disable)."""
        cls.description_handler = handler
