aggregator.py

from dataclasses import replace
from datetime import datetime
from decimal import Decimal
from typing import Any

from common.models.records import RawRecord
from common.utils.helper import sort_records_for_aggregation


class BaseAggregator:
    def __init__(self, config=None):
        """
        Base constructor to store configuration for all aggregators.
        """
        self.config = config

    def aggregate_lines(self, records: list[RawRecord]) -> list[RawRecord]:
        raise NotImplementedError("Implementation missing")

    @staticmethod
    def safe_float(value):
        """Convert to float, return 0.0 if the value is empty or invalid."""
        try:
            return float(value) if value else 0.0
        except ValueError:
            return 0.0

    @staticmethod
    def sum_round(value1, value2, decimals=8):
        """Sum two values and round the result to the specified number of decimals."""

        v1 = BaseAggregator.safe_float(value1)
        v2 = BaseAggregator.safe_float(value2)
        return round(v1 + v2, decimals)

    @staticmethod
    def values_equal(val1: Any, val2: Any) -> bool:
        """
        Compare two values for equality, considering None as equal.

        Works for str, int, float, Decimal, datetime, or None.
        """
        # Both None
        if val1 is None and val2 is None:
            return True

        # One is None, the other not
        if val1 is None or val2 is None:
            return False

        # Otherwise, normal equality
        return val1 == val2

    @staticmethod
    def _is_coin_buy(record: RawRecord) -> bool:
        """
        A coin buy occurs when a crypto asset is bought in exchange for
        fiat currency or another asset.
        """

        crypto_currencies = {
            "BTC",
            "ETH",
            "ADA",
            "SOL2",
            "LUNA2",
            "LUNA3",
            "DFI",
            "BNB",
            "XRP",
            "KFEE",
        }

        stablecoins = {"USDC", "USDT", "BUSD"}
        fiat_currencies = {"EUR", "USD"}

        buy_curr = record.buy_currency
        sell_curr = record.sell_currency

        # Buying crypto with fiat
        if buy_curr in crypto_currencies and sell_curr in fiat_currencies:
            return True

        # Buying crypto with stablecoin
        if buy_curr in crypto_currencies and sell_curr in stablecoins:
            return True

        return False

    @staticmethod
    def _is_coin_sell(record: RawRecord) -> bool:
        """
        A coin sell occurs when a crypto asset is sold in exchange for
        fiat currency or another asset.
        """

        crypto_currencies = {
            "BTC",
            "ETH",
            "ADA",
            "SOL2",
            "LUNA2",
            "LUNA3",
            "DFI",
            "BNB",
            "XRP",
            "KFEE",
        }

        stablecoins = {"USDC", "USDT", "BUSD"}
        fiat_currencies = {"EUR", "USD"}

        buy_curr = record.buy_currency
        sell_curr = record.sell_currency

        # Buying crypto with fiat
        if sell_curr in crypto_currencies and buy_curr in fiat_currencies:
            return True

        # Buying crypto with stablecoin
        if sell_curr in crypto_currencies and buy_curr in stablecoins:
            return True

        return False

    @staticmethod
    def _set_time(date_value: datetime, new_time: str) -> datetime:
        """
        Replace the time part of a datetime with a given HH:MM:SS value.
        """
        hours, minutes, seconds = map(int, new_time.split(":"))

        return datetime(
            year=date_value.year,
            month=date_value.month,
            day=date_value.day,
            hour=hours,
            minute=minutes,
            second=seconds,
        )


class CoinTrackingAggregator(BaseAggregator):
    def _is_aggregation_applicable(
        self, current_line: RawRecord, next_line: RawRecord
    ) -> bool:
        if (
            BaseAggregator.values_equal(current_line.type, next_line.type)
            and BaseAggregator.values_equal(
                current_line.buy_currency, next_line.buy_currency
            )
            and BaseAggregator.values_equal(
                current_line.sell_currency, next_line.sell_currency
            )
            and BaseAggregator.values_equal(
                current_line.fee_currency, next_line.fee_currency
            )
            and BaseAggregator.values_equal(current_line.exchange, next_line.exchange)
            and BaseAggregator.values_equal(current_line.group, next_line.group)
            and BaseAggregator.values_equal(current_line.comment, next_line.comment)
            and BaseAggregator.values_equal(
                current_line.datetime.date(), next_line.datetime.date()
            )  # nur Datum vergleichen
        ):
            return True
        else:
            return False

    def _adjust_timestamp(self, record: RawRecord) -> RawRecord:
        """
        Adjust the timestamp of a record based on its business meaning.

        Rules:
        - Deposits are set to 00:01:00
        - Coin buys are set to 23:55:00
        - Coin sells are set to 23:56:00
        """

        if record.type == "Deposit":
            new_datetime = BaseAggregator._set_time(record.datetime, "00:01:00")
            return replace(record, datetime=new_datetime)

        if record.type == "Trade":
            if self._is_coin_buy(record):
                new_datetime = BaseAggregator._set_time(record.datetime, "23:55:00")
                return replace(record, datetime=new_datetime)
            if self._is_coin_sell(record):
                new_datetime = BaseAggregator._set_time(record.datetime, "23:56:00")
                return replace(record, datetime=new_datetime)

        # Default: no change
        return record

    def _adjust_margin_timestamp(self, record: RawRecord) -> RawRecord:
        """
        Adjust the timestamp of a record based on its business meaning.

        Rules:
        - Margin Profit: 00:01:00
        - deactivated: Margin Fee: +1 min (except the old time is 23:59)
        """

        # Margin Fee: +1 Minute (Limit 23:59)
        # if record.type == "Margin Fee":
        #     if record.date.hour == 23 and record.date.minute == 59:
        #         return record
        #     new_date = record.date + timedelta(minutes=1)
        #     return replace(record, date=new_date)

        # Margin Profit: -1 Minute (Limit 00:00)
        if record.type == "Margin Profit":
            new_datetime = BaseAggregator._set_time(record.datetime, "00:01:00")
            return replace(record, datetime=new_datetime)

        # Default: no change
        return record

    def _sort_result(self, records: list[RawRecord]) -> list[RawRecord]:
        """
        Sort records chronologically by date.
        """
        return sorted(records, key=lambda r: r.datetime)

    def aggregate_lines(self, records: list[RawRecord]) -> list[RawRecord]:
        """
        The aggregation process consolidates multiple transactions of one day (and further criterias) into a single daily entry.
        The specific logic for the time adjustment is documented in the `_adjust_timestamp` method.
        """
        if len(records) <= 1:
            return records

        # Ensure data is in the correct order
        sort_records_for_aggregation(records)

        result: list[RawRecord] = []

        aggr_buy = Decimal("0")
        aggr_sell = Decimal("0")
        aggr_fee = Decimal("0")
        aggregation_happened = False
        aggr_count = 1

        i = 0
        while i < len(records) - 1:
            current = records[i]
            next_rec = records[i + 1]

            if self._is_aggregation_applicable(current, next_rec):
                aggr_buy += current.buy_amount
                aggr_sell += current.sell_amount
                aggr_fee += current.fee_amount
                aggregation_happened = True
                aggr_count += 1
            else:
                if aggregation_happened:
                    # Finalize the current record with aggregated values
                    current = self._finalize_aggregation(
                        current, aggr_buy, aggr_sell, aggr_fee, aggr_count
                    )

                    aggr_buy = Decimal("0")
                    aggr_sell = Decimal("0")
                    aggr_fee = Decimal("0")
                    aggr_count = 1

                current = self._adjust_margin_timestamp(current)
                result.append(current)
                aggregation_happened = False

            i += 1

        # letzte Zeile behandeln
        last = records[-1]
        if aggregation_happened:
            last = self._finalize_aggregation(
                last, aggr_buy, aggr_sell, aggr_fee, aggr_count
            )

        last = self._adjust_margin_timestamp(last)
        result.append(last)

        return self._sort_result(result)

    def _finalize_aggregation(
        self, base_rec: RawRecord, buy: Decimal, sell: Decimal, fee: Decimal, count: int
    ) -> RawRecord:
        """
        Finalizes the aggregation for a group of records by summing up amounts
        and using the Tx-ID to show the count of aggregated entries.
        """
        # We use replace to create a new instance with the final values
        final_rec = replace(
            base_rec,
            buy_amount=base_rec.buy_amount + buy,
            sell_amount=base_rec.sell_amount + sell,
            fee_amount=base_rec.fee_amount + fee,
            comment=f"Aggregated records: {count}",
            tx_id="",
            lpn="",
        )

        # Ensure the timestamp is adjusted according to the tool's rules
        final_rec = self._adjust_timestamp(final_rec)

        return final_rec


class AggregatorFactory:
    @staticmethod
    def get_aggregator(format: str) -> BaseAggregator:
        if format == "CoinTracking":
            return CoinTrackingAggregator()
        else:
            raise ValueError(f"Unknown format: {format}")