service.py

from __future__ import annotations

from dataclasses import dataclass

try:
    from functools import cache
except ImportError:
    from functools import lru_cache as cache

from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union, cast

from sift.annotations.v1.annotations_pb2 import AnnotationType
from sift.assets.v1.assets_pb2 import Asset
from sift.assets.v1.assets_pb2_grpc import AssetServiceStub
from sift.channels.v3.channels_pb2 import Channel as ChannelPb
from sift.channels.v3.channels_pb2_grpc import ChannelServiceStub
from sift.common.type.v1.user_pb2 import User as UserPb
from sift.rules.v1.rules_pb2 import (
    ANNOTATION,
    AnnotationActionConfiguration,
    BatchUpdateRulesRequest,
    BatchUpdateRulesResponse,
    CalculatedChannelConfig,
    ContextualChannels,
    CreateRuleRequest,
    GetRuleRequest,
    GetRuleResponse,
    Rule,
    RuleActionConfiguration,
    RuleAssetConfiguration,
    RuleConditionExpression,
    UpdateActionRequest,
    UpdateConditionRequest,
    UpdateRuleRequest,
)
from sift.rules.v1.rules_pb2_grpc import RuleServiceStub
from sift.users.v2.users_pb2_grpc import UserServiceStub

from sift_py._internal.cel import cel_in
from sift_py._internal.channel import channel_fqn as _channel_fqn
from sift_py._internal.channel import get_channels
from sift_py._internal.user import get_active_users
from sift_py.asset._internal.shared import list_assets_impl
from sift_py.grpc.transport import SiftChannel
from sift_py.ingestion._internal.channel import channel_reference_from_fqn
from sift_py.ingestion.channel import channel_fqn
from sift_py.rule.config import (
    ExpressionChannelReference,
    ExpressionChannelReferenceChannelConfig,
    RuleAction,
    RuleActionAnnotationKind,
    RuleActionCreateDataReviewAnnotation,
    RuleActionCreatePhaseAnnotation,
    RuleActionKind,
    RuleConfig,
)
from sift_py.yaml.rule import load_rule_modules


class RuleService:
    """
    A service for managing rules. Allows for loading rules from YAML and creating or updating them in the Sift API.

    Args:
        channel: The configured Sift channel.
        enable_caching: Enable caching on various API calls to speed up rule creation. Use this for short lived
            instantiations of the RuleService where assets, channels, users are unlikely to change.
    """

    _asset_service_stub: AssetServiceStub
    _channel_service_stub: ChannelServiceStub
    _rule_service_stub: RuleServiceStub
    _user_service_stub: UserServiceStub
    _enable_caching: bool

    def __init__(self, channel: SiftChannel, enable_caching=False):
        self._asset_service_stub = AssetServiceStub(channel)
        self._channel_service_stub = ChannelServiceStub(channel)
        self._rule_service_stub = RuleServiceStub(channel)
        self._user_service_stub = UserServiceStub(channel)
        self._enable_caching = enable_caching

    def load_rules_from_yaml(
        self,
        paths: List[Path],
        named_expressions: Optional[Dict[str, str]] = None,
    ) -> List[RuleConfig]:
        """
        Loads rules from a YAML spec, and creates or updates the rules in the Sift API.
        For more on rule YAML definitions, see `sift_py.ingestion.config.yaml.spec.RuleYamlSpec`.

        Args:
            paths: The list of YAML paths to load.
            named_expressions: The named expressions to substitute into the rules.

        Returns:
            The loaded RuleConfigs.
        """
        rule_configs = self._parse_rules_from_yaml(paths, named_expressions)

        # Create all the rules
        for rule_config in rule_configs:
            self.create_or_update_rule(rule_config)

        return rule_configs

    def create_external_rules_from_yaml(
        self,
        paths: List[Path],
        named_expressions: Optional[Dict[str, str]] = None,
    ) -> List[RuleIdentifier]:
        """
        Creates external rules from a YAML spec in the Sift API.
        For more on rule YAML definitions, see `sift_py.ingestion.config.yaml.spec.RuleYamlSpec`.
        If is_external is set in the YAML, this overrides that.

        Args:
            paths: The list of YAML paths to load.
            named_expressions: The named expressions to substitute into the rules.

        Returns:
            The loaded RuleConfigs as external rules.
        """
        rule_configs = self._parse_rules_from_yaml(paths, named_expressions)

        # Prepare all of the rules.
        for rule_config in rule_configs:
            rule_config.is_external = True
            if rule_config.rule_client_key:
                raise ValueError(
                    f"Rule of name '{rule_config.name}' requires rule_client_key to be empty"
                )

        return self.create_external_rules(rule_configs)

    def _parse_rules_from_yaml(
        self,
        paths: List[Path],
        named_expressions: Optional[Dict[str, str]] = None,
    ) -> List[RuleConfig]:
        """
        Parses rules from a YAML spec.
        For more on rule YAML definitions, see `sift_py.ingestion.config.yaml.spec.RuleYamlSpec`.

        Args:
            paths: The list of YAML paths to load.
            named_expressions: The named expressions to substitute into the rules.

        Returns:
            The parsed RuleConfigs.
        """
        module_rules = load_rule_modules(paths)

        rule_configs = []
        for rule_yaml in module_rules:
            rule_name = rule_yaml["name"]

            # First parse channel references
            yaml_channel_references = rule_yaml.get("channel_references", [])

            rule_channel_references: List[
                Union[ExpressionChannelReference, ExpressionChannelReferenceChannelConfig]
            ] = []

            for channel_ref in yaml_channel_references:
                for ref, channel_config in channel_ref.items():
                    if isinstance(channel_config, dict):
                        name = channel_config.get("name", "")
                        # NOTE: Component deprecated, but warning is thrown in the channel_fqn below
                        component = channel_config.get("component")
                    elif isinstance(channel_config, str):
                        channel_reference = channel_reference_from_fqn(channel_config)
                        name = _channel_fqn(
                            name=channel_reference.name, component=channel_reference.component
                        )
                        component = None
                    else:
                        raise ValueError(
                            f"Channel reference '{channel_config}' must be a string or a ChannelConfigYamlSpec"
                        )

                    rule_channel_references.append(
                        {
                            "channel_reference": ref,
                            "channel_identifier": channel_fqn(
                                {
                                    "channel_name": name,
                                    "component": component,
                                }
                            ),
                        }
                    )

            if not rule_channel_references:
                raise ValueError(f"Rule of name '{rule_yaml['name']}' requires channel_references")

            # Parse contextual channels
            yaml_contextual_channels: List[str] = rule_yaml.get("contextual_channels", [])
            contextual_channels: List[str] = []

            for channel_name in yaml_contextual_channels:
                if isinstance(channel_name, str):
                    contextual_channels.append(channel_name)
                else:
                    raise ValueError(f"Contextual channel '{channel_name}' must be a string")

            # Parse expression for named expressions and sub expressions
            expression = rule_yaml["expression"]
            if isinstance(expression, dict):
                expression_name = expression.get("name", "")
                if not named_expressions:
                    raise ValueError(
                        f"Rule '{rule_name}' requires named expressions, but none were provided."
                    )
                expression = named_expressions.get(expression_name, "")
                if not expression:
                    raise ValueError(f"Named expression '{expression_name}' could not be found.")

            yaml_subexprs = rule_yaml.get("sub_expressions", [])
            subexpr: Dict[str, Any] = {}
            for sub in yaml_subexprs:
                for iden, value in sub.items():
                    subexpr[iden] = value

            # Create rule actions
            tags = rule_yaml.get("tags", [])
            annotation_type = RuleActionAnnotationKind.from_str(rule_yaml["type"])
            action: RuleAction = RuleActionCreatePhaseAnnotation(tags)
            if annotation_type == RuleActionAnnotationKind.REVIEW:
                action = RuleActionCreateDataReviewAnnotation(
                    assignee=rule_yaml.get("assignee"),
                    tags=tags,
                )

            # Append rule config to list
            rule_configs.append(
                RuleConfig(
                    name=rule_name,
                    rule_client_key=rule_yaml.get("rule_client_key"),
                    description=rule_yaml.get("description", ""),
                    expression=str(expression),
                    action=action,
                    channel_references=rule_channel_references,
                    contextual_channels=contextual_channels,
                    asset_names=rule_yaml.get("asset_names", []),
                    sub_expressions=subexpr,
                    is_external=rule_yaml.get("is_external", False),
                    is_live=rule_yaml.get("is_live", False),
                )
            )

        return rule_configs

    def create_or_update_rules(self, rule_configs: List[RuleConfig]):
        """
        Create or update a list of rules via a list of RuleConfigs.
        See `sift_py.rule.config.RuleConfig` for more information on configuation parameters for rules.
        """
        for config in rule_configs:
            self.create_or_update_rule(config)

    def attach_asset(self, rule: Union[str, RuleConfig], asset_names: List[str]) -> RuleConfig:
        """
        Associates a rule with an asset by name. The asset must already exist in the Sift API.
        The provided rule may either be a rule client key, rule id, or a RuleConfig.
        """
        return self._attach_or_detach_asset(rule, asset_names, attach=True)

    def detach_asset(self, rule: Union[str, RuleConfig], asset_names: List[str]) -> RuleConfig:
        """
        Disassociates a rule from an asset by name. The asset must already exist in the Sift API.
        The provided rule may either be a rule client key, rule id, or a RuleConfig.
        """
        return self._attach_or_detach_asset(rule, asset_names, attach=False)

    def _attach_or_detach_asset(
        self, rule: Union[str, RuleConfig], asset_names: List[str], attach: bool
    ) -> RuleConfig:
        assets = self._get_assets(names=asset_names)
        if not assets:
            raise ValueError(
                f"Cannot find all assets in list '{asset_names}'. One of these assets does not exist."
            )

        if isinstance(rule, str):
            rule = cast(RuleConfig, self.get_rule(rule))
        elif not rule._rule_id and rule.rule_client_key:
            # Return provided rule with its rule_id
            # Needed to fix bug with updating an existing rule without an id
            existing_rule = cast(RuleConfig, self.get_rule(cast(str, rule.rule_client_key)))
            rule._rule_id = existing_rule._rule_id

        if attach:
            if not rule.asset_names:
                rule.asset_names = asset_names
            else:
                rule.asset_names.extend(asset_names)
        else:
            rule.asset_names = list(set(rule.asset_names) ^ set(asset_names))

        if not rule.asset_names:
            raise ValueError(f"Rule '{rule.name}' must be associated with at least one asset.")

        req = self._update_req_from_rule_config(rule)
        self._rule_service_stub.UpdateRule(req)

        return rule

    def create_or_update_rule(self, config: RuleConfig):
        """
        Create or update a rule via a RuleConfig. The config must contain a rule_client_key or an exception will be raised.
        If a rule with the given client key already exists it will be updated, otherwise it will be created.
        See `sift_py.rule.config.RuleConfig` for more information on configuation parameters for rules.
        """
        if not config.rule_client_key:
            raise ValueError(f"Rule of name '{config.name}' requires a rule_client_key")

        rule = self._get_rule_from_client_key(config.rule_client_key)
        if rule:
            self._update_rule(config, rule)
        else:
            self._create_rule(config)

    def create_external_rules(self, configs: List[RuleConfig]) -> List[RuleIdentifier]:
        """
        Create external rules via RuleConfigs. The configs must have is_external set to
        True or an exception will be raised. rule_client_key must be empty.

        Args:
            configs: The list of RuleConfigs to create as external rules.

        Returns:
            The list of RuleIdentifiers created.
        """
        for config in configs:
            if not config.is_external:
                raise ValueError(f"Rule of name '{config.name}' requires is_external to be set")
            if config.rule_client_key:
                raise ValueError(
                    f"Rule of name '{config.name}' requires rule_client_key to be empty"
                )

        update_rule_reqs = [self._update_req_from_rule_config(config) for config in configs]

        req = BatchUpdateRulesRequest(rules=update_rule_reqs)
        res = cast(BatchUpdateRulesResponse, self._rule_service_stub.BatchUpdateRules(req))

        if not res.success:
            reason = "\n".join(str(vr) for vr in res.validation_results)
            raise Exception("Failed to create external rules\n" + reason)

        return [RuleIdentifier(r.rule_id, r.name) for r in res.created_rule_identifiers]

    def _update_rule(self, updated_config: RuleConfig, rule: Rule):
        req = self._update_req_from_rule_config(updated_config, rule)
        self._rule_service_stub.UpdateRule(req)

    def _create_rule(self, config: RuleConfig):
        req = self._update_req_from_rule_config(config)
        self._rule_service_stub.CreateRule(CreateRuleRequest(update=req))

    def _update_req_from_rule(self, rule: Rule) -> UpdateRuleRequest:
        return UpdateRuleRequest(
            organization_id=rule.organization_id,
            rule_id=rule.rule_id,
            client_key=rule.client_key,
            name=rule.name,
            description=rule.description,
            conditions=[
                UpdateConditionRequest(
                    rule_condition_id=condition.rule_condition_id,
                    actions=[
                        UpdateActionRequest(
                            rule_action_id=action.rule_action_id,
                            action_type=action.action_type,
                            configuration=action.configuration,
                        )
                        for action in condition.actions
                    ],
                    expression=condition.expression,
                )
                for condition in rule.conditions
            ],
            asset_configuration=RuleAssetConfiguration(
                asset_ids=rule.asset_configuration.asset_ids,
            ),
        )

    def _update_req_from_rule_config(
        self, config: RuleConfig, rule: Optional[Rule] = None
    ) -> UpdateRuleRequest:
        if not config.expression:
            raise ValueError(
                "Cannot create a rule with an empty expression."
                "See `sift_py.rule.config.RuleConfig` for more information on rule configuration."
            )

        if not config.action:
            raise ValueError(
                "Cannot create a rule with no corresponding action."
                "See `sift_py.rule.config.RuleAction` for available actions."
            )

        # TODO: once we have TagService_ListTags we can do asset-agnostic rules via tags
        assets = self._get_assets(names=config.asset_names) if config.asset_names else None

        actions = []
        if config.action.kind() == RuleActionKind.NOTIFICATION:
            raise NotImplementedError(
                "Notification actions are not yet supported."
                "Please contact the Sift team for assistance."
            )
        elif config.action.kind() == RuleActionKind.ANNOTATION:
            if isinstance(config.action, RuleActionCreateDataReviewAnnotation):
                assignee = config.action.assignee
                user_id = None
                if assignee:
                    users = self._get_active_users(
                        filter=f"name=='{assignee}'",
                    )
                    if not users:
                        raise ValueError(f"Cannot find user '{assignee}'.")
                    if len(users) > 1:
                        raise ValueError(f"Multiple users found with name '{assignee}'.")
                    user_id = users[0].user_id

                action_config = UpdateActionRequest(
                    action_type=ANNOTATION,
                    configuration=RuleActionConfiguration(
                        annotation=AnnotationActionConfiguration(
                            assigned_to_user_id=user_id,
                            annotation_type=AnnotationType.ANNOTATION_TYPE_DATA_REVIEW,
                            # tag_ids=config.action.tags,  # TODO: Requires TagService
                        )
                    ),
                )
                actions.append(action_config)
            elif isinstance(config.action, RuleActionCreatePhaseAnnotation):
                action_config = UpdateActionRequest(
                    action_type=ANNOTATION,
                    configuration=RuleActionConfiguration(
                        annotation=AnnotationActionConfiguration(
                            annotation_type=AnnotationType.ANNOTATION_TYPE_PHASE,
                            # tag_ids=config.action.tags,  # TODO: Requires TagService
                        )
                    ),
                )

        # Get all channels that need validation (both expression and contextual)
        expression_channels = {ref["channel_identifier"] for ref in config.channel_references}
        contextual_channels = set(config.contextual_channels)
        all_channel_references = expression_channels | contextual_channels

        # Validate all channels exist in assets
        if assets and all_channel_references:
            names = [
                _channel_fqn(
                    name=channel_reference_from_fqn(ident).name,
                    component=channel_reference_from_fqn(ident).component,
                )
                for ident in all_channel_references
            ]

            # Create CEL search filters
            name_in = cel_in("name", names)

            # Validate channels are present within each asset
            for asset in assets:
                found_channels = self._get_channels(
                    filter=f"asset_id == '{asset.asset_id}' && {name_in}",
                )
                found_channels_names = [channel.name for channel in found_channels]

                missing_channels = set(names) ^ set(found_channels_names)
                if missing_channels:
                    raise RuntimeError(
                        f"Asset {asset.name} is missing channels required for rule {config.name}: {missing_channels}"
                    )

        # Create channel references map for expression channels
        channel_references = {}
        for channel_reference in config.channel_references:
            ref = channel_reference["channel_reference"]
            ident = channel_reference_from_fqn(channel_reference["channel_identifier"])
            channel_references[ref] = ident

        # Create channel references map for contextual channels
        contextual_channel_names = []
        for channel in config.contextual_channels:
            ident = channel_reference_from_fqn(channel)
            contextual_channel_names.append(ident)

        # Pass rule_id from config if retrived from existing rule
        # Will be overwritten by rule if passed to function
        # Fixes bug with missing rule_id in UpdateRuleRequest for existing rules
        if config._rule_id:
            rule_id = config._rule_id
        else:
            rule_id = ""

        organization_id = ""
        if rule:
            rule_id = rule.rule_id
            organization_id = rule.organization_id

        return UpdateRuleRequest(
            organization_id=organization_id,
            rule_id=rule_id,
            client_key=config.rule_client_key,
            name=config.name,
            description=config.description,
            conditions=[
                UpdateConditionRequest(
                    actions=actions,
                    expression=RuleConditionExpression(
                        calculated_channel=CalculatedChannelConfig(
                            expression=config.expression,
                            channel_references=channel_references,
                        )
                    ),
                )
            ],
            asset_configuration=RuleAssetConfiguration(
                asset_ids=[asset.asset_id for asset in assets] if assets else None,
            ),
            contextual_channels=ContextualChannels(channels=contextual_channel_names),
            is_external=config.is_external,
            is_live_evaluation_enabled=config.is_live,
        )

    def get_rule(self, rule: str) -> Optional[RuleConfig]:
        """
        Get a rule by rule id or client key. Returns a RuleConfig if the rule exists, otherwise None.
        """
        rule_pb = self._get_rule_from_client_key(rule) or self._get_rule_from_rule_id(rule)
        if not rule_pb:
            return None

        channel_references: List[ExpressionChannelReference] = []
        contextual_channels: List[str] = []
        expression = ""
        action: Optional[
            Union[RuleActionCreateDataReviewAnnotation, RuleActionCreatePhaseAnnotation]
        ] = None

        for condition in rule_pb.conditions:
            expression = condition.expression.calculated_channel.expression

            # Get expression channel references
            for ref, id in condition.expression.calculated_channel.channel_references.items():
                channel_references.append(
                    {
                        "channel_reference": ref,
                        "channel_identifier": id.name,
                    }
                )

            for action_config in condition.actions:
                annotation_type = action_config.configuration.annotation.annotation_type
                if annotation_type == AnnotationType.ANNOTATION_TYPE_PHASE:
                    action = RuleActionCreatePhaseAnnotation(
                        tags=[tag for tag in action_config.configuration.annotation.tag_ids],
                    )
                else:
                    assignee = action_config.configuration.annotation.assigned_to_user_id
                    action = RuleActionCreateDataReviewAnnotation(
                        assignee=assignee,
                        tags=[tag for tag in action_config.configuration.annotation.tag_ids],
                    )

        assets = self._get_assets(
            ids=[asset_id for asset_id in rule_pb.asset_configuration.asset_ids]
        )
        asset_names = [asset.name for asset in assets]

        contextual_channels = []
        for channel_ref in rule_pb.contextual_channels.channels:
            contextual_channels.append(channel_ref.name)

        rule_config = RuleConfig(
            name=rule_pb.name,
            description=rule_pb.description,
            rule_client_key=rule_pb.client_key,
            channel_references=channel_references,  # type: ignore
            contextual_channels=contextual_channels,
            asset_names=asset_names,
            action=action,
            expression=expression,
        )

        # rule_id currently required for an existing rule
        rule_config._rule_id = rule_pb.rule_id if rule_pb.rule_id else None

        return rule_config

    def _get_rule_from_client_key(self, client_key: str) -> Optional[Rule]:
        req = GetRuleRequest(client_key=client_key)
        try:
            res = cast(GetRuleResponse, self._rule_service_stub.GetRule(req))
            return res.rule or None
        except:
            return None

    def _get_rule_from_rule_id(self, rule_id: str) -> Optional[Rule]:
        req = GetRuleRequest(rule_id=rule_id)
        try:
            res = cast(GetRuleResponse, self._rule_service_stub.GetRule(req))
            return res.rule or None
        except:
            return None

    def _get_assets(self, names: List[str] = [], ids: List[str] = []) -> List[Asset]:
        if self._enable_caching:
            return self._get_assets_cached(tuple(sorted(names)), tuple(sorted(ids)))
        else:
            return list_assets_impl(self._asset_service_stub, names, ids)

    def _get_channels(self, filter: str) -> List[ChannelPb]:
        if self._enable_caching:
            return self._get_channels_cached(filter)
        else:
            return get_channels(channel_service=self._channel_service_stub, filter=filter)

    def _get_active_users(self, filter: str) -> List[UserPb]:
        if self._enable_caching:
            return self._get_active_users_cached(filter)
        else:
            return get_active_users(user_service=self._user_service_stub, filter=filter)

    @cache
    def _get_assets_cached(self, names: Tuple[str], ids: Tuple[str]) -> List[Asset]:
        return list_assets_impl(self._asset_service_stub, names, ids)

    @cache
    def _get_channels_cached(self, filter: str) -> List[ChannelPb]:
        return get_channels(channel_service=self._channel_service_stub, filter=filter)

    @cache
    def _get_active_users_cached(self, filter: str) -> List[UserPb]:
        return get_active_users(user_service=self._user_service_stub, filter=filter)


@dataclass
class RuleChannelReference:
    """
    Convenient wrapper to map rule names to relevant channel references
    when creating rules from yaml.
    """

    rule_name: str
    channel_references: Dict[str, Any]


@dataclass
class RuleIdentifier:
    """
    Wrapper around RuleIdentifier for working with external rules.
    """

    rule_id: str
    name: str