import pandas as pd

from importlib import import_module
from serenitas.utils import SerenitasFileHandler
from serenitas.utils.db import dawn_engine, dbconn

from .common import get_bilateral_trades, send_email, prev_business_day

import argparse
import datetime
import logging

fh = SerenitasFileHandler("collateral_calc.log")
logger = logging.getLogger("collateral")
if not logger.handlers:
    logger.addHandler(fh)
logger.setLevel(logging.WARNING)

parser = argparse.ArgumentParser()
parser.add_argument(
    "workdate",
    nargs="?",
    type=datetime.date.fromisoformat,
    default=datetime.date.today(),
)
parser.add_argument(
    "-d", "--download", action="store_true", help="download counterparty reports"
)
parser.add_argument(
    "-s", "--send-email", action="store_true", help="send email to Globeop"
)
args = parser.parse_args()
counterparties = (
    "citi",
    "baml_isda",
    "ms",
    "gs",
    "bnp",
    "baml_fcm",
    "wells",
    "cs",
    "gs_fcm",
    "jpm",
)
funds = ("Serenitas", "Brinker", "BowdSt", "Selene")
if args.download:
    from serenitas.utils.exchange import ExchangeMessage

    em = ExchangeMessage()
    for cp in counterparties:
        cp_mod = import_module(f".{cp}", "collateral")
        for fund in funds:
            cp_mod.download_files(em, fund=fund)

workdate = prev_business_day(args.workdate)

cp_dict = {
    "Serenitas": {
        "fcms": ("baml_fcm", "wells"),
        "isda_cps": ("baml_isda", "jpm", "ms", "gs", "bnp", "cs"),
    },
    "Brinker": {"fcms": (), "isda_cps": ("ms", "gs")},
    "BowdSt": {
        "fcms": ("gs_fcm",),
        "isda_cps": ("ms", "bnp", "gs", "baml_isda", "cs", "jpm"),
    },
    "Selene": {
        "fcms": ("baml_fcm",),
        "isda_cps": ("baml_isda", "gs", "ms", "jpm"),
    },
}


def run_collateral(cp, fund, positions, positions_rates, workdate, engine):
    cp_mod = import_module("." + cp, "collateral")
    lookback = 0
    while lookback < 2:
        try:
            return cp_mod.collateral(
                workdate,
                positions,
                engine=engine,
                fund=fund,
                positions_irs=positions_rates,
            )
        except FileNotFoundError as e:
            logger.warning(e)
            lookback += 1
            workdate = prev_business_day(workdate)
        except ValueError as e:
            logger.error(f"{fund} ({cp}) : {str(e)}")
            break
        else:
            break


df = {}
fcm_mapping = {
    ("baml_fcm", "Serenitas"): "BAML",
    ("wells", "Serenitas"): "WF",
    ("gs_fcm", "BowdSt"): "GS",
    ("baml_fcm", "Selene"): "BOAFC",
}
fund_mapping = {
    "Serenitas": "SERCGMAST",
    "Brinker": "BRINKER",
    "BowdSt": "BOWDST",
    "Selene": "ISOSEL",
}

for fund in funds:
    bilat_positions = get_bilateral_trades(workdate, fund_mapping[fund], dawn_engine)
    for fcm in cp_dict[fund]["fcms"]:
        positions = pd.read_sql_query(
            "SELECT security_id, security_desc, maturity, "
            "folder, notional, currency "
            "FROM  list_cds_positions_by_strat_fcm(%s, %s, %s)",
            dawn_engine,
            params=(workdate, fcm_mapping[fcm, fund], fund_mapping[fund]),
            index_col=["security_id", "maturity"],
            parse_dates=["maturity"],
        )
        positions_rates = pd.read_sql_query(
            "SELECT cleared_trade_id, notional, effective_date, maturity_date, payreceive, fixed_rate, folder, currency "
            "FROM irs "
            "WHERE fund=%s",
            dawn_engine,
            params=(fund_mapping[fund],),
            index_col=["cleared_trade_id"],
            parse_dates=["effective_date", "maturity_date"],
        )
        df[(fund, fcm.upper())] = run_collateral(
            fcm, fund, positions, positions_rates, workdate, dawn_engine
        )
    for cp in cp_dict[fund]["isda_cps"]:
        df[(fund, cp.upper())] = run_collateral(
            cp, fund, bilat_positions, None, workdate, dawn_engine
        )

df = pd.concat(df, names=["fund", "broker", "strategy"]).reset_index()
df.strategy = df.strategy.str.replace("^(M_|SER_)?", "", 1, regex=True)
df["fund"] = df.fund.map(fund_mapping)

df = df[["date", "broker", "strategy", "Amount", "Currency", "fund"]]
sql_str = (
    "INSERT INTO strategy_im VALUES(%s, %s, %s, %s, %s, %s) "
    "ON CONFLICT (date, fund, broker, strategy, currency) DO UPDATE "
    "SET amount=EXCLUDED.amount"
)
with dawn_engine.connect() as conn:
    conn.execute(sql_str, list(df.itertuples(index=False)))

if args.send_email:
    send_email(workdate, df[df.fund == "SERCGMAST"].drop("fund", axis=1))