from bbg_helpers import BBG_IP, init_bbg_session, retrieve_data
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 tickers

    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.get("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: connection, d: Dict[str, Tuple[str, str]], df: pd.DataFrame):
    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: connection):
    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, df: pd.DataFrame,
                date: datetime.date = datetime.date.today()) -> pd.DataFrame:
    df_mapping = 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"])
    df.columns = ["markit_ticker", "markit_tier"]
    df = df.set_index(["markit_ticker", "markit_tier"])
    df = df.join(df_mapping).groupby(level=["markit_ticker", "markit_tier"]).first()
    df.company_id = df.company_id.astype("Int64")
    return df

def get_basketid(conn: connection, name: str) -> Tuple[int, int]:
    with conn.cursor() as c:
        c.execute("SELECT max(series)+1 FROM index_version where index='BS'")
        serie,  = c.fetchone()
        c.execute("INSERT INTO index_version(index, series, version, redindexcode) "
                  "VALUES('BS', %s, 1, %s) RETURNING basketid",
                  (serie, name))
        basket_id, = c.fetchone()
    conn.commit()
    return serie, basket_id

def add_basket_constituents(conn: connection, basketid: int,
                            df: pd.DataFrame = pd.read_clipboard(header=None)):

    df_bbg = get_bbg_ids(conn, df.iloc[:, :2])
    df.columns = ["markit_ticker", "markit_tier", "weight"]
    df = df.set_index(["markit_ticker", "markit_tier"])
    df = df.join(df_bbg)
    with conn.cursor() as c:
        c.executemany("INSERT INTO basket_constituents VALUES(%s, %s, %s, %s)",
                      [(t.company_id, t.seniority, basketid, t.weight)
                       for t in df.itertuples(index=False)])
    conn.commit()

if __name__ == "__main__":
    df = pd.read_clipboard(header=None)
    from db import serenitas_pool
    conn = serenitas_pool.getconn()
    df = get_bbg_ids(conn, pd.read_clipboard(header=None))
    #missing tickers
    missing_tickers = df[df.company_id.isnull()]
    df_mapping = pd.read_clipboard(header=None)
    bbg_tickers = df_mapping[1].to_list()
    d = insert_bbg_tickers(conn, bbg_tickers)
    insert_bbg_markit_mapping(conn, d, df_mapping)
    serenitas_pool.putconn(conn)