from serenitas.utils.env import DAILY_DIR
from serenitas.utils.exchange import ExchangeMessage
import datetime
from collateral.baml_isda import download_from_secure_id
from bs4 import BeautifulSoup
from urllib.parse import urlsplit, parse_qs, urlunsplit
import logging

logger = logging.getLogger(__name__)


def download_messages(em, counterparty):
    for msg in em.get_msgs(
        20,
        path=["Interest", counterparty],
    ):
        if counterparty == "CITI":
            base_dir = DAILY_DIR / f"{counterparty}_reports" / "Interest Statements"
        else:
            base_dir = (
                DAILY_DIR
                / "Serenitas"
                / f"{counterparty}_reports"
                / "Interest Statements"
            )
        if msg.datetime_sent.date() > datetime.date(2021, 5, 1):
            if counterparty == "BAML":
                soup = BeautifulSoup(msg.body, features="lxml")
                a = soup.find("a")
                url = urlsplit(a["href"])
                query = parse_qs(url.query)
                base_url = urlunsplit(url[:2] + ("",) * 3)
                try:
                    download_from_secure_id(
                        query["id"][0], query["brand"][0], base_dir, base_url
                    )
                except ValueError as e:
                    logging.error(e)
                    continue
                continue
            for attach in msg.attachments:
                fname = attach.name
                if (counterparty == "CS") and not ("Interest" in fname):
                    continue
                p = base_dir / fname
                if not p.parent.exists():
                    p.parent.mkdir(parents=True)
                if not p.exists():
                    p.write_bytes(attach.content)


em = ExchangeMessage()
counterparties = ["BAML", "BNP", "CITI", "CS", "GS", "MS"]

for cp in counterparties:
    download_messages(em, cp)