from db import serenitas_pool
import datetime
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")
    print(sqlstr)
    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:
            r['indexfactor'] -= r['orig_weight']
            r['version'] += 1
            r['cumulativeloss'] += (100-recovery) * r['orig_weight'] /100
            r['lastdate'] = 'infinity'
            c.execute(insertstr, r)
            newids[r['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'

    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:
        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)
        serenitas_pool.putconn(conn,)