from bbg_helpers import BBG_IP, init_bbg_session, retrieve_data
from db import serenitas_pool
from psycopg2.extensions import connection
from typing import List, Dict, Tuple
import datetime
import pandas as pd

def insert_bbg_tickers(conn: connection, tickers: List[str]) -> Dict[str, Tuple[str, str]]:
    """ insert bbg tickers into the database

    Parameters
    ----------
    conn: psycopg2 connection
    ticker: list of bloomberg ticker

    Returns
    -------
    dict mapping list of tickers to (company_id, seniority)
    """

    fields = ["cds_company_id", "cds_company_name", "crncy",
              "cds_restructuring_type", "cds_corp_tkr", "sw_seniority",
              "cds_curve_info", "restructuring_type_short_code"]
    securities = [f"{t} Corp" for t in tickers]
    with init_bbg_session(BBG_IP) as session:
        ticker_data = retrieve_data(session, securities, fields)
    placeholder = ",".join(["%s"] * 8)
    sql_str = ("INSERT INTO bbg_issuers(name, company_id, ticker, currency,"
               "seniority, doc_clause, cds_curve, short_code) "
               f"VALUES({placeholder}) "
               "ON CONFLICT DO NOTHING")
    d = {}
    with conn.cursor() as c:
        for k, v in ticker_data.items():
            c.execute(sql_str, (v["cds_company_name"],
                                v["cds_company_id"],
                                v["cds_corp_tkr"],
                                v["crncy"],
                                v["sw_seniority"],
                                v["cds_restructuring_type"].replace("-", " "),
                                v["cds_curve_info"]["Ticker"].to_list(),
                                v["restructuring_type_short_code"]))
            d[k.split()[0]] = (v["cds_company_id"], v["sw_seniority"])
    conn.commit()
    return d

def insert_bbg_markit_mapping(conn, d):
    sql_str = ("INSERT INTO bbg_markit_mapping VALUES(%s, %s, %s, %s, %s, %s) "
               "ON CONFLICT DO NOTHING")
    with conn.cursor() as c:
        for markit_ticker, bbg_ticker in df.itertuples(index=False):
            company_id, seniority = d[bbg_ticker]
            mkit_tier = "SNRFOR" if seniority == "Senior" else "SUBLT2"
            c.execute(sql_str, ("infinity", company_id, markit_ticker,
                                mkit_tier, 100, seniority))
    conn.commit()

def backpopulate_short_codes(conn):
    with conn.cursor() as c:
        c.execute("SELECT cds_curve[6] from bbg_issuers")
        securities = [f"{t} Corp" for t, in c]
    with init_bbg_session(BBG_IP) as session:
        short_code = retrieve_data(session, securities,
                                   ["restructuring_type_short_code",
                                    "cds_company_id"])
    sql_str = "UPDATE bbg_issuers SET short_code=%s WHERE company_id=%s"
    with conn.cursor() as c:
        for v in short_code.values():
            if 'restructuring_type_short_code' in v:
                c.execute(sql_str,
                          (v["restructuring_type_short_code"], v["cds_company_id"]))
        c.executemany(sql_str,
                      [(v["restructuring_type_short_code"], v["cds_company_id"])
                       for v in short_code.values() if 'restructuring_type_short_code' in v])
    conn.commit()

def get_bbg_ids(conn: connection, date: datetime.date):
    df = pd.read_sql_query("SELECT company_id, markit_ticker, markit_tier, "
                           "seniority FROM historical_cds_issuers(%s) "
                           "WHERE markit_ticker IS NOT NULL",
                           conn, params=(date,), index_col=["markit_ticker", "markit_tier"])
    df2 = pd.read_clipboard(header=None)
    df2.columns = ["markit_ticker", "markit_tier"]
    df2 = df2.set_index(["markit_ticker", "markit_tier"])
    return df2.join(df).groupby(level=["markit_ticker", "markit_tier"]).first()

def add_basketid(conn: connection, df: pd.DataFrame, basketid: int):
    for company_id, seniority in df.itertuples(index=False):
        with conn.cursor() as c:
            c.execute("UPDATE bbg_issuers set index_list=index_list||%s"
                      "WHERE company_id=%s AND seniority=%s",
                      (basketid, company_id, seniority))
        conn.commit()

if __name__ == "__main__":
    df = pd.read_clipboard(header=None)
    bbg_tickers = df[1].to_list()
    conn = serenitas_pool.getconn()
    d = insert_bbg_tickers(conn, bbg_tickers)
    insert_bbg_markit_mapping(conn, d)