from serenitas.ops.funds import Service
from serenitas.ops.headers import HEADERS


class UMB(Service, service_name="UMB"):
    filepath_pattern = "Serenitas.ALL.{timestamp:%Y%m%d.%H%M%S}.csv"
    credential = "umb"

    @classmethod
    def set_headers(cls, *args):
        cls.headers = HEADERS["bond"] + [
            "DTC",
            "AccruedPayment",
            "PrincipalPayment",
            "CurrentFace",
        ]

    @classmethod
    def push_trade(cls, trade, action):
        cls.staging_queue.append(trade.to_umb(action))


def get_service(account_counterparty):
    match account_counterparty:
        case "UMB":
            return UMB
        case "BONY":
            return Service[account_counterparty]
        case "BBH":
            return Service["BRINKER"]
        case _:
            return None