import pandas as pd

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

from . import bus_day
from .common import get_dawn_trades, send_email
from pandas.tseries.offsets import BDay

import argparse
import logging

fh = SerenitasFileHandler("collateral_calc.log")
logger = logging.getLogger("collateral_calc")
logger.addHandler(fh)
logger.setLevel(logging.WARNING)

parser = argparse.ArgumentParser()
parser.add_argument(
    "workdate",
    nargs="?",
    type=lambda s: pd.datetime.strptime(s, "%Y-%m-%d").date(),
    default=pd.Timestamp.today().normalize(),
)
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", "ms", "gs", "baml_fcm", "baml_isda", "wells"]

if args.download:
    for cp in counterparties:
        cp_mod = import_module(f".{cp}", "collateral")
        cp_mod.download_files()

dawn_trades = get_dawn_trades(args.workdate, dawn_engine)

df = {}
mapping = {"baml_fcm": "BAML", "wells": "WF"}
for cp in counterparties:
    cp_mod = import_module("." + cp, "collateral")
    if cp in ["baml_fcm", "wells"]:
        positions = pd.read_sql_query(
            "SELECT security_id, security_desc, maturity, "
            "folder, notional, currency "
            "FROM  list_cds_positions_by_strat_fcm(%s, %s)",
            dawn_engine,
            params=(args.workdate.date(), mapping[cp]),
            index_col=["security_id", "maturity"],
        )
    else:
        positions = dawn_trades
    try:
        df[cp.upper()] = cp_mod.collateral(args.workdate, positions, dawn_engine)
    except FileNotFoundError as e:
        logger.info(e)
        df[cp.upper()] = cp_mod.collateral(
            args.workdate - bus_day, positions, dawn_engine
        )
    except ValueError as e:
        logger.error(e)
    if cp == "citi":
        args.workdate = args.workdate - BDay()

df = pd.concat(df, names=["broker", "strategy"]).reset_index()
df.strategy = df.strategy.str.replace("^(M_|SER_)?", "", 1)
df = df[["date", "broker", "strategy", "Amount", "Currency"]]
conn = dbconn("dawndb")
sql_str = (
    "INSERT INTO strategy_im VALUES(%s, %s, %s, %s, %s) "
    "ON CONFLICT (date, strategy, broker) DO UPDATE "
    "SET currency=EXCLUDED.currency, amount=EXCLUDED.amount"
)
with conn.cursor() as c:
    for t in df.itertuples(index=False):
        c.execute(sql_str, t)
conn.commit()
conn.close()

if args.send_email:
    send_email(args.workdate, df)