from io import StringIO
import csv
from serenitas.utils.env import DAILY_DIR
from serenitas.utils.remote import SftpClient
import datetime
from trade_dataclasses import CDSDeal, SwaptionDeal, TerminationDeal
from headers import MTM_HEADERS

deal_registry = {
    "cds": CDSDeal,
    "swaptions": SwaptionDeal,
    "termination": TerminationDeal,
}


def process_upload(tradeid, asset_type):
    obj = deal_registry["asset_type"].from_tradeid(tradeid).to_markit()
    buf = StringIO()
    csvwriter = csv.writer(buf)
    csvwriter.writerow(MTM_HEADERS[asset_type])
    csvwriter.writerows([obj.get(h, None) for h in MTM_HEADERS[asset_type]])
    buf = buf.getvalue().encode()
    fname = f"MTM.{datetime.datetime.now():%Y%m%d.%H%M%S}.{asset_type.capitalize()}.csv"
    sftp = SftpClient.from_creds("mtm")
    sftp.put(buf, fname)
    dest = DAILY_DIR / str(datetime.date.today()) / fname
    dest.write_bytes(buf)


def upload_mtm_trades(trade_type, tradeid, conn=None):
    match trade_type:
        case "swaption":
            process_upload(
                SwaptionDeal.from_tradeid(tradeid).to_markit(), trade_type, True
            )
        case "cds":
            process_upload(CDSDeal.from_tradeid(tradeid).to_markit(), trade_type, True)
        case "termination":
            process_upload(terminations(tradeid, conn), trade_type, True)


if __name__ == "__main__":
    conn = dbconn("dawndb")
    upload_trades(conn)