tenant-engine/src/tenant_engine/store.py

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass
from datetime import UTC, datetime
from typing import Any, Protocol

from tenant_engine.domain import (
    CapabilityRole,
    PlanAssignment,
    RoleGrant,
    Tenant,
    TenantLifecycle,
    TenantRetiredError,
)


class TenantNotFoundError(KeyError):
    pass


class VersionConflictError(RuntimeError):
    """The caller's `If-Match` version is not the tenant's current version."""

    def __init__(self, *, expected: int, actual: int) -> None:
        super().__init__(f"expected version {expected}, current version is {actual}")
        self.expected = expected
        self.actual = actual


class IdempotencyConflictError(RuntimeError):
    """An `Idempotency-Key` was reused for a materially different request."""


class TenantAlreadyExistsError(ValueError):
    pass


class GrantNotFoundError(KeyError):
    pass


class StoreUnavailableError(RuntimeError):
    """The store could not answer -- callers on a privileged decision path

    (flex-auth's live lookup) must treat this as deny, never as "zero roles".
    """


@dataclass(frozen=True, slots=True)
class DomainEvent:
    """Boundary contract's Audit Correlation Contract, in event form."""

    event_type: str
    tenant_id: str
    at: datetime
    payload: dict[str, Any]


class TenantStore(Protocol):
    """Swappable persistence seam -- domain/ and api/ depend on this, not a backend.

    Every method taking a `tenant_id` accepts either the tenant's internal
    id (caller-chosen at creation, e.g. from an admin tool) or its profile
    identifier (e.g. "tenant:friendly:binky") -- external callers like
    key-cape and flex-auth only ever have the identifier (it's the IAM
    Profile `tenant` claim value), never tenant-engine's internal id. See
    `_resolve_tenant_id` for why this had to be added after real
    cross-service testing caught the gap (found integrating key-cape's
    tenant_roles claim, KEY-WP-0005-T02).
    """

    def create_tenant(self, tenant: Tenant) -> None: ...

    def get_tenant(self, tenant_id: str) -> Tenant: ...

    def grant_role(self, grant: RoleGrant) -> None: ...

    def revoke_role(self, *, tenant_id: str, grant_id: str, at: datetime) -> RoleGrant: ...

    def active_roles(self, tenant_id: str) -> frozenset[CapabilityRole]: ...

    def assign_plan(self, assignment: PlanAssignment) -> None: ...

    def events(self) -> list[DomainEvent]: ...

    def mutate_tenant(
        self,
        *,
        tenant_id: str,
        expected_version: int,
        mutate: Callable[[Tenant], Tenant],
        event_type: str,
        evidence: dict[str, Any],
        idempotency_key: str,
        request_fingerprint: str,
    ) -> tuple[Tenant, bool]:
        """Atomically compare-and-swap a tenant record. Returns (tenant, replayed).

        One method carries all four concerns -- idempotency replay, version
        CAS, the mutation itself, and the audit event -- because they have to
        commit or fail together. Splitting them across store calls would leave
        a window where a crash yields a bumped version with no receipt (a
        retry then double-applies) or a receipt with no mutation.

        Order matters: a replayed `idempotency_key` short-circuits *before*
        the version check, because a genuine retry of an already-applied
        mutation necessarily carries a now-stale `If-Match`.
        """
        ...


class InMemoryTenantStore:
    def __init__(self) -> None:
        self._tenants: dict[str, Tenant] = {}
        self._by_identifier: dict[str, str] = {}
        self._grants: dict[str, dict[str, RoleGrant]] = {}
        self._plans: dict[str, PlanAssignment] = {}
        self._events: list[DomainEvent] = []
        # (tenant_id, idempotency_key) -> (request_fingerprint, result snapshot)
        self._receipts: dict[tuple[str, str], tuple[str, Tenant]] = {}

    def create_tenant(self, tenant: Tenant) -> None:
        if tenant.tenant_id in self._tenants:
            raise TenantAlreadyExistsError(tenant.tenant_id)
        if tenant.identifier in self._by_identifier:
            raise TenantAlreadyExistsError(tenant.identifier)
        self._tenants[tenant.tenant_id] = tenant
        self._by_identifier[tenant.identifier] = tenant.tenant_id
        self._grants[tenant.tenant_id] = {}
        self._emit(
            "tenant_created",
            tenant.tenant_id,
            {"identifier": tenant.identifier, "grouping": tenant.grouping},
        )

    def get_tenant(self, tenant_id: str) -> Tenant:
        return self._tenants[self._resolve(tenant_id)]

    def grant_role(self, grant: RoleGrant) -> None:
        resolved = self._resolve(grant.tenant_id)
        self._require_active(resolved, "grant a role")
        self._grants[resolved][grant.grant_id] = grant
        self._emit(
            "role_granted",
            resolved,
            {
                "grant_id": grant.grant_id,
                "role": grant.role.value,
                "grant_reason": grant.grant_reason,
                "correlation_id": grant.correlation_id,
            },
        )

    def revoke_role(self, *, tenant_id: str, grant_id: str, at: datetime) -> RoleGrant:
        resolved = self._resolve(tenant_id)
        try:
            grant = self._grants[resolved][grant_id]
        except KeyError:
            raise GrantNotFoundError(grant_id) from None
        revoked = grant.revoke(at=at)
        self._grants[resolved][grant_id] = revoked
        self._emit(
            "role_revoked",
            resolved,
            {"grant_id": grant_id, "role": revoked.role.value},
        )
        return revoked

    def active_roles(self, tenant_id: str) -> frozenset[CapabilityRole]:
        resolved = self._resolve(tenant_id)
        return frozenset(
            grant.role for grant in self._grants.get(resolved, {}).values() if grant.active
        )

    def assign_plan(self, assignment: PlanAssignment) -> None:
        resolved = self._resolve(assignment.tenant_id)
        self._require_active(resolved, "assign a plan")
        self._plans[resolved] = assignment
        self._emit("plan_assigned", resolved, {"plan_id": assignment.plan_id})

    def events(self) -> list[DomainEvent]:
        return list(self._events)

    def mutate_tenant(
        self,
        *,
        tenant_id: str,
        expected_version: int,
        mutate: Callable[[Tenant], Tenant],
        event_type: str,
        evidence: dict[str, Any],
        idempotency_key: str,
        request_fingerprint: str,
    ) -> tuple[Tenant, bool]:
        resolved = self._resolve(tenant_id)

        receipt = self._receipts.get((resolved, idempotency_key))
        if receipt is not None:
            fingerprint, snapshot = receipt
            if fingerprint != request_fingerprint:
                raise IdempotencyConflictError(idempotency_key)
            return snapshot, True

        current = self._tenants[resolved]
        if current.version != expected_version:
            raise VersionConflictError(expected=expected_version, actual=current.version)

        updated = mutate(current)
        self._tenants[resolved] = updated
        self._receipts[(resolved, idempotency_key)] = (request_fingerprint, updated)
        self._emit(event_type, resolved, {**evidence, "version": updated.version})
        return updated, False

    def _require_active(self, resolved_id: str, what: str) -> None:
        tenant = self._tenants[resolved_id]
        if tenant.lifecycle is not TenantLifecycle.ACTIVE:
            raise TenantRetiredError(f"cannot {what} on a retired tenant")

    def _resolve(self, tenant_id: str) -> str:
        """Resolve an internal tenant_id or a profile identifier to the

        canonical internal tenant_id every other private dict is keyed by.

        External callers (key-cape, flex-auth) only ever have the
        identifier (the IAM Profile `tenant` claim value) -- they have no
        way to know a tenant's internal tenant_id, which is caller-chosen
        at creation time and otherwise opaque. Found the hard way: a real
        cross-service check (key-cape's KEY-WP-0005-T02 against a live
        tenant-engine) returned tenant_not_found for a tenant that
        genuinely existed, because the caller only had the identifier.
        """
        resolved = self._by_identifier.get(tenant_id, tenant_id)
        if resolved not in self._tenants:
            raise TenantNotFoundError(tenant_id)
        return resolved

    def _emit(self, event_type: str, tenant_id: str, payload: dict[str, Any]) -> None:
        self._events.append(
            DomainEvent(event_type=event_type, tenant_id=tenant_id, at=datetime.now(UTC), payload=payload)
        )