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_factors 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_factors(index_id, version, indexfactor,"
        "cumulativeloss, lastdate) VALUES(%(index_id)s, %(version)s, %(indexfactor)s,"
        "%(cumulativeloss)s, %(lastdate)s) RETURNING basketid"
    )

    updatestr = "UPDATE index_factors 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(conn):
    with conn.cursor() as c:
        c.execute(
            "UPDATE INDEX_factors SET redindexcode=ivm.redindexcode FROM index_version_markit ivm "
            "RIGHT JOIN index_series ON index_series.series=ivm.series AND indexsubfamily=( "
            "CASE INDEX "
            "WHEN  'HY' THEN 'CDX.NA.HY' "
            "WHEN 'IG' THEN 'CDX.NA.IG' "
            "WHEN 'EU' THEN 'iTraxx Europe' "
            "WHEN 'XO' THEN 'iTraxx Europe Crossover' "
            "WHEN 'HYBB' THEN 'CDX.NA.HY.BB' END) "
            "WHERE ivm.VERSION=index_factors.VERSION AND index_series.index_id=index_factors.index_id "
            "AND index_factors.redindexcode is NULL"
        )
    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)
        update_redcodes(conn)
        serenitas_pool.putconn(conn)