from . import dbconn
import datetime
from env import DATA_DIR
from io import BytesIO
import lz4
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


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

    tree = ET.parse(curve_file)
    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": datetime.date.fromisoformat(effectiveasof),
    }
    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"], 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()