from __future__ import annotations

from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Awaitable, Callable, Iterable, Literal, Sequence


PolicyOutcome = Literal["allow", "deny", "approval"]


@dataclass(frozen=True)
class PolicyDecision:
    """Outcome of a policy check."""

    outcome: PolicyOutcome
    reason: str

    @classmethod
    def allow(cls, reason: str = "allowed") -> "PolicyDecision":
        return cls(outcome="allow", reason=reason)

    @classmethod
    def deny(cls, reason: str = "denied") -> "PolicyDecision":
        return cls(outcome="deny", reason=reason)

    @classmethod
    def requires_approval(cls, reason: str = "approval required") -> "PolicyDecision":
        return cls(outcome="approval", reason=reason)


@dataclass(frozen=True)
class PolicyRequest:
    tool_name: str
    args: dict[str, Any]
    active_tags: set[Enum]
    boundary: Enum | None
    agent_id: str | None
    run_id: str
    run_context: Any


PolicyCallable = Callable[[PolicyRequest], Awaitable[PolicyDecision]]


@dataclass(frozen=True)
class PolicySet:
    require: Sequence[PolicyCallable] = field(default_factory=tuple)
    any_of: Sequence[PolicyCallable] = field(default_factory=tuple)
    denied_message: str | None = None

    def is_empty(self) -> bool:
        return not self.require and not self.any_of


@dataclass(frozen=True)
class ToolRules:
    activates: frozenset[Enum] = field(default_factory=frozenset)
    blocked_by: frozenset[Enum] = field(default_factory=frozenset)
    boundary: Enum | None = None
    governed: bool = False
    policies: PolicySet = field(default_factory=PolicySet)

    def merge(self, other: "ToolRules") -> "ToolRules":
        return ToolRules(
            activates=frozenset(set(self.activates).union(other.activates)),
            blocked_by=frozenset(set(self.blocked_by).union(other.blocked_by)),
            boundary=other.boundary if other.boundary is not None else self.boundary,
            governed=self.governed or other.governed,
            policies=_merge_policies(self.policies, other.policies),
        )

    def is_empty(self) -> bool:
        return (
            not self.activates
            and not self.blocked_by
            and self.boundary is None
            and not self.governed
            and self.policies.is_empty()
        )


def _merge_policies(left: PolicySet, right: PolicySet) -> PolicySet:
    return PolicySet(
        require=tuple([*left.require, *right.require]),
        any_of=tuple([*left.any_of, *right.any_of]),
        denied_message=right.denied_message or left.denied_message,
    )


BoundaryRule = bool | Sequence[Enum]
BoundaryMap = dict[Enum, BoundaryRule]


def as_frozenset(values: Iterable[Enum] | None) -> frozenset[Enum]:
    return frozenset(values) if values else frozenset()
