from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Callable, Iterable, Sequence, TypeVar

from .types import PolicyCallable, PolicySet, ToolRules, as_frozenset

F = TypeVar("F", bound=Callable[..., Any])


def _get_rules(func: Callable[..., Any]) -> ToolRules:
    rules = getattr(func, "__rgltr__", None)
    if rules is None:
        rules = ToolRules()
        setattr(func, "__rgltr__", rules)
    return rules


def _apply_rules(func: Callable[..., Any], extra: ToolRules) -> Callable[..., Any]:
    existing = _get_rules(func)
    setattr(func, "__rgltr__", existing.merge(extra))
    return func


@dataclass(frozen=True)
class TagDirective:
    activates: frozenset[Any]
    blocked_by: frozenset[Any]

    def __call__(self, func: F) -> F:
        _apply_rules(
            func,
            ToolRules(
                activates=self.activates,
                blocked_by=self.blocked_by,
            ),
        )
        return func


@dataclass(frozen=True)
class BoundaryDirective:
    boundary: Any

    def __call__(self, func: F) -> F:
        _apply_rules(
            func,
            ToolRules(
                boundary=self.boundary,
            ),
        )
        return func


@dataclass(frozen=True)
class PolicyDirective:
    policies: PolicySet

    def __call__(self, func: F) -> F:
        _apply_rules(func, ToolRules(policies=self.policies))
        return func


@dataclass(frozen=True)
class GovernedDirective:
    def __call__(self, func: F) -> F:
        _apply_rules(func, ToolRules(governed=True))
        return func


def tag(
    *,
    activates: Iterable[Any] | None = None,
    blocked_by: Iterable[Any] | None = None,
) -> TagDirective:
    return TagDirective(
        activates=as_frozenset(activates),
        blocked_by=as_frozenset(blocked_by),
    )


def boundary(
    boundary: Any,
) -> BoundaryDirective:
    return BoundaryDirective(boundary=boundary)


def policy(
    *,
    require: Sequence[PolicyCallable] | None = None,
    any_of: Sequence[PolicyCallable] | None = None,
    denied_message: str | None = None,
) -> PolicyDirective:
    policy_set = PolicySet(require=tuple(require or ()), any_of=tuple(any_of or ()), denied_message=denied_message)
    return PolicyDirective(policies=policy_set)


def governed() -> GovernedDirective:
    return GovernedDirective()


def rules_for_tool(func: Callable[..., Any]) -> ToolRules:
    return _get_rules(func)


def apply_directives(base: ToolRules, directives: Sequence[Any]) -> ToolRules:
    rules = base
    for directive in directives:
        if isinstance(directive, TagDirective):
            rules = rules.merge(
                ToolRules(
                    activates=directive.activates,
                    blocked_by=directive.blocked_by,
                )
            )
        elif isinstance(directive, BoundaryDirective):
            rules = rules.merge(
                ToolRules(
                    boundary=directive.boundary,
                )
            )
        elif isinstance(directive, PolicyDirective):
            rules = rules.merge(ToolRules(policies=directive.policies))
        elif isinstance(directive, GovernedDirective):
            rules = rules.merge(ToolRules(governed=True))
        else:
            raise TypeError(f"Unsupported directive: {directive!r}")
    return rules
