# dao_treasury/sorting/rule.py

"""Module defining transaction sorting rules for the DAO treasury.

This module provides the `_SortRule` base class and subclasses for categorizing
`TreasuryTx` entries based on their attributes or a custom function. When a rule
is instantiated, it registers itself in the global `SORT_RULES` mapping under its
class and configures which transaction attributes to match via `_match_all`.

Examples:
    # Define a revenue rule for sales (assuming you only transact in DAI for sales)
    >>> from dao_treasury.sorting.rule import RevenueSortRule, SORT_RULES
    >>> RevenueSortRule(
    ...     txgroup='Sale',
    ...     token_address='0x6B175474E89094d879c81e570a000000000000',
    ...     symbol='DAI'
    ... )
    # Inspect rules registered for RevenueSortRule
    >>> len(SORT_RULES[RevenueSortRule])
    1

    # Iterate over all ExpenseSortRule instances
    >>> from dao_treasury.sorting.rule import ExpenseSortRule
    >>> for rule in SORT_RULES[ExpenseSortRule]:
    ...     print(rule.txgroup)

See Also:
    :const:`~dao_treasury.sorting.rule.SORT_RULES`
    :class:`~dao_treasury.sorting.rule._SortRule`
"""

from collections import defaultdict
from dataclasses import dataclass
from logging import getLogger
from typing import (
    TYPE_CHECKING,
    DefaultDict,
    Dict,
    Final,
    List,
    Optional,
    Type,
    TypeVar,
)

from brownie.convert.datatypes import EthAddress
from eth_typing import HexStr
from mypy_extensions import mypyc_attr

from dao_treasury._wallet import TreasuryWallet
from dao_treasury.types import SortFunction, SortRule, TxGroupDbid, TxGroupName

if TYPE_CHECKING:
    from dao_treasury.db import TreasuryTx


logger: Final = getLogger(__name__)
_log_debug: Final = logger.debug

SORT_RULES: DefaultDict[Type[SortRule], List[SortRule]] = defaultdict(list)
"""Mapping from sort rule classes to lists of instantiated rules, in creation order per class.

Each key is a subclass of :class:`~dao_treasury.types.SortRule` and the corresponding
value is the list of rule instances of that class.

Examples:
    >>> from dao_treasury.sorting.rule import RevenueSortRule, SORT_RULES
    >>> RevenueSortRule(txgroup='Interest', symbol='DAI')
    >>> SORT_RULES[RevenueSortRule][0].txgroup
    'Revenue:Interest'
"""

_match_all: Final[Dict[TxGroupName, List[str]]] = {}
"""An internal cache defining which matcher attributes are used for each `txgroup`."""

_MATCHING_ATTRS: Final = (
    "hash",
    "from_address",
    "from_nickname",
    "to_address",
    "to_nickname",
    "token_address",
    "symbol",
    "log_index",
)


@mypyc_attr(native_class=False)
@dataclass(kw_only=True, frozen=True)
class _SortRule:
    """Base class for defining transaction matching rules.

    When instantiated, a rule validates its inputs, determines which transaction
    attributes to match (or uses a custom function), and registers itself
    in the global `SORT_RULES` mapping under its class.

    Matched transactions are assigned to the specified `txgroup`.

    See Also:
        :const:`dao_treasury.sorting.rule.SORT_RULES`
    """

    txgroup: TxGroupName
    """Name of the transaction group to assign upon match."""

    hash: Optional[HexStr] = None
    """Exact transaction hash to match."""

    from_address: Optional[EthAddress] = None
    """Source wallet address to match."""

    from_nickname: Optional[str] = None
    """Sender nickname (alias) to match."""

    to_address: Optional[EthAddress] = None
    """Recipient wallet address to match."""

    to_nickname: Optional[str] = None
    """Recipient nickname (alias) to match."""

    token_address: Optional[EthAddress] = None
    """Token contract address to match."""

    symbol: Optional[str] = None
    """Token symbol to match."""

    log_index: Optional[int] = None
    """Log index within the transaction receipt to match."""

    func: Optional[SortFunction] = None
    """Custom matching function that takes a `TreasuryTx` and returns a bool or an awaitable that returns a bool."""

    def __post_init__(self) -> None:
        """Validate inputs, checksum addresses, and register the rule.

        - Ensures no duplicate rule exists for the same `txgroup`.
        - Converts address fields to checksummed format.
        - Determines which attributes will be used for direct matching.
        - Validates that exactly one of attribute-based or function-based matching is provided.
        - Registers the instance in :attr:`SORT_RULES` and :data:`_match_all`.
        """
        if self.txgroup in _match_all:
            raise ValueError(
                f"there is already a matcher defined for txgroup {self.txgroup}: {self}"
            )

        # ensure addresses are checksummed if applicable
        for attr in ["from_address", "to_address", "token_address"]:
            value = getattr(self, attr)
            if value is not None:
                checksummed = EthAddress(value)
                # NOTE: we must use object.__setattr__ to modify a frozen dataclass instance
                object.__setattr__(self, attr, checksummed)

        # define matchers used for this instance
        matchers = [attr for attr in _MATCHING_ATTRS if getattr(self, attr) is not None]
        _match_all[self.txgroup] = matchers

        if self.func is not None and matchers:
            raise ValueError(
                "You must specify attributes for matching or pass in a custom matching function, not both."
            )

        if self.func is None and not matchers:
            raise ValueError(
                "You must specify attributes for matching or pass in a custom matching function."
            )

        if self.func is not None and not callable(self.func):
            raise TypeError(f"func must be callable. You passed {self.func}")

        # append new instance to instances classvar
        # TODO: fix dataclass ClassVar handling in mypyc and reenable
        # self.__instances__.append(self)

        # append new instance under its class key
        SORT_RULES[type(self)].append(self)

    @property
    def txgroup_dbid(self) -> TxGroupDbid:
        """Compute the database ID for this rule's `txgroup`.

        Splits the `txgroup` string on ':' and resolves or creates the hierarchical
        `TxGroup` entries in the database, returning the final group ID.

        See Also:
            :class:`~dao_treasury.db.TxGroup`.
        """
        from dao_treasury.db import TxGroup

        txgroup = None
        for part in self.txgroup.split(":"):
            txgroup = TxGroup.get_dbid(part, txgroup)
        return txgroup

    async def match(self, tx: "TreasuryTx") -> bool:
        """Determine if the given transaction matches this rule.

        Args:
            tx: A `TreasuryTx` entity to test against this rule.

        Returns:
            True if the transaction matches the rule criteria; otherwise False.

        Examples:
            # match by symbol and recipient
            >>> rule = _SortRule(txgroup='Foo', symbol='DAI', to_address='0xabc...')
            >>> await rule.match(tx)  # where tx.symbol == 'DAI' and tx.to_address == '0xabc...'
            True

        See Also:
            :attr:`_match_all`
        """
        if matchers := _match_all[self.txgroup]:
            return all(
                getattr(tx, matcher) == getattr(self, matcher) for matcher in matchers
            )

        _log_debug("checking %s for %s", tx, self.func)
        match = self.func(tx)  # type: ignore [misc]
        return match if isinstance(match, bool) else await match


@mypyc_attr(native_class=False)
class _InboundSortRule(_SortRule):
    """Sort rule that applies only to inbound transactions (to the DAO's wallet).

    Checks that the transaction's `to_address` belongs to a known `TreasuryWallet`
    before applying the base matching logic.
    """

    async def match(self, tx: "TreasuryTx") -> bool:
        return (
            tx.to_address is not None
            and TreasuryWallet.check_membership(tx.to_address.address, tx.block)
            and await super(_InboundSortRule, self).match(tx)
        )


@mypyc_attr(native_class=False)
class _OutboundSortRule(_SortRule):
    """Sort rule that applies only to outbound transactions (from the DAO's wallet).

    Checks that the transaction's `from_address` belongs to a known `TreasuryWallet`
    before applying the base matching logic.
    """

    async def match(self, tx: "TreasuryTx") -> bool:
        return TreasuryWallet.check_membership(
            tx.from_address.address, tx.block
        ) and await super(_OutboundSortRule, self).match(tx)


@mypyc_attr(native_class=False)
class RevenueSortRule(_InboundSortRule):
    """Rule to categorize inbound transactions as revenue.

    Prepends 'Revenue:' to the `txgroup` name before registration.

    Examples:
        >>> RevenueSortRule(txgroup='Sale', to_address='0xabc...', symbol='DAI')
        # results in a rule with txgroup 'Revenue:Sale'
    """

    def __post_init__(self) -> None:
        """Prepends `self.txgroup` with 'Revenue:'."""
        object.__setattr__(self, "txgroup", f"Revenue:{self.txgroup}")
        super(RevenueSortRule, self).__post_init__()


@mypyc_attr(native_class=False)
class CostOfRevenueSortRule(_OutboundSortRule):
    """Rule to categorize outbound transactions as cost of revenue.

    Prepends 'Cost of Revenue:' to the `txgroup` name before registration.
    """

    def __post_init__(self) -> None:
        """Prepends `self.txgroup` with 'Cost of Revenue:'."""
        object.__setattr__(self, "txgroup", f"Cost of Revenue:{self.txgroup}")
        super(CostOfRevenueSortRule, self).__post_init__()


@mypyc_attr(native_class=False)
class ExpenseSortRule(_OutboundSortRule):
    """Rule to categorize outbound transactions as expenses.

    Prepends 'Expenses:' to the `txgroup` name before registration.
    """

    def __post_init__(self) -> None:
        """Prepends `self.txgroup` with 'Expenses:'."""
        object.__setattr__(self, "txgroup", f"Expenses:{self.txgroup}")
        super(ExpenseSortRule, self).__post_init__()


@mypyc_attr(native_class=False)
class OtherIncomeSortRule(_InboundSortRule):
    """Rule to categorize inbound transactions as other income.

    Prepends 'Other Income:' to the `txgroup` name before registration.
    """

    def __post_init__(self) -> None:
        """Prepends `self.txgroup` with 'Other Income:'."""
        object.__setattr__(self, "txgroup", f"Other Income:{self.txgroup}")
        super(OtherIncomeSortRule, self).__post_init__()


@mypyc_attr(native_class=False)
class OtherExpenseSortRule(_OutboundSortRule):
    """Rule to categorize outbound transactions as other expenses.

    Prepends 'Other Expenses:' to the `txgroup` name before registration.
    """

    def __post_init__(self) -> None:
        """Prepends `self.txgroup` with 'Other Expenses:'."""
        object.__setattr__(self, "txgroup", f"Other Expenses:{self.txgroup}")
        super(OtherExpenseSortRule, self).__post_init__()


@mypyc_attr(native_class=False)
class IgnoreSortRule(_SortRule):
    """Rule to ignore certain transactions.

    Prepends 'Ignore:' to the `txgroup` name before registration.
    """

    def __post_init__(self) -> None:
        """Prepends `self.txgroup` with 'Ignore:'."""
        object.__setattr__(self, "txgroup", f"Ignore:{self.txgroup}")
        super(IgnoreSortRule, self).__post_init__()


TRule = TypeVar(
    "TRule",
    RevenueSortRule,
    CostOfRevenueSortRule,
    ExpenseSortRule,
    OtherIncomeSortRule,
    OtherExpenseSortRule,
    IgnoreSortRule,
)
