import json
from datetime import datetime
from io import BytesIO
from typing import Optional

from django.db import models
from pandas.tseries.offsets import BDay
from wbcore.contrib.io.backends import AbstractDataBackend, register

from ..utils import process_request
from .mixin import DataBackendMixin


@register("Instrument Prices", provider_key="ubs", save_data_in_import_source=True, passive_only=True)
class DataBackend(DataBackendMixin, AbstractDataBackend):
    def __init__(
        self, import_credential: Optional[models.Model] = None, ubs_bank: Optional[models.Model] = None, **kwargs
    ):
        if not ubs_bank:
            raise ValueError("The ubs company objects needs to be passed to this backend")
        self.ubs_bank = ubs_bank
        if not import_credential or not import_credential.authentication_token:
            raise ValueError("UBS backend needs a valid import credential object")
        self.authentication_token = import_credential.authentication_token

    def get_files(
        self,
        execution_time: datetime,
        obj_external_ids: list[str] = None,
        **kwargs,
    ) -> BytesIO:
        execution_date = (execution_time - BDay(1)).date()
        endpoint = "https://neo.ubs.com/api/ged-amc/external/report/v1/valuation/{0}/{1}"
        if obj_external_ids:
            for external_id in obj_external_ids:
                res_json = process_request(
                    self.authentication_token, endpoint.format(external_id, execution_date.strftime("%Y-%m-%d"))
                )
                res_json.pop("constituents", None)
                if res_json:
                    content_file = BytesIO()
                    content_file.write(json.dumps(res_json).encode())
                    file_name = f"ubs_instrument_price_{external_id}_{execution_date:%Y-%m-%d}_{int(datetime.timestamp(execution_time))}.json"
                    yield file_name, content_file
