rules.py

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any, cast

from sift.common.type.v1.resource_identifier_pb2 import ResourceIdentifier
from sift.rule_evaluation.v1.rule_evaluation_pb2 import (
    EvaluateRulesRequest,
    EvaluateRulesResponse,
    RunTimeRange,
)
from sift.rule_evaluation.v1.rule_evaluation_pb2_grpc import RuleEvaluationServiceStub
from sift.rules.v1.rules_pb2 import (
    BatchDeleteRulesRequest,
    BatchGetRulesRequest,
    BatchGetRulesResponse,
    BatchUndeleteRulesRequest,
    BatchUpdateRulesRequest,
    BatchUpdateRulesResponse,
    CalculatedChannelConfig,
    ContextualChannels,
    CreateRuleRequest,
    CreateRuleResponse,
    DeleteRuleRequest,
    GetRuleRequest,
    GetRuleResponse,
    ListRulesRequest,
    RuleAssetConfiguration,
    RuleConditionExpression,
    UndeleteRuleRequest,
    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._internal.low_level_wrappers.reports import ReportsLowLevelClient
from sift_client.sift_types.rule import (
    Rule,
    RuleAction,
    RuleUpdate,
)
from sift_client.transport import GrpcClient, WithGrpcClient
from sift_client.util.util import count_non_none

if TYPE_CHECKING:
    from datetime import datetime

    from sift_client.sift_types.channel import ChannelReference
    from sift_client.sift_types.report import Report

# 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)
        response = cast("BatchGetRulesResponse", response)
        return [Rule._from_proto(rule) for rule in response.rules]

    async def create_rule(
        self,
        *,
        name: str,
        description: str,
        organization_id: str | None = None,
        client_key: str | None = None,
        asset_ids: list[str] | None = None,
        tag_ids: list[str] | None = None,
        contextual_channels: list[str] | None = None,
        is_external: bool,
        expression: str,
        channel_references: list[ChannelReference],
        action: RuleAction,
    ) -> Rule:
        """Create a new rule.

        Args:
            name: The name of the rule.
            description: The description of the rule.
            organization_id: The organization ID of the rule.
            client_key: The client key of the rule.
            asset_ids: The asset IDs of the rule.
            contextual_channels: Optional contextual channels of the rule.

        Returns:
            The rule ID of the created rule.
        """
        # Convert rule to UpdateRuleRequest
        expression_proto = RuleConditionExpression(
            calculated_channel=CalculatedChannelConfig(
                expression=expression,
                channel_references={
                    c.channel_reference: ChannelReferenceProto(name=c.channel_identifier)
                    for c in channel_references
                },
            )
        )
        conditions_request = [
            UpdateConditionRequest(
                expression=expression_proto, actions=[action._to_update_request()]
            )
        ]
        update_request = UpdateRuleRequest(
            name=name,
            description=description,
            is_enabled=True,
            organization_id=organization_id or "",
            client_key=client_key,
            is_external=is_external,
            conditions=conditions_request,
            asset_configuration=RuleAssetConfiguration(
                asset_ids=asset_ids or [],
                tag_ids=tag_ids or [],
            ),
            contextual_channels=ContextualChannels(
                channels=[ChannelReferenceProto(name=c) for c in 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=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",
        ]

        # 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.

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

        Returns:
            The updated Rule.
        """
        update.resource_id = rule.id_

        update_request = self._update_rule_request_from_update(rule, update, version_notes)

        response = await self._grpc_client.get_stub(RuleServiceStub).UpdateRule(update_request)
        updated_grpc_rule = cast("UpdateRuleResponse", response)
        # Get the updated rule
        return await self.get_rule(rule_id=updated_grpc_rule.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 = DeleteRuleRequest(**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 delete. 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 = BatchDeleteRulesRequest(**request_kwargs)
        await self._grpc_client.get_stub(RuleServiceStub).BatchDeleteRules(request)

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

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

        Returns:
            The restored Rule.

        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 = UndeleteRuleRequest(**request_kwargs)
        await self._grpc_client.get_stub(RuleServiceStub).UndeleteRule(request)
        # Get the restored rule
        return await self.get_rule(rule_id=rule_id, client_key=client_key)

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

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

        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 = BatchUndeleteRulesRequest(**request_kwargs)
        await self._grpc_client.get_stub(RuleServiceStub).BatchUndeleteRules(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,
        )

    async def evaluate_rules(
        self,
        *,
        run_id: str | None = None,
        assets: list[str] | None = None,
        all_applicable_rules: bool | None = None,
        run_start_time: datetime | None = None,
        run_end_time: datetime | None = None,
        rule_ids: list[str] | None = None,
        rule_version_ids: list[str] | None = None,
        report_template_id: str | None = None,
        tags: list[str] | None = None,
    ) -> Report | None:
        """Evaluate a rule.

        Args:
            run_id: The run ID to evaluate.
            assets: The assets to evaluate.
            run_start_time: The start time of the run.
            run_end_time: The end time of the run.
            all_applicable_rules: Whether to evaluate all rules applicable to the selected run, assets, or time range.
            rule_ids: The rule IDs to evaluate.
            rule_version_ids: The rule version IDs to evaluate.
            report_template_id: The report template ID to evaluate.
            tags: Optional tags to add to generated annotations.

        Returns:
            The result of the rule execution.
        """
        if count_non_none(run_id, assets, run_start_time, run_end_time) > 1:
            raise ValueError(
                "Pick only one run_id, assets, or (run_start_time and run_end_time) to select what to evaluate against."
            )

        all_applicable_rules = (
            None if not all_applicable_rules else True
        )  # Cast to None if False so we don't count it against other filters if they aren't opting in.
        if count_non_none(rule_ids, rule_version_ids, report_template_id, all_applicable_rules) > 1:
            raise ValueError(
                "Pick only one rule_ids, rule_version_ids, report_template_id, or all_applicable_rules to further filter which rules to evaluate."
            )

        kwargs: dict[str, Any] = {}
        if run_start_time and run_end_time:
            kwargs["run_time_range"] = RunTimeRange(
                run=run_id, start_time=run_start_time, end_time=run_end_time
            )
        if run_id:
            kwargs["run"] = ResourceIdentifier(id=run_id)
        if assets:
            kwargs["assets"] = assets
        if all_applicable_rules:
            kwargs["all_applicable_rules"] = all_applicable_rules
        if rule_ids:
            kwargs["rules"] = rule_ids
        if rule_version_ids:
            kwargs["rule_versions"] = rule_version_ids
        if report_template_id:
            kwargs["report_template"] = report_template_id
        if tags:
            kwargs["tags"] = tags

        request = EvaluateRulesRequest(**kwargs)
        response = await self._grpc_client.get_stub(RuleEvaluationServiceStub).EvaluateRules(
            request
        )
        response = cast("EvaluateRulesResponse", response)
        report_id = response.report_id
        if report_id:
            report = await ReportsLowLevelClient(self._grpc_client).get_report(report_id=report_id)
            return report
        return None