# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/integrations/Automation.ipynb.

# %% auto 0
__all__ = ['search_domo_groups_by_name', 'upsert_domo_group', 'search_or_upsert_domo_group', 'DJW_NoAccount',
           'search_domo_account_by_name', 'share_domo_account_with_domo_group', 'remove_partition_by_x_days',
           'get_company_domains']

# %% ../../nbs/integrations/Automation.ipynb 2
from typing import List

import domolibrary.client.DomoAuth as dmda
import domolibrary.client.DomoError as dmde
import domolibrary.classes.DomoAccount as dmacc

import domolibrary.classes.DomoGroup as dmdg

import domolibrary.utils.chunk_execution as dmce

import asyncio
import datetime as dt
import pandas as pd

import domolibrary.classes.DomoDataset as dmds

# %% ../../nbs/integrations/Automation.ipynb 5
async def search_domo_groups_by_name(
    auth: dmda.DomoAuth, group_names: List[str], is_hide_system_groups: bool = True
) -> List[dmdg.DomoGroup]:

    domo_groups = dmdg.DomoGroups(auth=auth)

    await domo_groups.get(is_hide_system_groups=is_hide_system_groups)

    return [dg for dg in domo_groups.groups if dg.name.lower() in group_names]


async def upsert_domo_group(
    auth: dmda.DomoAuth,
    group_name: str,
    description: str = f"updated via {dt.date.today()}",
    group_owner_names: List[str] = None,  # ["Role: Admin"]
    group_type: dmdg.GroupType_Enum = dmdg.GroupType_Enum["CLOSED"].value,
    is_hide_system_groups: bool = True,
    debug_api: bool = False,
) -> dmdg.DomoGroup:

    group_owner_names = group_owner_names or ["Role: Admin"]

    domo_group = await dmdg.DomoGroup.upsert(
        group_name=group_name,
        group_type=group_type,
        description=description,
        auth=auth,
        debug_api=debug_api,
    )

    domo_group_owners = await search_domo_groups_by_name(
        auth=auth,
        group_names=group_owner_names,
        is_hide_system_groups=is_hide_system_groups,
    )

    await domo_group.Membership.add_owners(add_owner_ls=domo_group_owners)

    return domo_group


async def search_or_upsert_domo_group(
    auth,
    group_name,
    upsert_if_not_exist: bool = True,
    group_type=dmdg.GroupType_Enum["CLOSED"].value,
    debug_api: bool = False,
    is_hide_system_groups: bool = True,
) -> dmdg.DomoGroup:
    try:
        domo_groups = await search_domo_groups_by_name(
            group_names=[group_name],
            auth=auth,
            is_hide_system_groups=is_hide_system_groups,
        )
        return domo_groups[0]

    except dmde.DomoError as e:
        if upsert_if_not_exist:
            return await upsert_domo_group(
                auth=auth,
                group_name=group_name,
                group_type=group_type,
                description=f"updated via {dt.date.today()}",
                debug_api=debug_api,
            )
        raise e from e

# %% ../../nbs/integrations/Automation.ipynb 7
class DJW_NoAccount(dmde.ClassError):
    def __init__(self, account_name, domo_instance):
        super().__init__(
            f"unable to retrieve account - {account_name} from {domo_instance}"
        )


async def search_domo_account_by_name(
    auth: dmda.DomoAuth, account_name: str
) -> dmacc.DomoAccount:


    domo_accounts = dmacc.DomoAccounts(auth=auth)

    await domo_accounts.get()

    domo_account = next(
        (
            domo_account
            for domo_account in domo_accounts.accounts
            if domo_account.name == account_name
        ),
        None,
    )

    if not domo_account:
        raise DJW_NoAccount(account_name, auth.domo_instance)

    return domo_account


async def share_domo_account_with_domo_group(
    auth: dmda.DomoAuth,
    account_name: str,
    group_name: str,
    upsert_group_if_no_exist: bool = True,
    is_hide_system_groups: bool = True,
    debug_api: bool = False,
    access_level=dmacc.ShareAccount_AccessLevel("default"),
) -> str:

    share_domo_group = await search_or_upsert_domo_group(
        auth=auth,
        group_name=group_name,
        upsert_if_not_exist=upsert_group_if_no_exist,
        is_hide_system_groups=is_hide_system_groups,
    )

    domo_account = await search_domo_account_by_name(
        auth=auth,
        account_name=account_name,
    )

    return await domo_account.Access.share(
        group_id=share_domo_group.id,
        debug_api=debug_api,
        access_level=access_level,
    )

# %% ../../nbs/integrations/Automation.ipynb 9
async def remove_partition_by_x_days(
    auth: dmda.DomoAuth,
    dataset_id: str,
    x_last_days: int = 0,
    separator: str = None,
    date_index: int = 0,
    date_format: str = "%Y-%m-%d",
):
    domo_ds = dmds.DomoDataset(auth=auth, id=dataset_id)

    list_partition = await domo_ds.list_partitions(auth=auth, dataset_id=dataset_id)

    today = dt.date.today()
    days_ago = today - dt.timedelta(days=x_last_days)
    for i in list_partition:
        compare_date = ""
        if separator is not None and separator != "":
            compare_date = i["partitionId"].split(separator)[date_index]
        else:
            compare_date = i["partitionId"]

        try:
            d = dt.datetime.strptime(compare_date, date_format).date()
        except ValueError:
            d = None
        if d is not None and d < days_ago:
            print(
                auth.domo_instance,
                ": 🚀  Removing partition key : ",
                (i["partitionId"]),
                " in ",
                dataset_id,
            )
            await domo_ds.delete_partition(
                dataset_partition_id=i["partitionId"], dataset_id=dataset_id, auth=auth
            )

# %% ../../nbs/integrations/Automation.ipynb 10
async def get_company_domains(
    auth: dmda.DomoAuth,
    dataset_id: str,
    handle_err_fn: callable,
    sql: str = "select domain from table",
    global_admin_username: str = None,
    global_admin_password: str = None,
    execution_env: str = None,
    debug_api: bool = False,
) -> pd.DataFrame:
    ds = await dmds.DomoDataset.get_by_id(
        auth=auth, id=dataset_id, debug_api=debug_api
    )

    print(f"⚙️ START - Retrieving company list \n{ds.display_url()}")
    print(f"⚙️ SQL = {sql}")

    df = await ds.query_dataset_private(
        auth=auth,
        dataset_id=dataset_id,
        sql=sql,
        loop_until_end=True,
        debug_api=debug_api,
    )

    df["domo_instance"] = df["domain"].apply(lambda x: x.replace(".domo.com", ""))

    if global_admin_username:
        df["domo_username"] = global_admin_username
    if global_admin_password:
        df["domo_password"] = global_admin_password

    if execution_env:
        df["env"] = execution_env or "manual"

    if df.empty:
        raise Exception("no companies retrieved")

    print(
        f"\n⚙️ SUCCESS 🎉 Retrieved company list \nThere are {len(df.index)} companies to update"
    )

    return df
