import datetime
from serenitas.ops.funds import Service
from serenitas.ops.trade_dataclasses import IAMDeal
from serenitas.analytics.dates import prev_business_day


def gen_old_iam(fund, cob, conn):
    """'Finds and cancels old IAM uploads"""
    with conn.cursor() as c:
        c.execute(
            "DELETE FROM iams WHERE trade_date=%s AND fund=%s RETURNING *",
            (
                cob,
                fund,
            ),
        )
        for row in c:
            d = row._asdict()
            yield IAMDeal.from_dict(**d)
        conn.commit()


def gen_new_iam(fund, cob, conn):
    """Generates new IAM deals"""
    with conn.cursor() as c:
        c.execute(
            "SELECT * FROM list_iam(%s, %s)",
            (
                cob,
                fund,
            ),
        )
        for row in c:
            d = row._asdict() | {"trade_date": cob}
            yield IAMDeal.from_dict(**d)


def gen_new_iam_offsets(fund, cob, conn):
    """'Generates offsets if the sma has already updated these iam deals"""
    _ignore = {"BOWDST": ("GS_FCM",)}
    with conn.cursor() as c:
        c.execute(
            "SELECT broker, currency, fund, sum(-start_money) AS start_money "
            "FROM list_iam(%s, %s) GROUP BY (broker, currency, fund);",
            (
                cob,
                fund,
            ),
        )
        for row in c:
            if row.broker not in _ignore[fund]:
                d = row._asdict() | {
                    "is_offset": True,
                    "folder": "M_CSH_CASH",
                    "trade_date": cob,
                }
                yield IAMDeal.from_dict(**d)


def gen_matured_iam(fund, cob, conn):
    """Sets previous days as matured"""
    with conn.cursor() as c:
        c.execute(
            "UPDATE iams SET maturity=%s WHERE maturity is NULL AND trade_date<=%s AND fund=%s RETURNING *",
            (cob, prev_business_day(cob), fund),
        )
        for row in c:
            d = row._asdict() | {}
            yield IAMDeal.from_dict(**d)
        conn.commit()


def build_iam(fund, cob, conn, upload):
    service = Service[fund]
    for old_iam in gen_old_iam(fund, cob, conn):
        service.push_trade(old_iam, "CANCEL")
    for new_iam in gen_new_iam(fund, cob, conn):
        new_iam.stage()
    if fund in ("BOWDST"):
        for new_iam_offset in gen_new_iam_offsets(fund, cob, conn):
            new_iam_offset.stage()
    for iam in IAMDeal.commit(returning=True):
        service.push_trade(iam, "NEW")
    for update_iam in gen_matured_iam(fund, cob, conn):
        service.push_trade(update_iam, "UPDATE")
    buf, dest = service.build_buffer(trade_type="iam")
    if upload:
        service.upload(buf, dest.name)
    service().clear()


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(
        "cob",
        nargs="?",
        type=datetime.date.fromisoformat,
        default=prev_business_day(datetime.date.today()),
    )
    parser.add_argument(
        "-n", "--no-upload", action="store_true", help="do not upload to CTM"
    )
    args = parser.parse_args()
    for fund in ("BOWDST",):
        build_iam(fund, args.cob, conn, not args.no_upload)