import lusid
import pandas as pd
from lusidtools import cocoon
from lusidtools.cocoon.utilities import checkargs, strip_whitespace
from lusidtools.cocoon.validator import Validator
from lusidtools.cocoon.async_tools import run_in_executor
import asyncio
from lusidtools.cocoon.dateorcutlabel import DateOrCutLabel
import logging


class BatchLoader:
    """
    This class contains all the methods used for loading data in batches. The @run_in_executor decorator makes the
    synchronous functions awaitable
    """

    @staticmethod
    @run_in_executor
    def load_instrument_batch(
        api_factory: lusid.utilities.ApiClientFactory, instrument_batch: list, **kwargs
    ) -> lusid.models.UpsertInstrumentsResponse:
        """
        Upserts a batch of instruments to LUSID

        :param lusid.utilities.ApiClientFactory api_factory: The api factory to use
        :param list[lusid.models.InstrumentDefinition] instrument_batch: The batch of instruments to upsert

        :return: UpsertInstrumentsResponse: The response from LUSID
        """

        # Ensure that the list of allowed unique identifiers exists
        if "unique_identifiers" not in list(kwargs.keys()):
            unique_identifiers = cocoon.instruments.get_unique_identifiers(
                api_factory=api_factory
            )
        else:
            unique_identifiers = kwargs["unique_identifiers"]

        @checkargs
        def get_alphabetically_first_identifier_key(
            instrument: lusid.models.InstrumentDefinition, unique_identifiers: list
        ):
            """
            Gets the alphabetically first occurring unique identifier on an instrument and use it as the correlation
            id on the request

            :param lusid.models.InstrumentDefinition instrument: The instrument to create a correlation id for
            :param list[str] unique_identifiers: The list of allowed unique identifiers

            :return: str: The correlation id to use on the request
            """
            unique_identifiers_populated = list(
                set(unique_identifiers).intersection(
                    set(list(instrument.identifiers.keys()))
                )
            )
            unique_identifiers_populated.sort()
            first_unique_identifier_alphabetically = unique_identifiers_populated[0]
            return f"{first_unique_identifier_alphabetically}: {instrument.identifiers[first_unique_identifier_alphabetically].value}"

        return api_factory.build(lusid.api.InstrumentsApi).upsert_instruments(
            instruments={
                get_alphabetically_first_identifier_key(
                    instrument, unique_identifiers
                ): instrument
                for instrument in instrument_batch
            }
        )

    @staticmethod
    @run_in_executor
    def load_quote_batch(
        api_factory: lusid.utilities.ApiClientFactory, quote_batch: list, **kwargs
    ) -> lusid.models.UpsertQuotesResponse:
        """
        Upserts a batch of quotes into LUSID

        :param lusid.utilities.ApiClientFactory api_factory: The api factory to use
        :param list[lusid.models.UpsertQuoteRequest] quote_batch: The batch of quotes to upsert
        :param str scope: The scope to upsert the quotes into

        :return: lusid.models.UpsertQuotesResponse: The response from LUSID
        """

        if "scope" not in list(kwargs.keys()):
            raise KeyError(
                "You are trying to load quotes without a scope, please ensure that a scope is provided."
            )

        return api_factory.build(lusid.api.QuotesApi).upsert_quotes(
            scope=kwargs["scope"],
            quotes={
                "_".join(
                    [
                        quote.quote_id.quote_series_id.instrument_id,
                        quote.quote_id.quote_series_id.instrument_id_type,
                        str(quote.quote_id.effective_at),
                    ]
                ): quote
                for quote in quote_batch
            },
        )

    @staticmethod
    @run_in_executor
    def load_transaction_batch(
        api_factory: lusid.utilities.ApiClientFactory, transaction_batch: list, **kwargs
    ) -> lusid.models.UpsertPortfolioTransactionsResponse:
        """
        Upserts a batch of transactions into LUSID

        :param lusid.utilities.ApiClientFactory api_factory: The api factory to use
        :param str scope: The scope of the Transaction Portfolio to upsert the transactions into
        :param str code: The code of the Transaction Portfolio, together with the scope this uniquely identifies the portfolio
        :param list[lusid.models.TransactionRequest] transaction_batch: The batch of transactions to upsert

        :return: lusid.models.UpsertPortfolioTransactionsResponse: The response from LUSID
        """

        if "scope" not in list(kwargs.keys()):
            raise KeyError(
                "You are trying to load transactions without a scope, please ensure that a scope is provided."
            )

        if "code" not in list(kwargs.keys()):
            raise KeyError(
                "You are trying to load transactions without a portfolio code, please ensure that a code is provided."
            )

        return api_factory.build(
            lusid.api.TransactionPortfoliosApi
        ).upsert_transactions(
            scope=kwargs["scope"], code=kwargs["code"], transactions=transaction_batch
        )

    @staticmethod
    @run_in_executor
    def load_holding_batch(
        api_factory: lusid.utilities.ApiClientFactory, holding_batch: list, **kwargs
    ) -> lusid.models.HoldingsAdjustment:
        """
        Upserts a batch of holdings into LUSID

        :param lusid.utilities.ApiClientFactory api_factory: The api factory to use
        :param list[lusid.models.AdjustHoldingRequest] holding_batch: The batch of holdings
        :param kwargs: 'scope', 'code', 'effective_at' The parameters required for the API call

        :return: lusid.models.HoldingsAdjustment: The response from LUSID
        """

        if "scope" not in list(kwargs.keys()):
            raise KeyError(
                "You are trying to load transactions without a scope, please ensure that a scope is provided."
            )

        if "code" not in list(kwargs.keys()):
            raise KeyError(
                "You are trying to load transactions without a portfolio code, please ensure that a code is provided."
            )

        if "effective_at" not in list(kwargs.keys()):
            raise KeyError(
                """There is no mapping for effective_at in the required mapping, please add it"""
            )

        return api_factory.build(lusid.api.TransactionPortfoliosApi).set_holdings(
            scope=kwargs["scope"],
            code=kwargs["code"],
            effective_at=str(DateOrCutLabel(kwargs["effective_at"])),
            holding_adjustments=holding_batch,
        )

    @staticmethod
    @run_in_executor
    def load_portfolio_batch(
        api_factory: lusid.utilities.ApiClientFactory, portfolio_batch: list, **kwargs
    ) -> lusid.models.Portfolio:
        """
        Upserts a batch of portfolios to LUSID

        :param lusid.utilities.ApiClientFactory api_factory: The api factory to use
        :param list[lusid.models.CreateTransactionPortfolioRequest] portfolio_batch: The batch of portfolios to create
        :param kwargs: 'scope', 'code' arguments required for the API call

        :return: lusid.models.Portfolio: The response from LUSID
        """

        if "scope" not in list(kwargs.keys()):
            raise KeyError(
                "You are trying to load transactions without a scope, please ensure that a scope is provided."
            )

        if "code" not in list(kwargs.keys()):
            raise KeyError(
                "You are trying to load transactions without a portfolio code, please ensure that a code is provided."
            )

        try:
            return api_factory.build(lusid.api.PortfoliosApi).get_portfolio(
                scope=kwargs["scope"], code=kwargs["code"]
            )
        except lusid.exceptions.ApiException as e:
            if e.status == 404:
                return api_factory.build(
                    lusid.api.TransactionPortfoliosApi
                ).create_portfolio(
                    scope=kwargs["scope"], transaction_portfolio=portfolio_batch[0]
                )
            # Add in here upsert portfolio properties if it does exist


async def load_data(
    api_factory: lusid.utilities.ApiClientFactory,
    data_frame: pd.DataFrame,
    mapping_required: dict,
    mapping_optional: dict,
    property_columns: list,
    properties_scope: str,
    instrument_identifier_mapping: dict,
    file_type: str,
    domain_lookup: dict,
    **kwargs,
):
    """
    This function populates the required models from a DataFrame and loads the data into LUSID

    :param lusid.utilities.ApiClientFactory api_factory: The api factory to use
    :param pd.DataFrame data_frame: The DataFrame containing the data to load
    :param dict mapping_required: The required mapping
    :param dict mapping_optional: The optional mapping
    :param list property_columns: The property columns to add as property values
    :param str properties_scope: The scope to add the property values in
    :param dict instrument_identifier_mapping: The mapping for the identifiers
    :param int batch_size: The batch size to use
    :param str file_type: The file type to load
    :param dict domain_lookup: The domain lookup
    :param kwargs: Arguments specific to each call e.g. effective_at for holdings

    :return: BatchLoader StaticMethod: A static method on BatchLoader
    """

    # Get the top level model used for this request e.g. lusid.models.InstrumentDefintion for upsert instruments
    top_level_model = getattr(lusid.models, domain_lookup[file_type]["top_level_model"])

    # Verify that all the required attributes for this top level model exist in the provided required mapping
    cocoon.utilities.verify_all_required_attributes_mapped(
        swagger_dict=cocoon.utilities.get_swagger_dict(
            api_url=api_factory.api_client.configuration.host
        ),
        mapping=mapping_required,
        model_object=top_level_model,
        exempt_attributes=["identifiers", "properties", "instrument_identifiers"],
    )

    # Get the data types of the columns to be added as properties
    property_dtypes = data_frame.loc[:, property_columns].dtypes

    unique_identifiers = kwargs["unique_identifiers"]

    # Iterate over the DataFrame creating the single requests
    single_requests = []
    for index, row in data_frame.iterrows():

        # Create the property values for this row
        if domain_lookup[file_type]["domain"] is None:
            properties = None
        else:
            properties = cocoon.properties.create_property_values(
                row=row,
                scope=properties_scope,
                domain=domain_lookup[file_type]["domain"],
                dtypes=property_dtypes,
            )

        # Create identifiers for this row if applicable
        # If no instrument identifier mapping is provided return None as no identifiers are required
        if instrument_identifier_mapping is None or not bool(
            instrument_identifier_mapping
        ):
            identifiers = None
        else:
            identifiers = cocoon.instruments.create_identifiers(
                index=index,
                row=row,
                file_type=file_type,
                instrument_identifier_mapping=instrument_identifier_mapping,
                unique_identifiers=unique_identifiers,
                full_key_format=kwargs["full_key_format"],
            )

        # Construct the from the mapping, properties and identifiers the single request object and add it to the list
        single_requests.append(
            cocoon.utilities.populate_model(
                model_object=top_level_model,
                required_mapping=mapping_required,
                optional_mapping=mapping_optional,
                row=row,
                properties=properties,
                identifiers=identifiers,
            )
        )

    # Dynamically call the correct async function to use based on the file type
    return await getattr(BatchLoader, f"load_{file_type}_batch")(
        api_factory,
        single_requests,
        # Any specific arguments e.g. 'code' for transactions, 'effective_at' for holdings is passed in via **kwargs
        **kwargs,
    )


async def construct_batches(
    api_factory: lusid.utilities.ApiClientFactory,
    data_frame: pd.DataFrame,
    mapping_required: dict,
    mapping_optional: dict,
    property_columns: list,
    properties_scope: str,
    instrument_identifier_mapping: dict,
    batch_size: int,
    file_type: str,
    domain_lookup: dict,
    **kwargs,
):

    """
    This constructs the batches and asynchronously sends them to be loaded into LUSID

    :param lusid.utilities.ApiClientFactory api_factory: The api factory to use
    :param pd.DataFrame data_frame: The DataFrame containing the data to load
    :param dict mapping_required: The required mapping
    :param dict mapping_optional: The optional mapping
    :param list property_columns: The property columns to add as property values
    :param str properties_scope: The scope to add the property values in
    :param dict instrument_identifier_mapping: The mapping for the identifiers
    :param int batch_size: The batch size to use
    :param str file_type: The file type to load
    :param dict domain_lookup: The domain lookup
    :param kwargs: Arguments specific to each call e.g. effective_at for holdings

    :return: dict: Contains the success responses and the errors (where an API exception has been raised)
    """

    # Get the different behaviours required for different entities e.g quotes can be batched without worrying about portfolios
    batching_no_portfolios = [
        file_type
        for file_type, settings in domain_lookup.items()
        if not settings["portfolio_specific"]
    ]
    batching_with_portfolios = [
        file_type
        for file_type, settings in domain_lookup.items()
        if settings["portfolio_specific"]
    ]

    if file_type in batching_no_portfolios:

        # Everything can be sent up asynchronously, prepare batches based on batch size alone
        async_batches = [
            data_frame.iloc[i : i + batch_size]
            for i in range(0, len(data_frame), batch_size)
        ]

        # Nest the async batches inside a single synchronous batch
        sync_batches = [
            {
                "async_batches": async_batches,
                "codes": [None] * len(async_batches),
                "effective_at": [None] * len(async_batches),
            }
        ]

    elif file_type in batching_with_portfolios:

        if "effective_at" in domain_lookup[file_type]["required_call_attributes"]:

            # Get unique effective dates
            unique_effective_dates = list(
                data_frame[mapping_required["effective_at"]].unique()
            )

            # Create a group for each effective date as they can not be batched asynchronously
            effective_at_groups = [
                data_frame.loc[
                    data_frame[mapping_required["effective_at"]] == effective_at
                ]
                for effective_at in unique_effective_dates
            ]

            # Create a synchronous batch for each effective date
            sync_batches = [
                {
                    # Different portfolio codes can be batched asynchronously inside the synchronous batch
                    "async_batches": [
                        effective_at_group.loc[
                            data_frame[mapping_required["code"]] == code
                        ]
                        for code in list(
                            effective_at_group[mapping_required["code"]].unique()
                        )
                    ],
                    "codes": list(
                        effective_at_group[mapping_required["code"]].unique()
                    ),
                    "effective_at": [
                        list(
                            effective_at_group[
                                mapping_required["effective_at"]
                            ].unique()
                        )[0]
                    ]
                    * len(list(effective_at_group[mapping_required["code"]].unique())),
                }
                for effective_at_group in effective_at_groups
            ]

        else:

            unique_portfolios = list(data_frame[mapping_required["code"]].unique())

            # Different portfolio codes can be batched asynchronously
            async_batches = [
                data_frame.loc[data_frame[mapping_required["code"]] == code]
                for code in unique_portfolios
            ]

            # Inside the synchronous batch split the values for each portfolio into appropriate batch sizes
            sync_batches = [
                {
                    "async_batches": [
                        async_batch.iloc[i : i + batch_size]
                        for async_batch in async_batches
                    ],
                    "codes": [str(code) for code in unique_portfolios],
                    "effective_at": [None] * len(async_batches),
                }
                for i in range(
                    0,
                    max([len(async_batch) for async_batch in async_batches]),
                    batch_size,
                )
            ]

    # Asynchronously load the data into LUSID
    responses = [
        await asyncio.gather(
            *[
                load_data(
                    api_factory=api_factory,
                    data_frame=async_batch,
                    mapping_required=mapping_required,
                    mapping_optional=mapping_optional,
                    property_columns=property_columns,
                    properties_scope=properties_scope,
                    instrument_identifier_mapping=instrument_identifier_mapping,
                    file_type=file_type,
                    domain_lookup=domain_lookup,
                    code=code,
                    effective_at=effective_at,
                    **kwargs,
                )
                for async_batch, code, effective_at in zip(
                    sync_batch["async_batches"],
                    sync_batch["codes"],
                    sync_batch["effective_at"],
                )
                if not async_batch.empty
            ],
            return_exceptions=True,
        )
        for sync_batch in sync_batches
    ]

    responses_flattened = [
        response for responses_sub in responses for response in responses_sub
    ]

    # Raise any internal exceptions rather than propagating them to the response
    for response in responses_flattened:
        if isinstance(response, Exception) and not isinstance(
            response, lusid.exceptions.ApiException
        ):
            raise response

    # Collects the exceptions as failures and successful calls as values
    return {
        "errors": [r for r in responses_flattened if isinstance(r, Exception)],
        "success": [r for r in responses_flattened if not isinstance(r, Exception)],
    }


@checkargs
def load_from_data_frame(
    api_factory: lusid.utilities.ApiClientFactory,
    scope: str,
    data_frame: pd.DataFrame,
    mapping_required: dict,
    mapping_optional: dict,
    file_type: str,
    identifier_mapping: dict = None,
    property_columns: list = None,
    properties_scope: str = None,
    batch_size: int = None,
    remove_white_space: bool = True,
    instrument_name_enrichment: bool = False,
):
    """
    Handles loading data from a DataFrame into LUSID

    :param lusid.utilities.ApiClientFactory api_factory: The api factory to use
    :param str scope: The scope of the resource to load the data into
    :param Pandas DataFrame data_frame: The DataFrame containing the data
    :param dict{str, str} mapping_required: The dictionary mapping the DataFrame columns to LUSID's required attributes
    :param dict{str, str} mapping_optional: The dictionary mapping the DataFrame columns to LUSID's optional attributes
    :param str file_type: The type of file e.g. transactions, instruments, holdings, quotes, portfolios
    :param dict{str, str} identifier_mapping: The dictionary mapping of LUSID instrument identifiers to identifiers in the DataFrame
    :param list[str] property_columns: The columns to create properties for
    :param str properties_scope: The scope to add the properties to
    :param int batch_size: The size of the batch to use when using upsert calls e.g. upsert instruments, upsert quotes etc.
    :param bool remove_white_space: remove whitespace either side of each value in the dataframe
    :param bool instrument_name_enrichment: request additional identifier information from open-figi
    :return: dict responses: The responses from loading the data into LUSID
    """
    # remove whitespace from dataframe
    if remove_white_space:
        data_frame = strip_whitespace(data_frame)

    # A mapping between the file type and relevant attributes e.g. domain, top_level_model etc.
    domain_lookup = cocoon.utilities.load_json_file("config/domain_settings.json")

    # Convert the file type to lower case & singular as well as checking it is of the allowed value
    file_type = (
        Validator(file_type, "file_type")
        .make_singular()
        .make_lower()
        .check_allowed_value(list(domain_lookup.keys()))
        .value
    )

    # Set defaults aligned with the data type of each argument, this allows for users to provide None
    identifier_mapping = (
        Validator(identifier_mapping, "identifier_mapping")
        .set_default_value_if_none(default={})
        .discard_dict_keys_none_value()
        .value
    )

    properties_scope = (
        Validator(properties_scope, "properties_scope")
        .set_default_value_if_none(default=scope)
        .value
    )

    property_columns = (
        Validator(property_columns, "property_columns")
        .set_default_value_if_none(default=[])
        .value
    )

    batch_size = (
        Validator(batch_size, "batch_size")
        .set_default_value_if_none(domain_lookup[file_type]["default_batch_size"])
        .override_value(
            not domain_lookup[file_type]["batch_allowed"],
            domain_lookup[file_type]["default_batch_size"],
        )
        .value
    )

    # Discard mappings where the provided value is None
    mapping_required = (
        Validator(mapping_required, "mapping_required")
        .discard_dict_keys_none_value()
        .value
    )

    mapping_optional = (
        Validator(mapping_optional, "mapping_optional")
        .discard_dict_keys_none_value()
        .value
    )

    required_call_attributes = domain_lookup[file_type]["required_call_attributes"]
    if "scope" in required_call_attributes:
        required_call_attributes.remove("scope")

    # Check that all required parameters exist
    Validator(
        required_call_attributes, "required_attributes_for_call"
    ).check_subset_of_list(list(mapping_required.keys()), "required_mapping")

    if instrument_name_enrichment:

        loop = cocoon.async_tools.start_event_loop_new_thread()

        data_frame, mapping_required = asyncio.run_coroutine_threadsafe(
            cocoon.instruments.enrich_instruments(
                api_factory=api_factory,
                data_frame=data_frame,
                instrument_identifier_mapping=identifier_mapping,
                mapping_required=mapping_required,
                constant_prefix="$",
            ),
            loop,
        ).result()

        # Stop the additional event loop
        loop.stop()

    """
    Unnest and populate defaults where a mapping is provided with column and/or default fields in a nested dictionary
    
    e.g.
    {'name': {
        'column': 'instrument_name',
        'default': 'unknown_name'
        }
    }
    
    rather than simply
    {'name': 'instrument_name'}
    """
    (
        data_frame,
        mapping_required,
    ) = cocoon.utilities.handle_nested_default_and_column_mapping(
        data_frame=data_frame, mapping=mapping_required, constant_prefix="$"
    )
    (
        data_frame,
        mapping_optional,
    ) = cocoon.utilities.handle_nested_default_and_column_mapping(
        data_frame=data_frame, mapping=mapping_optional, constant_prefix="$"
    )

    # Get all the DataFrame columns as well as those that contain at least one null value
    data_frame_columns = list(data_frame.columns.values)
    nan_columns = [
        column for column in data_frame_columns if data_frame[column].isna().any()
    ]

    # Validate that none of the provided columns are missing or invalid
    Validator(
        mapping_required, "mapping_required"
    ).get_dict_values().filter_list_using_first_character("$").check_subset_of_list(
        data_frame_columns, "DataFrame Columns"
    ).check_no_intersection_with_list(
        nan_columns, "Columns with Missing Values"
    )

    Validator(
        mapping_optional, "mapping_optional"
    ).get_dict_values().filter_list_using_first_character("$").check_subset_of_list(
        data_frame_columns, "DataFrame Columns"
    )

    Validator(
        identifier_mapping, "identifier_mapping"
    ).get_dict_values().filter_list_using_first_character("$").check_subset_of_list(
        data_frame_columns, "DataFrame Columns"
    )

    Validator(property_columns, "property_columns").check_subset_of_list(
        data_frame_columns, "DataFrame Columns"
    )

    # Converts higher level data types such as dictionaries and lists to strings
    data_frame = data_frame.applymap(cocoon.utilities.convert_cell_value_to_string)

    # Check for and create missing property defintions
    data_frame = cocoon.properties.create_missing_property_definitions_from_file(
        api_factory=api_factory,
        properties_scope=properties_scope,
        file_type=file_type,
        data_frame=data_frame,
        property_columns=property_columns,
        domain_lookup=domain_lookup,
    )

    # Start a new event loop in a new thread, this is required to run inside a Jupyter notebook
    loop = cocoon.async_tools.start_event_loop_new_thread()

    # Keyword arguments to be used in requests to the LUSID API
    keyword_arguments = {
        "scope": scope,
        # This handles that identifiers need to be specified differently based on the request type, allowing users
        # to provide either the entire key e.g. "Instrument/default/Figi" or just the code "Figi" for any request
        "full_key_format": domain_lookup[file_type]["full_key_format"],
        # Gets the allowed unique identifiers
        "unique_identifiers": cocoon.instruments.get_unique_identifiers(
            api_factory=api_factory
        ),
    }

    # Get the responses from LUSID
    responses = asyncio.run_coroutine_threadsafe(
        construct_batches(
            api_factory=api_factory,
            data_frame=data_frame,
            mapping_required=mapping_required,
            mapping_optional=mapping_optional,
            property_columns=property_columns,
            properties_scope=properties_scope,
            instrument_identifier_mapping=identifier_mapping,
            batch_size=batch_size,
            file_type=file_type,
            domain_lookup=domain_lookup,
            **keyword_arguments,
        ),
        loop,
    ).result()

    # Stop the additional event loop
    loop.stop()

    # Prefix the responses with the file type
    return {file_type + "s": responses}
