#! /usr/bin/python
import argparse
import csv
import datetime
from process_queue import build_termination, upload_file
from env import DAILY_DIR
from utils.db import dbconn
from io import StringIO


def build_termination(
    base_dir,
    dawndb,
    dealid,
    fee,
    *,
    termination_date=datetime.date.today(),
    termination_amount=None,
    termination_cp=None,
    deal_type="CreditDefaultSwapDeal",
):
    """ if termination_amount is None, assume full termination
    if termination_cp is None assume termination, otherwise assignment
    """
    if deal_type == "CreditDefaultSwapDeal":
        table = "cds"
    elif deal_type == "SwaptionDeal":
        table = "swaptions"
    else:
        raise ValueError("Unkown deal_type: {deal_type}")

    with dawndb.cursor() as c:
        c.execute(
            "SELECT dealid, cp_code, notional, termination_amount, "
            "globeop_id, currency "
            f"FROM {table} where id=%s",
            (dealid,),
        )
        dealid, cp_code, notional, partial_amounts, globeopid, ccy = c.fetchone()
        if partial_amounts is not None:
            remaining_notional = notional - sum(partial_amounts)
        else:
            remaining_notional = notional
        termination_amount = termination_amount or remaining_notional
        if deal_type == "CreditDefaultSwapDeal":
            c.execute(
                f"UPDATE {table} "
                "SET termination_amount=termination_amount||%s::float8, "
                "termination_cp=termination_cp||%s::text, "
                "termination_fee=termination_fee||%s::float8, "
                "termination_date=termination_date||%s::date "
                "WHERE dealid=%s",
                (
                    termination_amount,
                    termination_cp or cp_code,
                    fee,
                    termination_date,
                    dealid,
                ),
            )
        else:
            c.execute(
                f"UPDATE {table} "
                "SET termination_amount=%s::float8, "
                "termination_cp=%s::text, "
                "termination_fee=%s::float8, "
                "termination_date=%s::date "
                "WHERE dealid=%s",
                (
                    termination_amount,
                    termination_cp or cp_code,
                    fee,
                    termination_date,
                    dealid,
                ),
            )
    dawndb.commit()
    headers = [
        "DealType",
        "DealId",
        "Action",
        "Client",
        "SubAction",
        "PartialTermination",
        "TerminationAmount",
        "TerminationDate",
        "FeesPaid",
        "FeesReceived",
        "DealFunction",
        "Reserved",
        "ClientReference",
    ]
    if deal_type == "CreditDefaultSwapDeal":
        headers += ["TradeDate", "EffectiveDate", "FirstCouponDate"]
    else:
        headers += ["Reserved"] * 3
    headers += ["FeePaymentDate", "SpecialInstructions"]

    if termination_cp is not None:
        headers += ["AssignedCounterparty"]
    else:
        headers += ["Reserved"]
    if deal_type == "CreditDefaultSwapDeal" and termination_cp is not None:
        headers += [
            "AssignmentFee",
            "AssignedFeeTradeDate",
            "AssignedFeeValueDate",
            "AssignedCustodian",
            "AssignedCashAccount",
            "Reserved",
            "FeeCurrency",
        ]
    else:
        headers += ["Reserved"] * 7
    headers += ["GoTradeId"]
    if deal_type == "CreditDefaultSwapDeal":
        headers += ["FeeComments", "ZeroOutInterestCashFlows"]
    else:
        headers += ["Reserved"] * 2
    headers += ["Reserved"] * 4
    if deal_type == "SwaptionDeal":
        headers += ["Reserved"] * 2 + ["InMoney", "FeeCurrency"]
    elif deal_type == "CreditDefaultSwapDeal":
        if termination_cp is None:
            headers += ["Reserved"] * 3
        else:
            headers += ["AssignedDealFunction"] + ["Reserved"] * 2
        headers += ["InitialMargin", "InitialMarginCurrency"]
        if termination_cp is None:
            headers += ["Reserved"] * 4 + ["CreditEventOccured"]
    d = {
        "DealType": deal_type,
        "GoTradeId": int(globeopid[3:9]),
        "Action": "Update",
        "Client": "Serenitas",
        "SubAction": "Termination",
        "PartialTermination": "Y"
        if remaining_notional - termination_amount > 0
        else "N",
        "TerminationAmount": termination_amount,
        "TerminationDate": termination_date,
        "FeesPaid": -fee if fee < 0 else None,
        "FeesReceived": fee if fee > 0 else None,
        "FeePaymentDate": (termination_date + 3 * bus_day).date(),
    }
    if "FeeCurrency" in headers:
        d["FeeCurrency"] = ccy
    if termination_cp is not None:
        d["AssignedCounterparty"] = termination_cp
    buf = StringIO()
    csvwriter = csv.DictWriter(buf, headers)
    csvwriter.writeheader()
    csvwriter.writerow(d)
    timestamp = datetime.datetime.now()
    trade_type = f"{deal_type}A" if termination_cp is not None else f"{deal_type}T"
    file_path = (
        base_dir
        / str(timestamp.date())
        / f"Serenitas.ALL.{timestamp:%Y%m%d.%H%M%S}.{trade_type}.csv"
    )
    file_path.write_bytes(buf.getvalue().encode())
    return file_path


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="helper script to terminate or assign a trade"
    )
    parser.add_argument("dealid", help="Serenitas id", type=int)
    parser.add_argument(
        "fee", help="termination fee ( >0 we receive money)", type=float
    )
    parser.add_argument(
        "--date",
        "-d",
        help="termination date, default to today's date",
        type=datetime.date.fromisoformat,
        default=datetime.date.today(),
    )
    parser.add_argument(
        "--amount",
        "-a",
        help="termination amount, if missing full termination",
        default=None,
        required=False,
    )
    parser.add_argument(
        "--counterparty",
        "-c",
        help="termination counterparty, if missing, original counterparty is used (termination)",
        default=None,
        required=False,
    )
    args = parser.parse_args()
    dawndb = dbconn("dawndb")
    if "tranche" in __file__:
        deal_type = "CreditDefaultSwapDeal"
    elif "swaption" in __file__:
        deal_type = "SwaptionDeal"
    else:
        raise RuntimeError("Incorrect deal type")

    p = build_termination(
        DAILY_DIR,
        dawndb,
        args.dealid,
        args.fee,
        termination_date=args.date,
        termination_cp=args.counterparty,
        termination_amount=args.amount,
        deal_type=deal_type,
    )
    upload_file(p)