import logging
from datetime import datetime, timezone
from typing import Dict, List, Optional, Tuple
from warnings import warn

import pandas as pd
from cognite.client import CogniteClient
from cognite.client.data_classes import Asset, AssetHierarchy, AssetList
from deepdiff import DeepDiff
from rdflib import Graph, Namespace
from rdflib.term import URIRef

from cognite.neat.core.data_classes import AssetTemplate, Property
from cognite.neat.core.data_classes.config import EXCLUDE_PATHS
from cognite.neat.core.data_classes.transformation_rules import TransformationRules
from cognite.neat.core.loader.graph_store import NeatGraphStore
from cognite.neat.core.utils import chunker, datetime_utc_now, remove_namespace

ORPHANAGE = {
    "external_id": "orphanage",
    "name": "Orphanage",
    "parent_external_id": None,
    "description": "Used to store all assets which parent does not exist",
    "metadata": {
        "type": "Orphanage",
        "cdfResourceType": "Asset",
        "identifier": "orphanage",
        "active": "true",
        "start_time": "",
        "update_time": "",
    },
    "data_set_id": None,
    "labels": ["Orphanage", "non-historic"],
}


def _get_class_instance_ids(graph: Graph, class_: str, namespace: Namespace, limit: int = -1) -> List[URIRef]:
    """Get instances ids for a given class

    Parameters
    ----------
    graph : Graph
        Graph containing class instances
    class_ : str
        Class for which instances are to be found
    namespace : Namespace
        Namespace of given class (to avoid writing long URIs)
    limit : int
        Max number of instances to return, by default -1 meaning all instances

    Returns
    -------
    List[URIRef]
        List of class instance URIs
    """

    query_statement = "SELECT DISTINCT ?subject WHERE { ?subject a <class> .} LIMIT X".replace(
        "class", namespace[class_]
    ).replace("LIMIT X", "" if limit == -1 else f"LIMIT {limit}")
    logging.debug(f"Query statement: {query_statement}")
    return [res[0] for res in list(graph.query(query_statement))]


def _get_class_instance(graph: Graph, instance: URIRef) -> List[tuple]:
    """Get instance by means of tuples containing property-value pairs

    Returns
    -------
    List
        list property-value pairs
    """

    query_statement = "SELECT DISTINCT ?predicate ?object WHERE {<subject> ?predicate ?object .}".replace(
        "subject", instance
    )
    result = list(graph.query(query_statement))

    # Adds instance id for sake of keep chain of custody
    result += [(URIRef("http://purl.org/dc/terms/identifier"), instance)]

    return result


def _get_class_property_pairs(transformation_rules: TransformationRules) -> dict[List[Property]]:
    """Define classes in terms of their properties

    Parameters
    ----------
    transformation_rules : TransformationRules
        _description_

    Returns
    -------
    dict
        Dict containing keys as class names and list of their properties
    """

    classes = {}

    for property_ in transformation_rules.properties.keys():
        class_ = transformation_rules.properties[property_].class_name

        if class_ in classes:
            classes[class_] += [transformation_rules.properties[property_]]
        else:
            classes[class_] = [transformation_rules.properties[property_]]

    return classes


def _define_asset_class_mapping(transformation_rules: TransformationRules) -> dict:
    """Define mapping from class to asset properties

    Parameters
    ----------
    transformation_rules : TransformationRules
        Transformation rules

    Returns
    -------
    tuple[dict, dict[Any, dict[str, list[Any]]]]
        Dict containing mapping from class to asset properties
    """
    solution2cdf_mapping_rules = _get_class_property_pairs(transformation_rules)

    asset_class_mapping = {}

    for class_, properties in solution2cdf_mapping_rules.items():
        asset_class_mapping[class_] = {
            "external_id": [],
            "name": [],
            "description": [],
            "parent_external_id": [],
            "metadata": [],
        }

        for property_ in properties:
            if "Asset" in property_.cdf_resource_type and property_.property_name != "*":
                for resource_type_property in property_.resource_type_property:
                    if (
                        resource_type_property in asset_class_mapping[class_]
                        and property_.property_name not in asset_class_mapping[class_][resource_type_property]
                    ):
                        asset_class_mapping[class_][resource_type_property] += [property_.property_name]

                    if property_.property_name not in asset_class_mapping[class_]["metadata"]:
                        asset_class_mapping[class_]["metadata"] += [property_.property_name]

    return asset_class_mapping


def _remap_class_properties(class_instance: dict, asset_class_mapping: dict) -> Tuple[dict, set, set]:
    """Remaps original class instance properties to asset properties (e.g., external_id, name, description, metadata)

    Parameters
    ----------
    class_instance : dict
        Dictionary containing class instance properties and values originating from RDF
        stripped from namespaces
    asset_class_mapping : dict
        Property mapping from class to asset

    Returns
    -------
    Tuple[dict, set, set]
        Remapped class instance, set of missing asset properties and set of missing asset metadata
    """
    # Make distinction between missing properties that map into Asset fields
    # and missing RDF properties that are defined by sheet
    instance_properties = list(class_instance.keys())
    missing_properties = set()

    for property_group, ordered_properties in asset_class_mapping.items():
        if property_group != "metadata" and ordered_properties:
            if matching_property := next(
                (a for a in ordered_properties if a in instance_properties),
                None,
            ):
                class_instance[property_group] = class_instance[matching_property]
            else:
                missing_properties.add(property_group)

    missing_metadata = set(asset_class_mapping["metadata"]).difference(set(instance_properties))

    return class_instance, missing_properties, missing_metadata


def _class2asset_instance(
    class_: str,
    class_instance: dict,
    asset_class_mapping: dict,
    data_set_id: int,
    orphanage_asset_external_id: str = None,
    external_id_prefix: str = None,
    fallback_property: str = "identifier",
    empty_name_default: str = "Missing Name",
    add_missing_metadata: bool = True,
) -> dict:
    """Converts class instance to asset instance dictionary

    Parameters
    ----------
    class_ : str
        Class name which instance is being converted to asset instance
    class_instance : dict
        Dictionary containing class instance properties and values originating from RDF
        stripped from namespaces
    asset_class_mapping : dict
        Property mapping from class to asset
    data_set_id : int
        data set id to which asset belongs
    orphanage_asset_id : str, optional
        Orphanage asset external id, by default None
    external_id_prefix : str, optional
        External id prefix to be added to any external id, by default None
    fallback_property : str, optional
        Property from class instance to be used as fallback in case of missing properties, by default "identifier"


    Returns
    -------
    dict
        Asset instance dictionary
    """

    remapped_class_instance, missing_properties, missing_metadata = _remap_class_properties(
        class_instance, asset_class_mapping
    )

    # setting class instance type to class name
    remapped_class_instance["type"] = class_
    # This will be a default case since we want to use original identifier as external_id
    # We are though dropping namespace from the original identifier (avoiding long-tail URIs)
    if "external_id" in missing_properties or asset_class_mapping["external_id"] == []:
        try:
            __extracted_from___class2asset_instance_49(
                remapped_class_instance,
                fallback_property,
                "external_id",
                class_,
            )
        except Exception:
            __extracted_from___class2asset_instance_56(fallback_property, class_, remapped_class_instance)
    # This should not be the use case however to still have name of the object we are using
    # fallback property here as well (typically identifier)
    if "name" in missing_properties:
        try:
            __extracted_from___class2asset_instance_49(remapped_class_instance, fallback_property, "name", class_)
        except Exception:
            __extracted_from___class2asset_instance_56(fallback_property, class_, remapped_class_instance)
    # If object is expected to have parent, but parent is not provided, it is added to orphanage
    # This is typically sign of objects not following proposed ontology/data model/schema
    if "parent_external_id" in missing_properties and orphanage_asset_external_id:
        remapped_class_instance["parent_external_id"] = orphanage_asset_external_id

    if "name" in remapped_class_instance and remapped_class_instance["name"] == "":
        remapped_class_instance["name"] = empty_name_default

    # To maintain shape across of all assets of specific type we are adding missing metadata
    # keys as empty strings, this was request by Statnett
    # Generally this is bad practice, but more of a workaround of their bad data
    if missing_metadata and add_missing_metadata:
        msg = f"Adding missing metadata keys with values set to empty string for {class_}"
        msg += f" instance <{remapped_class_instance['identifier']}>. "
        logging.debug(msg)
        for key in missing_metadata:
            remapped_class_instance[key] = ""
            logging.debug(f"\tKey {key} added to <{remapped_class_instance['identifier']}> metadata!")

    asset_instance = AssetTemplate(
        **remapped_class_instance, external_id_prefix=external_id_prefix, data_set_id=data_set_id
    )

    # Removing field external_id_prefix from asset instance dictionary as it is only
    # convenience field for external_id and parent_external_id update in AssetTemplate
    return asset_instance.dict(exclude={"external_id_prefix"})


# TODO Rename this here and in `__class2asset_instance`
def __extracted_from___class2asset_instance_49(remapped_class_instance, fallback_property, arg2, class_):
    remapped_class_instance[arg2] = remapped_class_instance[fallback_property]
    msg = f"Missing external_id for {class_} instance <{remapped_class_instance['identifier']}>. "
    msg += f"Using value <{remapped_class_instance[fallback_property]}> provided "
    msg += f"by property <{fallback_property}>!"

    logging.debug(msg)


# TODO Rename this here and in `__class2asset_instance`
def __extracted_from___class2asset_instance_56(fallback_property, class_, remapped_class_instance):
    msg = f"Fallback property <{fallback_property}> not found for {class_} "
    msg += f"instance <{remapped_class_instance['identifier']}>."
    logging.error(msg)
    raise ValueError(msg)


def _list2dict(class_instance: list) -> dict:
    """Converting list of class instance properties and values to dictionary

    Parameters
    ----------
    class_instance : list
        Class instance properties and values originating from RDF as list of tuples

    Returns
    -------
    dict
        Class instance properties and values as dictionary
    """

    class_instance_dict = {}
    for property_value_pair in class_instance:
        property_ = remove_namespace(property_value_pair[0])

        # Remove namespace from URIRef values, otherwise convert Literal to string
        # ideally this should react upon property type provided in sheet
        # however Assets only support string values
        value = (
            remove_namespace(property_value_pair[1])
            if isinstance(property_value_pair[1], URIRef)
            else str(property_value_pair[1])
        )

        if property_ in class_instance_dict and value not in class_instance_dict[property_]:
            class_instance_dict[property_] = (
                class_instance_dict[property_] + [value]
                if isinstance(class_instance_dict[property_], list)
                else [class_instance_dict[property_], value]
            )
        else:
            class_instance_dict[property_] = value

    return class_instance_dict


def rdf2assets(
    graph_store: NeatGraphStore,
    transformation_rules: TransformationRules,
    stop_on_exception: bool = False,
    use_orphanage: bool = True,
) -> Dict[str, Asset]:
    """Creates assets from RDF graph

    Parameters
    ----------
    graph : Graph
        Graph containing RDF data
    transformation_rules : TransformationRules
        Instance of TransformationRules class containing transformation rules

    Returns
    -------
    Dict[str, Asset]
        Dictionary of assets with external_id as key
    """

    orphanage_asset_external_id = (
        f"{transformation_rules.metadata.externalIdPrefix}orphanage"
        if transformation_rules.metadata.externalIdPrefix
        else "orphanage"
    )

    graph = graph_store.get_graph()
    # Step 1: Create rdf to asset property mapping
    logging.info("Generating rdf to asset property mapping")
    asset_class_mapping = _define_asset_class_mapping(transformation_rules)

    # Step 4: Get ids of classes
    logging.info("Get ids of instances of classes")
    assets = {}
    class_ids = {
        class_: _get_class_instance_ids(graph, class_, transformation_rules.metadata.namespace)
        for class_ in asset_class_mapping
    }
    # Step 5: Create Assets based on class instances
    logging.info("Create Assets based on class instances")
    for class_ in asset_class_mapping:
        # TODO: Rename class_id to instance_id
        class_ns = transformation_rules.metadata.namespace[class_]
        logging.debug(f"Processing class <{class_ns}> . Number of instances: {len(class_ids[class_])}")
        progress_counter = 0
        # loading all instances into cache
        try:
            query = f"SELECT ?instance ?prop ?value WHERE {{ ?instance rdf:type <{class_ns}> . ?instance ?prop ?value . }} order by ?instance "
            logging.info(query)
            response_df = graph_store.query_to_dataframe(query)
        except Exception as e:
            logging.error(f"Error while loading instances of class <{class_ns}> into cache. Reason: {e}")
            # traceback.print_exc()
            if stop_on_exception:
                raise e
            continue

        grouped_df = response_df.groupby("instance")

        for instance_id, group_df in grouped_df:
            try:
                instance_property_values = group_df.filter(items=["property", "value"]).values.tolist()
                instance_property_values += [(URIRef("http://purl.org/dc/terms/identifier"), URIRef(instance_id))]

                # this will strip namespace from property names and values
                class_instance = _list2dict(instance_property_values)

                # class instance is repaired and converted to asset dictionary
                asset = _class2asset_instance(
                    class_,
                    class_instance,
                    asset_class_mapping[class_],
                    transformation_rules.metadata.data_set_id,
                    "orphanage" if use_orphanage else None,  # we need only base external id
                    transformation_rules.metadata.externalIdPrefix or None,
                )

                # adding labels and timestamps
                asset["labels"] = [asset["metadata"]["type"], "non-historic"]
                asset["metadata"]["start_time"] = str(datetime.now(timezone.utc))
                asset["metadata"]["update_time"] = str(datetime.now(timezone.utc))

                # log every 10000 assets
                if progress_counter % 10000 == 0:
                    logging.info(" Next 10000 Assets processed")

                assets[asset["external_id"]] = asset
                progress_counter += 1
            except Exception as ValidationError:
                logging.error(
                    f"Skipping class <{class_}> instance <{remove_namespace(instance_id)}>, reason:\n{ValidationError}\n"
                )
                if stop_on_exception:
                    raise ValidationError

        logging.debug(f"Class <{class_}> processed")

    if orphanage_asset_external_id not in assets:
        _extracted_from_rdf2asset_dictionary_95(orphanage_asset_external_id, transformation_rules, assets)
    logging.info("Assets dictionary created")
    return assets


def rdf2asset_dictionary(
    graph_store: NeatGraphStore,
    transformation_rules: TransformationRules,
    stop_on_exception: bool = False,
    use_orphanage: bool = True,
) -> Dict[str, Asset]:
    warn("'rdf2asset_dictionary' is deprecated, please use 'rdf2assets' instead!")
    logging.warning("'rdf2asset_dictionary' is deprecated, please use 'rdf2assets' instead!")
    return rdf2assets(graph_store, transformation_rules, stop_on_exception, use_orphanage)


# TODO Rename this here and in `rdf2asset_dictionary`
def _extracted_from_rdf2asset_dictionary_95(orphanage_asset_external_id, transformation_rules, assets):
    print(f"Orphanage with external id {orphanage_asset_external_id} not found in asset hierarchy!")
    logging.warning(f"Orphanage with external id {orphanage_asset_external_id} not found in asset hierarchy!")
    logging.warning(f"Adding default orphanage with external id {orphanage_asset_external_id}")

    # updating default external id for orphanage in case there are additional prefix to be added
    ORPHANAGE["external_id"] = orphanage_asset_external_id
    ORPHANAGE["data_set_id"] = transformation_rules.metadata.data_set_id
    ORPHANAGE["metadata"]["start_time"] = str(datetime_utc_now())
    ORPHANAGE["metadata"]["update_time"] = str(datetime_utc_now())
    assets[orphanage_asset_external_id] = ORPHANAGE


def _asset2dict(asset: Asset) -> dict:
    """Return asset as dict representation

    Parameters
    ----------
    asset : Asset
        Instance of Asset class

    Returns
    -------
    dict
        Asset in dict representation
    """

    return {
        "external_id": asset.external_id,
        "name": asset.name,
        "description": asset.description,
        "parent_external_id": asset.parent_external_id,
        "data_set_id": asset.data_set_id,
        "metadata": asset.metadata,
    }


def _flatten_labels(labels: List[Dict[str, str]]) -> set[str]:
    """Flatten labels"""
    return {label["externalId"] for label in labels}


def _is_historic(labels) -> bool:
    """Check if asset is historic"""
    return "historic" in labels


def _categorize_cdf_assets(
    client: CogniteClient, data_set_id: int, partitions: int
) -> Tuple[Optional[pd.DataFrame], dict[str, set]]:
    """Categorize CDF assets

    Parameters
    ----------
    client : CogniteClient
        Instance of CogniteClient
    data_set_id : int
        Id of data set
    partitions : int
        Number of partitions

    Returns
    -------
    Tuple[Optional[pd.DataFrame], dict]
        CDF assets as pandas dataframe and dictionary with categorized assets
    """
    cdf_assets = client.assets.list(data_set_ids=data_set_id, limit=None, partitions=partitions).to_pandas()

    logging.info(f"Number of assets in CDF {len(cdf_assets)} that have been fetched")

    if cdf_assets.empty:
        return None, {"non-historic": set(), "historic": set()}

    cdf_columns = set(cdf_assets.columns)
    expected_columns = {"external_id", "labels", "parent_external_id", "data_set_id", "name", "description", "metadata"}

    cdf_assets = cdf_assets[list(expected_columns.intersection(cdf_columns))]
    cdf_assets = cdf_assets.where(pd.notnull(cdf_assets), None)
    cdf_assets.labels = cdf_assets.labels.apply(_flatten_labels).values  # type: ignore
    cdf_assets["is_historic"] = cdf_assets.labels.apply(_is_historic).values

    categorized_asset_ids = {
        "historic": set(cdf_assets[cdf_assets.is_historic].external_id.values),
        "non-historic": set(cdf_assets[~cdf_assets.is_historic].external_id.values),
    }

    cdf_assets.drop(["is_historic"], axis=1, inplace=True)
    msg = f"CDF assets categorized into {len(categorized_asset_ids['historic'])} historic"
    msg += f" and {len(categorized_asset_ids['non-historic'])} non-historic assets"
    logging.info(msg)

    return cdf_assets, categorized_asset_ids


def order_assets(assets: dict[str, dict]) -> list[Asset]:
    """Order assets in a way that parent assets are created before child assets

    Parameters
    ----------
    assets : dict[str, dict]
        List of assets to be created

    Returns
    -------
    list[Asset]
        Ordered list of assets
    """
    hierarchy = AssetHierarchy([Asset(**asset) for asset in assets.values()], ignore_orphans=True)
    insert_dct = hierarchy.groupby_parent_xid()
    subtree_count = hierarchy.count_subtree(insert_dct)

    hierarchy = None

    asset_creation_order = pd.DataFrame.from_dict(subtree_count, orient="index", columns=["order"]).sort_values(
        by="order", ascending=False
    )
    asset_creation_order["external_id"] = asset_creation_order.index

    hierarchy = AssetList([Asset(**asset) for asset in assets.values()]).to_pandas()
    hierarchy = hierarchy.where(pd.notnull(hierarchy), None)
    hierarchy = hierarchy.merge(asset_creation_order, left_on="external_id", right_on="external_id")
    hierarchy = hierarchy.sort_values(by="order", ascending=False)
    hierarchy.reset_index(drop=True, inplace=True)
    hierarchy.labels = hierarchy.labels.apply(_flatten_labels)
    hierarchy.drop(["order"], axis=1, inplace=True)

    return [Asset(**row.to_dict()) for _, row in hierarchy.iterrows()]


def _assets_to_create(rdf_assets: dict, asset_ids: set) -> list[Asset]:
    """Return list of assets to be created

    Parameters
    ----------
    rdf_assets : dict
        Dictionary containing assets derived from knowledge graph (RDF)
    asset_ids : set
        Set of asset ids to be created

    Returns
    -------
    list[Asset]
        Ordered list of assets to be created
    """
    start_time = datetime_utc_now()
    if asset_ids:
        logging.info("Wrangling assets to be created into their final form")
        ordered_assets = order_assets({external_id: rdf_assets[external_id] for external_id in asset_ids})

        logging.info(f"Wrangling completed in {(datetime_utc_now() - start_time).seconds} seconds")
        return ordered_assets
    return []


def _assets_to_update(
    rdf_assets: dict, cdf_assets: pd.DataFrame, asset_ids: set, exclude_paths: list = EXCLUDE_PATHS
) -> list[Asset]:
    """Return list of assets to be updated

    Parameters
    ----------
    rdf_assets : dict
        Dictionary containing assets derived from knowledge graph (RDF)
    cdf_assets : pd.DataFrame
        Dataframe containing assets from CDF
    asset_ids : set
        Candidate assets to be updated
    exclude_paths : list, optional
        Paths not to be checked when diffing rdf and cdf assets, by default EXCLUDE_PATHS

    Returns
    -------
    list[Asset]
        List of assets to be updated
    """

    start_time = datetime_utc_now()
    assets = []
    if not asset_ids:
        return []
    logging.info("Wrangling assets to be updated into their final form")
    cdf_asset_subset = {
        row["external_id"]: row
        for row in cdf_assets[cdf_assets["external_id"].isin(asset_ids)].to_dict(orient="records")
    }
    for external_id in asset_ids:
        cdf_asset = cdf_asset_subset[external_id]
        diffing_result = DeepDiff(cdf_asset, rdf_assets[external_id], exclude_paths=exclude_paths)

        if diffing_result and "root['metadata']['active']" not in diffing_result.affected_paths:
            asset = Asset(**rdf_assets[external_id])
            try:
                asset.metadata["start_time"] = cdf_asset[external_id]["metadata"]["start_time"]
            except KeyError:
                asset.metadata["start_time"] = str(datetime.now(timezone.utc))
            asset.metadata["update_time"] = str(datetime.now(timezone.utc))
            assets.append(asset)

    logging.info(f"Wrangling of {len(assets)} completed in {(datetime_utc_now() - start_time).seconds} seconds")
    return assets


def _assets_to_resurrect(rdf_assets: dict, cdf_assets: pd.DataFrame, asset_ids: set) -> list[Asset]:
    """Returns list of assets to be resurrected

    Parameters
    ----------
    rdf_assets : dict
        Dictionary containing assets derived from knowledge graph (RDF)
    cdf_assets : pd.DataFrame
        Dataframe containing assets from CDF
    asset_ids : set
        Set of asset ids to be resurrected

    Returns
    -------
    list[Asset]
        List of assets to be resurrected
    """
    start_time = datetime_utc_now()
    assets = []
    if not asset_ids:
        return []
    logging.info("Wrangling assets to be resurrected into their final form")
    cdf_asset_subset = {
        row["external_id"]: row
        for row in cdf_assets[cdf_assets["external_id"].isin(asset_ids)].to_dict(orient="records")
    }
    for external_id in asset_ids:
        cdf_asset = cdf_asset_subset[external_id]

        asset = Asset(**rdf_assets[external_id])

        try:
            asset.metadata["start_time"] = cdf_asset[external_id]["metadata"]["start_time"]
        except KeyError:
            asset.metadata["start_time"] = str(datetime.now(timezone.utc))
        asset.metadata["update_time"] = str(datetime.now(timezone.utc))
        asset.metadata["resurrection_time"] = str(datetime.now(timezone.utc))
        assets.append(asset)

    logging.info(f"Wrangling of {len(assets)} completed in {(datetime_utc_now() - start_time).seconds} seconds")
    return assets


def _assets_to_decommission(cdf_assets, asset_ids) -> list[Asset]:
    start_time = datetime_utc_now()

    assets = []
    if not asset_ids:
        return []
    logging.info("Wrangling assets to be decommissioned into their final form")
    cdf_asset_subset = {
        row["external_id"]: row
        for row in cdf_assets[cdf_assets["external_id"].isin(asset_ids)].to_dict(orient="records")
    }

    for external_id in asset_ids:
        cdf_asset = cdf_asset_subset[external_id]

        cdf_asset["metadata"]["update_time"] = str(datetime.now(timezone.utc))
        cdf_asset["metadata"].pop("resurrection_time", None)
        cdf_asset["metadata"]["end_time"] = str(datetime.now(timezone.utc))
        cdf_asset["metadata"]["active"] = "false"
        try:
            cdf_asset["labels"].remove("non-historic")
        except KeyError:
            logging.info(f"Asset {external_id} missed label 'non-historic'")
        cdf_asset["labels"].add("historic")

        assets.append(Asset(**cdf_asset))

    logging.info(f"Wrangling completed in {(datetime_utc_now() - start_time).seconds} seconds")
    return assets


def categorize_assets(
    client: CogniteClient, rdf_assets: dict, data_set_id: int, partitions: int = 40
) -> dict[str, list[Asset]]:
    """Categorize assets on those that are to be created, updated and decommissioned

    Parameters
    ----------
    client : CogniteClient
        Instance of CogniteClient
    rdf_assets : dict
        Dictionary containing asset external_id - asset pairs
    data_set_id : int
        Dataset id to which assets are to be/are stored
    partitions : int, optional
        Number of partitions to use when fetching assets from CDF, by default 40

    Returns
    -------
    Dict[str, list]
        dictionary containing asset category - list of asset pairs
    """
    # TODO: Cache categorized assets somewhere instead of creating them

    cdf_assets, categorized_asset_ids = _categorize_cdf_assets(client, data_set_id, partitions)

    rdf_asset_ids = set(rdf_assets.keys())

    # ids to create
    create_ids = rdf_asset_ids.difference(
        categorized_asset_ids["historic"].union(categorized_asset_ids["non-historic"])
    )

    # ids potentially to update
    update_ids = rdf_asset_ids.intersection(categorized_asset_ids["non-historic"])

    # ids to decommission
    decommission_ids = categorized_asset_ids["non-historic"].difference(rdf_asset_ids)

    # ids to resurrect
    resurrect_ids = categorized_asset_ids["historic"].intersection(rdf_asset_ids)

    logging.info(f"Number of assets to create: { len(create_ids)}")
    logging.info(f"Number of assets to potentially update: { len(update_ids)}")
    logging.info(f"Number of assets to decommission: { len(decommission_ids)}")
    logging.info(f"Number of assets to resurrect: { len(resurrect_ids)}")

    return {
        "create": _assets_to_create(rdf_assets, create_ids),
        "update": _assets_to_update(rdf_assets, cdf_assets, update_ids),
        "resurrect": _assets_to_resurrect(rdf_assets, cdf_assets, resurrect_ids),
        "decommission": _assets_to_decommission(cdf_assets, decommission_ids),
    }


def _micro_batch_push(
    client: CogniteClient, assets: list, batch_size: int = 1000, push_type: str = "update", message: str = "Updated"
):
    """Updates assets in batches of 1000

    Parameters
    ----------
    client : CogniteClient
        Instance of CogniteClient
    assets : list
        List of assets to be created or updated
    batch_size : int, optional
        Size of batch, by default 1000
    push_type : str, optional
        Type of push, either "update" or "create", by default "update"
    message : str, optional
        Message to logged, by default "Updated"
    """
    total = len(assets)
    counter = 0

    for batch in chunker(assets, batch_size):
        counter += len(batch)
        start_time = datetime_utc_now()
        if push_type == "update":
            client.assets.update(batch)
        elif push_type == "create":
            client.assets.create_hierarchy(batch)
        else:
            logging.info(f"push_type {push_type} not supported")
            raise ValueError(f"push_type {push_type} not supported")

        delta_time = (datetime_utc_now() - start_time).seconds

        msg = f"{message} {counter} of {total} assets, batch processing time: {delta_time:.2f} "
        msg += f"seconds ETC: {delta_time * (total - counter) / (60*batch_size) :.2f} minutes"
        logging.info(msg)


def upload_assets(client: CogniteClient, categorized_assets: Dict[str, list], batch_size: int = 5000):
    """Uploads categorized assets to CDF

    Parameters
    ----------
    client : CogniteClient
        Instance of CogniteClient
    categorized_assets : Dict[str, list]
        dictionary containing asset category - list of asset pairs
    batch_size : int, optional
        Size of batch, by default 5000
    """
    if batch_size:
        logging.info(f"Uploading assets in batches of {batch_size}")
        if categorized_assets["create"]:
            _micro_batch_push(client, categorized_assets["create"], batch_size, push_type="create", message="Created")

        if categorized_assets["update"]:
            _micro_batch_push(
                client,
                categorized_assets["update"],
                batch_size,
                message="Updated",
            )

        if categorized_assets["resurrect"]:
            _micro_batch_push(
                client,
                categorized_assets["resurrect"],
                batch_size,
                message="Resurrected",
            )

        if categorized_assets["decommission"]:
            _micro_batch_push(
                client,
                categorized_assets["decommission"],
                batch_size,
                message="Decommissioned",
            )

    else:
        logging.info("Batch size not set, pushing all assets to CDF in one go!")
        if categorized_assets["create"]:
            client.assets.create_hierarchy(categorized_assets["create"])

        if categorized_assets["update"]:
            client.assets.create_hierarchy(categorized_assets["update"], upsert=True, upsert_mode="replace")

        if categorized_assets["resurrect"]:
            client.assets.create_hierarchy(categorized_assets["resurrect"], upsert=True, upsert_mode="replace")

        if categorized_assets["decommission"]:
            client.assets.create_hierarchy(categorized_assets["decommission"], upsert=True, upsert_mode="replace")
