rules.py

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any, cast

from sift.rules.v1.rules_pb2 import (
    ArchiveRuleRequest,
    BatchArchiveRulesRequest,
    BatchGetRulesRequest,
    BatchUnarchiveRulesRequest,
    BatchUpdateRulesRequest,
    BatchUpdateRulesResponse,
    CalculatedChannelConfig,
    ContextualChannels,
    CreateRuleRequest,
    CreateRuleResponse,
    GetRuleRequest,
    GetRuleResponse,
    ListRulesRequest,
    RuleAssetConfiguration,
    RuleConditionExpression,
    UnarchiveRuleRequest,
    UpdateConditionRequest,
    UpdateRuleRequest,
    UpdateRuleResponse,
)
from sift.rules.v1.rules_pb2 import (
    ChannelReference as ChannelReferenceProto,
)
from sift.rules.v1.rules_pb2_grpc import RuleServiceStub

from sift_client._internal.low_level_wrappers.base import LowLevelClientBase
from sift_client.sift_types.rule import (
    Rule,
    RuleCreate,
    RuleUpdate,
)
from sift_client.transport import GrpcClient, WithGrpcClient

if TYPE_CHECKING:
    from sift_client.sift_types.channel import ChannelReference

# Configure logging
logger = logging.getLogger(__name__)


class RulesLowLevelClient(LowLevelClientBase, WithGrpcClient):
    """Low-level client for the RulesAPI.

    This class provides a thin wrapper around the autogenerated bindings for the RulesAPI.
    """

    def __init__(self, grpc_client: GrpcClient):
        """Initialize the RulesLowLevelClient.

        Args:
            grpc_client: The gRPC client to use for making API calls.
        """
        super().__init__(grpc_client)

    async def get_rule(self, rule_id: str | None = None, client_key: str | None = None) -> Rule:
        """Get a rule by rule_id or client_key.

        Args:
            rule_id: The rule ID to get.
            client_key: The client key to get.

        Returns:
            The Rule.

        Raises:
            ValueError: If neither rule_id nor client_key is provided.
        """
        request_kwargs: dict[str, Any] = {}
        if rule_id is not None:
            request_kwargs["rule_id"] = rule_id
        if client_key is not None:
            request_kwargs["client_key"] = client_key

        request = GetRuleRequest(**request_kwargs)
        response = await self._grpc_client.get_stub(RuleServiceStub).GetRule(request)
        grpc_rule = cast("GetRuleResponse", response).rule
        return Rule._from_proto(grpc_rule)

    async def batch_get_rules(
        self, rule_ids: list[str] | None = None, client_keys: list[str] | None = None
    ) -> list[Rule]:
        """Get multiple rules by rule_ids or client_keys.

        Args:
            rule_ids: List of rule IDs to get.
            client_keys: List of client keys to get.

        Returns:
            List of Rules.

        Raises:
            ValueError: If neither rule_ids nor client_keys is provided.
        """
        if rule_ids is None and client_keys is None:
            raise ValueError("Either rule_ids or client_keys must be provided")

        request_kwargs: dict[str, Any] = {}
        if rule_ids is not None:
            request_kwargs["rule_ids"] = rule_ids
        if client_keys is not None:
            request_kwargs["client_keys"] = client_keys

        request = BatchGetRulesRequest(**request_kwargs)
        response = await self._grpc_client.get_stub(RuleServiceStub).BatchGetRules(request)
        return [Rule._from_proto(rule) for rule in response.rules]

    async def create_rule(
        self,
        *,
        create: RuleCreate,
    ) -> Rule:
        """Create a new rule.

        Args:
            create: The RuleCreate model with the rule configuration.

        Returns:
            The created Rule.
        """
        # Convert rule to UpdateRuleRequest
        expression_proto = RuleConditionExpression(
            calculated_channel=CalculatedChannelConfig(
                expression=create.expression,
                channel_references={
                    c.channel_reference: ChannelReferenceProto(name=c.channel_identifier)
                    for c in create.channel_references
                },
            )
        )
        conditions_request = [
            UpdateConditionRequest(
                expression=expression_proto,
                actions=[create.action._to_update_request()],
            )
        ]
        update_request = UpdateRuleRequest(
            name=create.name,
            description=create.description,
            is_enabled=True,
            organization_id=create.organization_id or "",
            client_key=create.client_key,
            is_external=create.is_external,
            conditions=conditions_request,
            asset_configuration=RuleAssetConfiguration(
                asset_ids=create.asset_ids or [],
                tag_ids=create.asset_tag_ids or [],
            ),
            contextual_channels=ContextualChannels(
                channels=[ChannelReferenceProto(name=c) for c in create.contextual_channels or []]
            ),  # type: ignore
        )

        request = CreateRuleRequest(update=update_request)
        created_rule = cast(
            "CreateRuleResponse",
            await self._grpc_client.get_stub(RuleServiceStub).CreateRule(request),
        )
        return await self.get_rule(rule_id=created_rule.rule_id, client_key=create.client_key)

    def _update_rule_request_from_update(
        self, rule: Rule, update: RuleUpdate, version_notes: str | None = None
    ) -> UpdateRuleRequest:
        """Create an update request from a rule and update.

        This helper exists because the Rule update protos need a pattern that is less generic than the normal update + mask pattern of other types.
        """
        model_dump = update.model_dump(exclude_unset=True, exclude_none=False)

        update_dict = {
            "version_notes": version_notes,
        }
        nontrivial_updates = [
            "expression",
            "channel_references",
            "action",
            "contextual_channels",
            "asset_ids",
            "asset_tag_ids",
        ]
        # Need to manually copy fields that will be reset even if not provided in update dict.
        copy_unset_fields = ["description", "name"]

        # Populate the trivial fields first.
        update_dict.update(
            {
                updated_field: value
                for updated_field, value in model_dump.items()
                if updated_field not in nontrivial_updates
            }
        )
        # Populate the fields that weren't updated but will be reset if not provided in request.
        for field in copy_unset_fields:
            if field not in model_dump:
                update_dict[field] = getattr(rule, field)

        # Special handling for the more complex fields.
        # Also, these must always be set.
        expression = model_dump.get("expression", rule.expression)
        channel_references: list[ChannelReference] = (
            update.channel_references
            if "channel_references" in model_dump
            else rule.channel_references
        ) or []
        action = update.action if "action" in model_dump else rule.action
        if bool(expression) != bool(channel_references):
            raise ValueError(
                "Expression and channel_references must both be provided or both be None"
            )
        expression_proto = RuleConditionExpression(
            calculated_channel=(
                CalculatedChannelConfig(
                    expression=expression,
                    channel_references={
                        c.channel_reference: ChannelReferenceProto(name=c.channel_identifier)
                        for c in channel_references
                    },
                )
                if expression
                else None
            )
        )
        conditions_request = [
            UpdateConditionRequest(
                expression=expression_proto,
                actions=[action._to_update_request()] if action else None,
            )
        ]
        update_dict["conditions"] = conditions_request  # type: ignore
        if "contextual_channels" in model_dump:
            update_dict["contextual_channels"] = ContextualChannels(  # type: ignore
                channels=[ChannelReferenceProto(name=c) for c in update.contextual_channels or []]
            )

        # This always needs to be set, so handle the defaults.
        update_dict["asset_configuration"] = RuleAssetConfiguration(  # type: ignore
            asset_ids=(update.asset_ids if "asset_ids" in model_dump else rule.asset_ids or []),
            tag_ids=(
                update.asset_tag_ids if "asset_tag_ids" in model_dump else rule.asset_tag_ids or []
            ),
        )

        update_request = UpdateRuleRequest(
            rule_id=rule.id_,
            **update_dict,  # type: ignore
        )

        return update_request

    async def update_rule(
        self, rule: Rule, update: RuleUpdate, version_notes: str | None = None
    ) -> Rule:
        """Update a rule. Also handles archive/unarchive to behave similar to other low-level clients.

        Args:
            rule: The rule to update.
            update: The update to apply.
            version_notes: Notes to include in the rule version.

        Returns:
            The updated Rule.
        """

        should_update_archive = "is_archived" in update.model_fields_set

        update.resource_id = rule.id_
        if not should_update_archive or (
            should_update_archive and len(update.model_fields_set) > 1
        ):
            update_request = self._update_rule_request_from_update(rule, update, version_notes)

            response = await self._grpc_client.get_stub(RuleServiceStub).UpdateRule(update_request)
            _ = cast("UpdateRuleResponse", response)

        if should_update_archive:
            if update.is_archived:
                await self.archive_rule(rule_id=rule.id_)
            else:
                await self.unarchive_rule(rule_id=rule.id_)

        # Get the updated rule
        return await self.get_rule(rule_id=rule.id_)

    async def batch_update_rules(self, rules: list[RuleUpdate]) -> BatchUpdateRulesResponse:
        """Batch update rules.

        Args:
            rules: List of rule updates to apply.

        Returns:
            The batch update response.
        """
        update_requests = []
        for rule_update in rules:
            rule = await self.get_rule(rule_id=rule_update.resource_id)
            request = self._update_rule_request_from_update(rule, rule_update)
            update_requests.append(request)

        request = BatchUpdateRulesRequest(rules=update_requests)  # type: ignore
        response = await self._grpc_client.get_stub(RuleServiceStub).BatchUpdateRules(request)
        return cast("BatchUpdateRulesResponse", response)

    async def archive_rule(self, rule_id: str | None = None, client_key: str | None = None) -> None:
        """Archive a rule.

        Args:
            rule_id: The rule ID to archive.
            client_key: The client key to archive.

        Raises:
            ValueError: If neither rule_id nor client_key is provided.
        """
        if rule_id is None and client_key is None:
            raise ValueError("Either rule_id or client_key must be provided")

        request_kwargs: dict[str, Any] = {}
        if rule_id is not None:
            request_kwargs["rule_id"] = rule_id
        if client_key is not None:
            request_kwargs["client_key"] = client_key

        request = ArchiveRuleRequest(**request_kwargs)
        await self._grpc_client.get_stub(RuleServiceStub).ArchiveRule(request)

    async def batch_archive_rules(
        self, rule_ids: list[str] | None = None, client_keys: list[str] | None = None
    ) -> None:
        """Batch archive rules.

        Args:
            rule_ids: List of rule IDs to archive.
            client_keys: List of client keys to archive. If both are provided, rule_ids will be used.

        Raises:
            ValueError: If neither rule_ids nor client_keys is provided.
        """
        if rule_ids is None and client_keys is None:
            raise ValueError("Either rule_ids or client_keys must be provided")

        request_kwargs: dict[str, Any] = {}
        if rule_ids is not None:
            request_kwargs["rule_ids"] = rule_ids
        if client_keys is not None:
            request_kwargs["client_keys"] = client_keys

        request = BatchArchiveRulesRequest(**request_kwargs)
        await self._grpc_client.get_stub(RuleServiceStub).BatchArchiveRules(request)

    async def unarchive_rule(
        self, rule_id: str | None = None, client_key: str | None = None
    ) -> Rule:
        """Unarchive a rule.

        Args:
            rule_id: The rule ID to unarchive.
            client_key: The client key to unarchive.

        Returns:
            The unarchived Rule.

        Raises:
            ValueError: If neither rule_id nor client_key is provided.
        """
        request_kwargs: dict[str, Any] = {}
        if rule_id is not None:
            request_kwargs["rule_id"] = rule_id
        if client_key is not None:
            request_kwargs["client_key"] = client_key

        request = UnarchiveRuleRequest(**request_kwargs)
        await self._grpc_client.get_stub(RuleServiceStub).UnarchiveRule(request)
        # Get the unarchived rule
        return await self.get_rule(rule_id=rule_id, client_key=client_key)

    async def batch_unarchive_rules(
        self, rule_ids: list[str] | None = None, client_keys: list[str] | None = None
    ) -> None:
        """Batch unarchive rules.

        Args:
            rule_ids: List of rule IDs to unarchive.
            client_keys: List of client keys to unarchive.

        Raises:
            ValueError: If neither rule_ids nor client_keys is provided.
        """
        request_kwargs: dict[str, Any] = {}
        if rule_ids is not None:
            request_kwargs["rule_ids"] = rule_ids
        if client_keys is not None:
            request_kwargs["client_keys"] = client_keys

        request = BatchUnarchiveRulesRequest(**request_kwargs)
        await self._grpc_client.get_stub(RuleServiceStub).BatchUnarchiveRules(request)

    async def list_rules(
        self,
        *,
        filter_query: str | None = None,
        order_by: str | None = None,
        page_size: int | None = None,
        page_token: str | None = None,
    ) -> tuple[list[Rule], str | None]:
        """List rules."""
        request_kwargs: dict[str, Any] = {}
        if filter_query is not None:
            request_kwargs["filter"] = filter_query
        if order_by is not None:
            request_kwargs["order_by"] = order_by
        if page_size is not None:
            request_kwargs["page_size"] = page_size
        if page_token is not None:
            request_kwargs["page_token"] = page_token

        request = ListRulesRequest(**request_kwargs)
        response = await self._grpc_client.get_stub(RuleServiceStub).ListRules(request)
        return [Rule._from_proto(rule) for rule in response.rules], response.next_page_token

    async def list_all_rules(
        self,
        *,
        filter_query: str | None = None,
        order_by: str | None = None,
        max_results: int | None = None,
        page_size: int | None = None,
    ) -> list[Rule]:
        """List all rules."""
        return await self._handle_pagination(
            self.list_rules,
            kwargs={"filter_query": filter_query},
            page_size=page_size,
            order_by=order_by,
            max_results=max_results,
        )