from . import dbconn
import datetime
from io import BytesIO
import lz4
import os
import pandas as pd
from pandas.tseries.offsets import BDay
from psycopg2 import sql
import requests
import xml.etree.ElementTree as ET
import zipfile
from yieldcurve import YC, ql_to_jp
from pickle import dumps


def downloadMarkitIRData(download_date=datetime.date.today(), currency="USD"):
    conn = dbconn("serenitasdb")
    basedir = os.path.join(os.environ["DATA_DIR"], "Yield Curves")
    filename = f"InterestRates_{currency}_{download_date:%Y%m%d}"
    if not os.path.exists(os.path.join(basedir, filename + ".xml")):
        r = requests.post(f"http://www.markit.com/news/{filename}.zip")
        if "zip" in r.headers["content-type"]:
            with zipfile.ZipFile(BytesIO(r.content)) as z:
                z.extractall(path=os.path.join(os.environ["DATA_DIR"], "Yield Curves"))
        else:
            raise ValueError(r.content.decode().rstrip())

    tree = ET.parse(
        os.path.join(os.environ["DATA_DIR"], "Yield Curves", filename + ".xml")
    )
    deposits = zip(
        [e.text for e in tree.findall("./deposits/*/tenor")],
        [float(e.text) for e in tree.findall("./deposits/*/parrate")],
    )
    swaps = zip(
        [e.text for e in tree.findall("./swaps/*/tenor")],
        [float(e.text) for e in tree.findall("./swaps/*/parrate")],
    )
    effectiveasof = tree.find("./effectiveasof").text
    MarkitData = {
        "deposits": list(deposits),
        "swaps": list(swaps),
        "effectiveasof": pd.Timestamp(effectiveasof).date(),
    }
    ql_yc = YC(
        currency=currency,
        MarkitData=MarkitData,
        evaluation_date=MarkitData["effectiveasof"],
    )
    jp_yc = ql_to_jp(ql_yc)
    sql_str = f"INSERT INTO {currency}_curves VALUES(%s, %s) ON CONFLICT DO NOTHING"
    with conn.cursor() as c:
        c.execute(
            sql_str,
            (MarkitData["effectiveasof"], lz4.block.compress(jp_yc.__getstate__())),
        )
    instruments = MarkitData["deposits"] + MarkitData["swaps"]
    names = sql.SQL(", ").join([sql.Identifier(r[0]) for r in instruments])
    values = sql.SQL(", ").join(
        sql.Placeholder() * (len(instruments) + 1)
    )  # +1 for effective_date
    insert_str = sql.SQL(
        f"INSERT INTO {currency}_rates(effective_date, {{}}) VALUES({{}}) "
        "ON CONFLICT DO NOTHING"
    ).format(names, values)
    with conn.cursor() as c:
        c.execute(
            insert_str, [MarkitData["effectiveasof"]] + [r[1] for r in instruments]
        )
    conn.commit()
    conn.close()