# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/utils/upload_data.ipynb.

# %% auto 0
__all__ = ['upload_data', 'upload_data_with_date']

# %% ../../nbs/utils/upload_data.ipynb 2
import httpx
import pandas as pd
import asyncio

import domolibrary.client.Logger as lc
import domolibrary.client.DomoAuth as dmda
import domolibrary.classes.DomoDataset as dmds


# %% ../../nbs/utils/upload_data.ipynb 3
async def loop_upload(
    upload_df: pd.DataFrame,
    consol_ds: dmds.DomoDataset,
    partition_key: str,
    upload_method: str,
    logger: lc.Logger,
    debug_api: bool = False,
    debug_prn: bool = False,
    debug_fn: bool = True,
    max_retry: int = 2,
    is_index: bool = False
):
    base_msg = f"{partition_key} in {consol_ds.auth.domo_instance}" if partition_key else f"in {consol_ds.auth.domo_instance}"

    if debug_fn:
        print(
            f"starting upload of {len(upload_df)} rows to {base_msg} with {max_retry} attempts")

    retry_attempt = 0

    while retry_attempt <= max_retry:
        try:
            retry_attempt += 1

            if debug_fn:
                print(f"attempt {retry_attempt} for {base_msg}")

            return await consol_ds.upload_data(
                upload_df=upload_df,
                upload_method="REPLACE" if partition_key else upload_method,
                partition_key=partition_key,
                is_index=is_index,
                debug_api=debug_api,
                debug_prn=debug_prn,
            )

        except Exception as e:
            print(e)

            if retry_attempt <= max_retry:
                message = f"⚠️ upload_data : unexpected error: {e} in {partition_key} during retry_attempt {retry_attempt}"

                logger.log_warning(message)
            else:
                message = f"💀⚠️ upload_data : failed to upload {partition_key} to {consol_ds.auth.domo_instance}"
                logger.log_error(message)
            
            print(message)


# %% ../../nbs/utils/upload_data.ipynb 4
async def upload_data(
    # instance where the data_fn function will execute against
    data_fn,  # data function to execute
    instance_auth: dmda.DomoAuth,  # instance to run the data function against
    consol_ds: dmds.DomoDataset,  # dataset where data should be accumulated
    # if partition key supplied, will replace existing partition
    partition_key: str = None,
    upload_method: str = 'REPLACE',
    is_index: bool = False,  # index dataset
    debug_prn: bool = False,
    debug_fn: bool = True,
    debug_api: bool = False,
    logger: lc.Logger = None,
    max_retry: int = 2  # number of times to attempt upload
):
    logger = logger or lc.Logger(app_name="upload_data")

    try:
        message = f"🏁 starting {instance_auth.domo_instance} - {data_fn.__name__}"

        if debug_fn:
            print(message)
        logger.log_info(message)

        instance_session = httpx.AsyncClient()

        upload_df = await data_fn(instance_auth, instance_session, debug_api=debug_api)

        if upload_df is None or len(upload_df.index) == 0:
            print("no data to upload")
            return None

        if debug_fn:
            print(upload_df.head())

        res = await loop_upload(
            upload_df=upload_df,
            consol_ds=consol_ds,
            partition_key=partition_key,
            upload_method=upload_method,
            debug_api=debug_api,
            debug_prn=debug_prn,
            debug_fn=debug_fn,
            max_retry=max_retry,
            logger=logger,
            is_index=False
        )

        if res.is_success:
                message = f"🚀 success upload of {partition_key} to {consol_ds.id} in {consol_ds.auth.domo_instance} in {data_fn.__name__}"

        else:
            message = f"💣 upload_data successful status but failed to upload {partition_key} - {res.status} - {res.response} in {data_fn.__name__}"
            
        if debug_fn:
            print(message)
        
        logger.log_info(message)
        
        return res

    finally:
        if is_index:

            res = await consol_ds.index_dataset(debug_api=debug_api, session=instance_session)
            if res.is_success:
                message = f"🥫 successfully indexed {consol_ds.name} in {consol_ds.auth.domo_instance}"
            else:
                message = f"💀⚠️ failure to index {consol_ds.name} in {consol_ds.auth.domo_instance}"
            if debug_fn:
                print(message)
            logger.log_info(message)
        await instance_session.aclose()


# %% ../../nbs/utils/upload_data.ipynb 7
async def upload_data_with_date(
    instance_auth,
    consol_auth,
    data_fn,
    consol_ds,
    partition_date_col,
    partition_delimiter,
    start_date,
    end_date,
    debug: bool = False,
    debug_prn: bool = False,
):

    instance_session = httpx.AsyncClient()

    print(
        f"'🎬 upload_with_data: starting retrieval {start_date}, {end_date}, {instance_auth.domo_instance}"
    )

    upload_df = await data_fn(
        instance_auth=instance_auth,
        session=instance_session,
        start_date=start_date,
        end_date=end_date,
        debug=debug,
    )

    await instance_session.aclose()

    if not isinstance(upload_df, pd.DataFrame):
        print(f"🛑 error no data returned {instance_auth.domo_instance}")
        print(upload_df)
        return None

    if debug_prn:
        print(
            f"🧻 upload_with_data: starting upload {len(upload_df)} rows for {instance_auth.domo_instance}"
        )

    task = []

    for index, partition_set in upload_df.drop_duplicates(
        subset=[partition_date_col]
    ).iterrows():
        partition_date = partition_set[partition_date_col]

        partition_key = (
            f"{instance_auth.domo_instance}{partition_delimiter}{str(partition_date)}"
        )

        task.append(
            consol_ds.upload_data(
                upload_df=upload_df[(upload_df[partition_date_col] == partition_date)],
                upload_method="REPLACE",
                partition_key=partition_key,
                is_index=False,
                debug_api=debug_api,
                debug_prn=debug_prn,
            )
        )

    res = await asyncio.gather(*task)

    if debug_prn:
        print(
            f"🎉 upload_with_data : finished uploading {len(upload_df)} rows for {instance_auth.domo_instance}"
        )
    return res
