# AUTOGENERATED! DO NOT EDIT! File to edit: ../../notebooks/CLI_Helper.ipynb.

# %% auto 0
__all__ = ['logger', 'requires_auth_token', 'humanize_date', 'humanize_number', 'humanize_size', 'get_example_for_type',
           'get_example_output_format', 'customize_output_format', 'separate_integers_and_strings',
           'echo_formatted_output', 'display_formated_table', 'requires_totp', 'get_phone_registration_status']

# %% ../../notebooks/CLI_Helper.ipynb 3
import logging
from typing import *

# %% ../../notebooks/CLI_Helper.ipynb 4
import os
from contextlib import contextmanager
import functools
import ast
import datetime as dt

import pandas as pd
import typer
import humanize
from tabulate import tabulate

from airt.client import Client, User
from airt.logger import get_logger, set_level
from airt.constant import SERVER_URL, CLIENT_NAME, SERVICE_TOKEN, SERVICE_USERNAME

# %% ../../notebooks/CLI_Helper.ipynb 6
logger = get_logger(__name__)

# %% ../../notebooks/CLI_Helper.ipynb 12
@contextmanager
def authenicate_user():
    Client(
        auth_token=os.environ[SERVICE_TOKEN], server=os.environ.get(SERVER_URL, None)
    )

    yield

# %% ../../notebooks/CLI_Helper.ipynb 15
def requires_auth_token(func):
    @functools.wraps(func)
    def wrapper_decorator(*args, **kwargs):
        try:

            if ("debug" in kwargs) and kwargs["debug"]:
                set_level(logging.DEBUG)
            else:
                set_level(logging.WARNING)

            with authenicate_user():
                # Do something before
                return func(*args, **kwargs)
                # Do something after

        except KeyError as e:
            typer.echo(
                message=f"KeyError: The environment variable {e} is not set.", err=True
            )

            if f"'{SERVICE_TOKEN}'" in str(e):
                typer.echo(
                    f"\nPlease run the command '{CLIENT_NAME} token' to get the application token and set it in the "
                    f"environment variable `{SERVICE_TOKEN}`."
                )
                typer.echo(f"\nTry '{CLIENT_NAME} token --help' for help.")

            raise typer.Exit(code=1)

        except Exception as e:
            typer.echo(message=f"Error: {e}", err=True)
            if ("Invalid OTP" in str(e)) or ("OTP is required" in str(e)):
                raise ValueError(e)
            raise typer.Exit(code=1)

    return wrapper_decorator

# %% ../../notebooks/CLI_Helper.ipynb 18
def humanize_date(s: pd.Series) -> pd.Series:
    return s.apply(
        lambda date: humanize.naturaltime(
            dt.datetime.now() - dt.datetime.strptime(date, "%Y-%m-%dT%H:%M:%S")  # type: ignore
        )
        if date
        else "None"
    )

# %% ../../notebooks/CLI_Helper.ipynb 20
def humanize_number(s: pd.Series) -> pd.Series:
    return s.apply(
        lambda num: humanize.intcomma(int(num)) if pd.notna(num) else "unknown"
    )

# %% ../../notebooks/CLI_Helper.ipynb 22
def humanize_size(s: pd.Series) -> pd.Series:
    return s.apply(
        lambda size: humanize.naturalsize(size) if pd.notna(size) else "unknown"
    )

# %% ../../notebooks/CLI_Helper.ipynb 24
def get_example_for_type(xs: pd.Series) -> str:
    """Get example output for the given series

    Args:
        xs: Input series

    Returns:
        The valid formatting example for the series
    """
    if pd.api.types.is_float_dtype(xs):
        return "€{:,.2f}"
    if pd.api.types.is_integer_dtype(xs):
        return "{:,d}"

    return "{}"

# %% ../../notebooks/CLI_Helper.ipynb 26
def get_example_output_format(df: pd.DataFrame) -> Dict[str, str]:
    """Get example output format for the dataframe

    Args:
        df: Input dataframe

    Returns:
        The example output format for the dataframe
    """
    return {c: get_example_for_type(df[c]) for c in df.columns}  # type: ignore

# %% ../../notebooks/CLI_Helper.ipynb 28
def customize_output_format(format_str: str, df: pd.DataFrame) -> pd.DataFrame:
    """Customize output format

    Args:
        format_str: A dict mapping of column names into their python format
        df: Input dataframe

    Returns:
        The formatted pandas DataFrame

    Raises:
        Error: If the formatting string is not a valid python expression
        Error: If invalid column name is passed
        Error: If invalid formatting string is passed
        Error: If wrong formatting is passed for a column
    """
    try:
        formatters = ast.literal_eval(format_str)
    except Exception as e:
        typer.echo(f"Not a valid python expression: {format_str}", err=True)
        typer.echo(
            f"An example of a valid formatting string: {get_example_output_format(df)}",
            err=True,
        )
        raise typer.Exit(code=1)

    if not isinstance(formatters, dict):
        typer.echo(f"The format string is not a dictionary: {formatters}", err=True)
        typer.echo(
            f"An example of a valid formatting string: {get_example_output_format(df)}",
            err=True,
        )
        raise typer.Exit(code=1)

    if not (set(formatters.keys()) <= set(df.columns)):
        typer.echo(
            f"The following columns are not valid: {set(formatters.keys()) - set(df.columns)}. Only the following columns are valid: {set(df.columns)}",
            err=True,
        )
        typer.echo(
            f"An example of a valid formatting string: {get_example_output_format(df)}",
            err=True,
        )
        raise typer.Exit(code=1)

    df_copy = df.copy()
    try:
        for k, v in formatters.items():
            df_copy[k] = df_copy[k].apply(lambda x: v.format(x))
    except Exception as e:
        typer.echo(f"Formatting is wrong for {k}: {v}", err=True)
        typer.echo(
            f"An example of a valid formatting string: {get_example_output_format(df)}",
            err=True,
        )
        raise typer.Exit(code=1)
    return df_copy[formatters.keys()]  # type: ignore

# %% ../../notebooks/CLI_Helper.ipynb 34
def separate_integers_and_strings(xs: List[str]) -> List[Union[int, str]]:
    """Seperate integers and strings from the list of strings

    Args:
        xs: List containing string inputs

    Returns:
        A list containing the integers and strings
    """
    return [int(v) if v.isdigit() else v for v in xs]

# %% ../../notebooks/CLI_Helper.ipynb 36
def echo_formatted_output(df: pd.DataFrame):
    """Echo the formatted output to the terminal

    Args:
        df: Input DataFrame
    """
    if len(df.columns) > 1:
        typer.echo(tabulate(df, headers="keys", tablefmt="plain", showindex=False))  # type: ignore
    else:
        single_col_results = df.iloc[:, 0].astype(str).to_list()
        typer.echo("\n".join(single_col_results))

# %% ../../notebooks/CLI_Helper.ipynb 40
def display_formated_table(func):
    """A decorator function to format the CLI table output"""

    @functools.wraps(func)
    def wrapper(*args, **kwargs):

        # Do something before
        result_dict = func(*args, **kwargs)
        # Do something after

        df = result_dict["df"]
        quite_column_name = (
            result_dict["quite_column_name"]
            if "quite_column_name" in result_dict
            else "uuid"
        )

        if kwargs["format"]:
            df = customize_output_format(kwargs["format"], df)
            echo_formatted_output(df)

        elif "quiet" in kwargs and kwargs["quiet"]:
            ids = df[quite_column_name].astype(str).to_list()
            typer.echo("\n".join(ids))

        else:
            typer.echo(
                tabulate(
                    result_dict["df"],
                    headers="keys",
                    tablefmt="plain",
                    showindex=False,
                    missingval="<none>",
                )
            )

    return wrapper

# %% ../../notebooks/CLI_Helper.ipynb 42
def requires_totp_or_otp(
    message_template_name: str, no_retries: int = 3, requires_auth_token: bool = True
):
    """A decorator function to prompt users to enter a valid totp or otp"""

    def wrapper_fn(f):
        @functools.wraps(f)
        def new_wrapper(*args, **kwargs):
            if requires_auth_token:
                try:
                    Client(
                        auth_token=os.environ[SERVICE_TOKEN],
                        server=os.environ.get(SERVER_URL, None),
                    )
                except KeyError as e:
                    typer.echo(
                        message=f"KeyError: The environment variable {e} is not set.",
                        err=True,
                    )

                    if f"'{SERVICE_TOKEN}'" in str(e):
                        typer.echo(
                            f"\nPlease run the command '{CLIENT_NAME} token' to get the application token and set it in the "
                            f"environment variable `{SERVICE_TOKEN}`."
                        )
                        typer.echo(f"\nTry '{CLIENT_NAME} token --help' for help.")

                    raise typer.Exit(code=1)

            # Non-Interactive mode
            if kwargs["otp"] is not None:
                try:
                    return f(*args, **kwargs)
                except Exception as e:
                    raise typer.Exit(code=1)
            else:
                # Interactive mode
                if not requires_auth_token:
                    try:
                        return f(*args, **kwargs)
                    except ValueError as e:
                        pass
                typer.echo("\nPlease choose an option\n\n")
                while True:
                    typer.echo(
                        "[1] Use the dynamically generated six-digit verification code from the authenticator application\n"
                    )
                    typer.echo(
                        "[2] Request the OTP via SMS to the registered phone number\n"
                    )
                    typer.echo(
                        "If you cannot access the authenticator application and your registered phone number, please contact your administrator.\n"
                    )

                    user_option = typer.prompt("Enter your option")

                    if user_option in ["1", "2"]:
                        break
                    typer.echo("Please enter a valid option")

                if user_option == "1":
                    for i in range(no_retries):
                        _totp = typer.prompt(
                            "Please enter the OTP displayed in the authenticator app"
                        )
                        kwargs["otp"] = _totp
                        try:
                            return f(*args, **kwargs)
                        except ValueError as e:
                            pass
                    raise typer.Exit(code=1)
                else:
                    if requires_auth_token:
                        username = User.details()["username"]
                    else:
                        username = (
                            kwargs["username"]
                            if kwargs["username"] is not None
                            else os.environ.get(SERVICE_USERNAME)
                        )
                    sms_status = User.send_sms_otp(
                        username=username, message_template_name=message_template_name
                    )
                    typer.echo(f"\n{sms_status}\n")
                    for i in range(no_retries):
                        _sms_otp = typer.prompt(
                            f"Please enter the One-Time Password (OTP) you received on your registered phone number"
                        )
                        kwargs["otp"] = _sms_otp
                        try:
                            return f(*args, **kwargs)
                        except ValueError as e:
                            pass
                    raise typer.Exit(code=1)

        return new_wrapper

    return wrapper_fn

# %% ../../notebooks/CLI_Helper.ipynb 44
def requires_totp(no_retries: int = 3):
    """A decorator function to prompt users to enter a valid totp"""

    def wrapper_fn(f):
        @functools.wraps(f)
        def new_wrapper(*args, **kwargs):
            # Non-Interactive mode
            if kwargs["otp"] is not None:
                try:
                    return f(*args, **kwargs)
                except Exception as e:
                    raise typer.Exit(code=1)
            else:
                # Interactive mode
                _activation_otp = None
                for i in range(no_retries):
                    if _activation_otp is not None:
                        kwargs["otp"] = _activation_otp
                    try:
                        return f(*args, **kwargs)
                    except ValueError as e:
                        _activation_otp = typer.prompt(
                            "Please enter the OTP displayed in the authenticator app"
                        )
                raise typer.Exit(code=1)

        return new_wrapper

    return wrapper_fn

# %% ../../notebooks/CLI_Helper.ipynb 46
PHONE_REGISTRATION_STATUS = {
    "not_registered": (
        "\n\nPlease take a moment to register and verify your phone number. If you forget your password or cannot access your account, you can request "
        "the OTP to your registered phone number to regain access. \nTo register a new phone number, please set the token in the AIRT_SERVICE_TOKEN environment "
        f"variable and execute the below command with your phone number and follow the on-screen instructions:\n\n{CLIENT_NAME} user register-phone-number --phone-number"
    ),
    "not_validated": (
        "\n\nYour phone number is added to your account but not yet verified. Please take a moment to register and verify your phone number. If you forget "
        "your password or cannot access your account, you can request the OTP to your registered phone number to regain access. \nTo register a new phone number, please set "
        "the token in the AIRT_SERVICE_TOKEN environment variable and execute the below command with your phone number and follow the on-screen instructions:"
        f"\n\n{CLIENT_NAME} user register-phone-number"
        "\n\nIn case you want to register a new number, please execute the below command with your new phone number:"
        f"\n\n{CLIENT_NAME} user register-phone-number --phone-number"
    ),
}

# %% ../../notebooks/CLI_Helper.ipynb 47
def get_phone_registration_status(xs: Dict[str, Union[str, bool]]) -> Optional[str]:
    """Get the phone number registration status

    Args:
        xs: A dict containing the user details

    Returns:
        None, if the user phone number is registred and validated. Else, a message containing the current state of the phone number registration process.
    """
    if not xs["phone_number"]:
        return PHONE_REGISTRATION_STATUS["not_registered"]
    if not xs["is_phone_number_verified"]:
        return PHONE_REGISTRATION_STATUS["not_validated"]
    return None
