from serenitas.utils.misc import rename_keys
import datetime
from collections import defaultdict
from psycopg2.errors import UniqueViolation
from serenitas.analytics.dates import bus_day
from serenitas.ops.funds import Service

_iam_brokers = {
    "BAML_ISDA": "BOANNY",
    "CS": "CSITLN",
    "GS": "GOLINY",
    "BNP": "BNPBNY",
    "MS": "MSCILN",
    "JPM": "JPCBNY",
    "GS_FCM": "GOLDNY",
}
_custodian = {"GOLDNY": "GS"}
_cash_account = {"GOLDNY": "057363418ICE-CDS"}


def iam_serialize(obj, action, trade_date):
    rename_keys(
        obj,
        {
            "dealid": "Deal Id",
            "counterparty": "Counterparty",
            "currency": "Currency",
            "start_money": "StartMoney",
            "trade_date": "Trade Date",
        },
    )

    deal_fields = {
        "Deal Type": "IamDeal",
        "Client": "HEDGEMARK",
        "Fund": "BOS_PAT_BOWDOIN",
        "Custodian": _custodian.get(obj["Counterparty"], "BNY"),
        "Cash Account": _cash_account.get(obj["Counterparty"], 751254),
        "State": "Valid",
        "SettlementDate": obj["Trade Date"],
        "Basis": "ACT/360",
        "MarginType": "Variation",
        "DealFunction": "OTC",
        "Folder": "M_CSH_CASH" if obj["strategy"] == "CSH_CASH" else obj["strategy"],
        "Action": action,
        "TransactionIndicator": "DEPOSIT" if obj["StartMoney"] > 0 else "LOAN",
        "ExpirationDate": trade_date
        if action == "UPDATE"
        else None,  # We are updating old IAMDeals to expire
        "CallNoticeIndicator": "24H" if action == "NEW" else None,  # This is a new IAM
    }
    obj["StartMoney"] = abs(obj["StartMoney"])
    obj.update(deal_fields)
    return obj


def gen_offset(totals, trade_date):
    offsets = [
        (trade_date, "NEW", "CSH_CASH", broker, None, -amount, "USD", True)
        for broker, amount in totals.items()
    ]
    return offsets


def gen_strat_alloc(conn, trade_date):
    new_iam = []
    totals = defaultdict(int)  # Need to figure out how much to offset
    with conn.cursor() as c:
        strategy_allocation = (
            "SELECT broker, amount, currency, strategy FROM (SELECT *, rank() "
            "OVER(PARTITION BY broker,fund ORDER BY date desc) FROM strategy_im si "
            "WHERE fund = 'BOWDST' AND date<=%s ORDER BY date DESC) test WHERE RANK=1 and abs(amount) >= .01;"
        )
        c.execute(strategy_allocation, (trade_date,))
        for row in c:
            new_iam.append(
                (
                    trade_date,
                    "NEW",
                    row.strategy,
                    _iam_brokers[row.broker],
                    None,
                    row.amount,
                    row.currency,
                    False,
                )
            )
            if row.broker != "GS_FCM":  # HM doesn't book FCM, no need to offset
                totals[_iam_brokers[row.broker]] += row.amount
    new_iam += gen_offset(totals, trade_date)
    return new_iam


def process_trades(conn, trade_date):
    # We will be grabbing the IAM deals from the DB for the new trades and for the updated trades
    deals = []
    actions = [
        (
            "NEW",
            "UPDATE iam_tickets set uploaded=True where maturity is null and trade_date =%s "
            "and action='NEW' and not uploaded returning *",
            (trade_date,),
        ),
        (
            "UPDATE",
            "UPDATE iam_tickets set maturity=%s where trade_date != %s and maturity is null and action='NEW' returning *",
            (trade_date, trade_date),
        ),
    ]
    for action, query, params in actions:
        with conn.cursor() as c:
            c.execute(query, params)
            for row in c:
                deals.append(iam_serialize(row._asdict(), action, trade_date))
    return deals


def gen_iam_deals(conn, trade_date):
    new_iam = gen_strat_alloc(conn, trade_date)
    iam_deals = []
    with conn.cursor() as c:
        insert_query = (
            "INSERT INTO iam_tickets(trade_date, action, strategy, counterparty, maturity, start_money, currency, booked_offset) "
            "VALUES (%s, %s, %s, %s, %s, %s, %s, %s)"
        )
        try:
            c.executemany(insert_query, new_iam)
        except UniqueViolation:
            # We already uploaded the IAM tickets today in that case, we need to update and cancel the old uploads
            conn.rollback()
            c.execute(
                "DELETE FROM iam_tickets where trade_date=%s returning *", (trade_date,)
            )
            cancel_iam_deals = [
                iam_serialize(row._asdict(), "CANCEL", trade_date) for row in c
            ]
            iam_deals += cancel_iam_deals
            c.executemany(insert_query, new_iam)
    iam_deals += process_trades(conn, trade_date)
    return iam_deals


def process_upload(conn, trade_date, iam_deals):
    bowdst = Service["BOWDST"]()
    bowdst.staging_queue.extend(iam_deals)
    if bowdst.staging_queue:
        buf, dest = bowdst.build_buffer("iam")
        bowdst.upload(buf, dest.name)


if __name__ == "__main__":
    import argparse
    from serenitas.utils.db import dbconn

    conn = dbconn("dawndb")
    parser = argparse.ArgumentParser(description="Generate IAM file for globeop")
    parser.add_argument(
        "date",
        nargs="?",
        type=datetime.date.fromisoformat,
        default=(datetime.date.today() - bus_day).date(),
    )
    args = parser.parse_args()
    iam_deals = gen_iam_deals(conn, args.date)
    process_upload(conn, args.date, iam_deals)
    conn.commit()