from serenitas.utils.db import serenitas_pool
from sys import argv


def get_recovery(company_id: int, seniority: str, conn):
    with conn.cursor() as c:
        c.execute(
            "SELECT recovery * 100, auction_date FROM defaulted "
            "WHERE id=%s AND seniority=%s",
            (company_id, seniority),
        )
        return c.fetchone()


def affected_indices(company_id: int, seniority: str, conn):
    """returns the list of indices containing company_id"""
    sqlstr = (
        "SELECT b.*, a.curr_weight*b.indexfactor AS orig_weight "
        "FROM basket_constituents_current a "
        "JOIN (SELECT * FROM index_version WHERE lastdate='infinity') b "
        "USING (basketid) "
        "WHERE company_id=%s AND seniority=%s"
    )
    with conn.cursor() as c:
        c.execute(sqlstr, (company_id, seniority))
        recordslist = c.fetchall()
    return recordslist


def create_newindices(recordslist, recovery, lastdate, conn):
    """create the new indices versions and update the old"""
    insertstr = """INSERT INTO index_version(Index, Series, Version, IndexFactor,
    CumulativeLoss, lastdate) Values(%(index)s, %(series)s, %(version)s, %(indexfactor)s,
    %(cumulativeloss)s, %(lastdate)s) RETURNING basketid"""

    updatestr = "UPDATE index_version SET lastdate=%s WHERE basketid=%s"

    with conn.cursor() as c:
        newids = {}
        for r in recordslist:
            d = r._asdict()
            d["indexfactor"] -= d["orig_weight"]
            d["version"] += 1
            d["cumulativeloss"] += (100 - recovery) * d["orig_weight"] / 100
            d["lastdate"] = "infinity"
            c.execute(insertstr, d)
            newids[d["basketid"]] = c.fetchone()[0]
        for oldid in newids.keys():
            c.execute(updatestr, (lastdate, oldid))
    conn.commit()
    return newids


def update_indexmembers(newids, company_id, seniority, conn):
    with conn.cursor() as c:
        for oldid, newid in newids.items():
            c.execute(
                "INSERT INTO basket_constituents "
                "(SELECT company_id, seniority, %s, weight "
                "FROM basket_constituents "
                "WHERE basketid=%s AND NOT (company_id=%s AND seniority=%s))",
                (newid, oldid, company_id, seniority),
            )
    conn.commit()


def update_redcodes(index_type, conn):
    if index_type == "HY":
        index_subfamily = "CDX.NA.HY"
    elif index_type == "IG":
        index_subfamily = "CDX.NA.IG"
    elif index_type == "EU":
        index_subfamily = "iTraxx Europe"
    elif index_type == "XO":
        index_subfamily = "iTraxx Europe Crossover"
    elif index_type == "HYBB":
        index_subfamily = "CDX.NA.HY.BB"
    elif index_type == "BS":
        return

    with conn.cursor() as c:
        c.execute(
            "UPDATE index_version SET redindexcode=index_version_markit.redindexcode "
            "FROM index_version_markit WHERE index_version.series=index_version_markit.series "
            "AND index_version.version=index_version_markit.version "
            "AND index_version.index=%s "
            "AND index_version_markit.indexsubfamily=%s"
            "AND index_version.redindexcode IS NULL",
            (index_type, index_subfamily),
        )
    conn.commit()


if __name__ == "__main__":
    if len(argv) == 1:
        print(
            """Usage:
    python handle_default.py <company_id> <seniority>
For instance:
    python handle_default.py 210065 Senior"""
        )
    else:
        from markit_red import download_report, update_redcodes as upd_redcodes

        fname = next(download_report("REDIndexCodes"))
        upd_redcodes(fname[0])
        conn = serenitas_pool.getconn()
        company_id = int(argv[1])
        seniority = argv[2]
        recovery, lastdate = get_recovery(company_id, seniority, conn)
        recordslist = affected_indices(company_id, seniority, conn)
        newids = create_newindices(recordslist, recovery, lastdate, conn)
        update_indexmembers(newids, company_id, seniority, conn)
        for index in set(r.index for r in recordslist):
            update_redcodes(index, conn)
        serenitas_pool.putconn(conn)