# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/CLI_DataSource.ipynb (unless otherwise specified).

__all__ = ["logger"]

# Cell

from typing import *

# Internal Cell

import os

import typer
from typer import echo
from tabulate import tabulate
import datetime as dt
import pandas as pd

from ..client import Client
from . import helper
from ..logger import get_logger, set_level

# Internal Cell

app = typer.Typer()

local_ds_app = typer.Typer()

app.add_typer(
    local_ds_app,
    name="local",
    short_help="Pull the data source from your local machine to airt server.",
    help="Pull the data source from your local machine to airt server. "
    "Currently, we support reading and pushing the data from a CSV file. "
    "We plan to support more data source formats in the future.",
)

# Cell

logger = get_logger(__name__)

# Internal Cell


@app.command()
@helper.requires_auth_token
def s3(
    uri: str = typer.Argument(
        ..., help="The AWS S3 bucket location of the Parquet files as a string."
    ),
    access_key: Optional[str] = typer.Option(
        None,
        help="The access key for the S3 bucket. If None (default value), then the value of environment variable AWS_ACCESS_KEY_ID is used.",
    ),
    secret_key: Optional[str] = typer.Option(
        None,
        help="The secret key for the S3 bucket. If None (default value), then the value of environment variable AWS_SECRET_ACCESS_KEY is used.",
    ),
    tag: Optional[str] = typer.Option(
        None,
        help="The tag name for the data source. If None (default value), then the tag latest will be assigned to the data source.",
    ),
    quiet: bool = typer.Option(
        False,
        "--quiet",
        "-q",
        help="Output data id only.",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        "-d",
        help="Set logger level to DEBUG and output everything.",
    ),
):
    """Pull the data source from a S3 bucket to airt server."""

    from ..client import DataSource

    ds = DataSource.s3(uri=uri, access_key=access_key, secret_key=secret_key, tag=tag)

    progress = ds.pull()

    if quiet:
        progress.wait()

        typer.echo(f"{ds.id}")
    else:
        typer.echo(f"Pulling data ID: {ds.id}")

        progress.progress_bar()


# Internal Cell


@app.command()
@helper.requires_auth_token
def db(
    host: str = typer.Option(..., help="The name of the remote database host machine."),
    database: str = typer.Option(
        ..., help="The logical name of the database to establish the connection."
    ),
    table: str = typer.Option(..., help="The name of the table in the database."),
    port: int = typer.Option(
        3306,
        help="The port for the database server. If the value is not passed then the default port number will be used (e.g. for MySQL, 3306 will be used)",
    ),
    engine: str = typer.Option(
        "mysql",
        help="""The name of the database engine. If the value is not passed then the
            default database engine for MySQL will be used.""",
    ),
    username: Optional[str] = typer.Option(
        None,
        "--username",
        "-u",
        help="A valid database user name. If not set (default value root), it will try to use the value from environment variable AIRT_CLIENT_DB_USERNAME.",
    ),
    password: Optional[str] = typer.Option(
        None,
        "--password",
        "-p",
        help="The password for the specified user. If not set (default value ), it will try to use the value from environment variable AIRT_CLIENT_DB_PASSWORD.",
    ),
    tag: Optional[str] = typer.Option(
        None,
        help="The tag name for the data source. If None (default value), then the tag latest will be assigned to the data source.",
    ),
    quiet: bool = typer.Option(
        False,
        "--quiet",
        "-q",
        help="Output data id only.",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        "-d",
        help="Set logger level to DEBUG and output everything.",
    ),
):
    """Creates an object that encapsulates the data from a database."""

    from ..client import DataSource

    ds = DataSource.db(
        host=host,
        database=database,
        port=port,
        table=table,
        engine=engine,
        username=username,
        password=password,
        tag=tag,
    )

    progress = ds.pull()

    if quiet:
        progress.wait()
        typer.echo(f"{ds.id}")
    else:
        typer.echo(f"Pulling data ID: {ds.id}")
        progress.progress_bar()


# Internal Cell


@app.command()
@helper.requires_auth_token
def ls(
    offset: int = typer.Option(
        0,
        "--offset",
        "-o",
        help="The number of records to offset at the beginning of the datasource list. If None, then the default value 0 will be used.",
    ),
    limit: int = typer.Option(
        100,
        "--limit",
        "-l",
        help="The maximum number of records to return from the server. If None, then the default value 100 will be used.",
    ),
    disabled: bool = typer.Option(
        False,
        "--disabled",
        help="If set to True, then only the deleted datasources will be displayed. Else, the default value "
        "False will be used to display only the list of active datasources.",
    ),
    quiet: bool = typer.Option(
        False,
        "--quiet",
        "-q",
        help="Output only ids of datasource separated by space",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        "-d",
        help="Set logger level to DEBUG and output everything.",
    ),
) -> None:
    """List available data sources."""

    from ..client import DataSource

    df = DataSource.ls(offset=offset, limit=limit, disabled=disabled)

    df["created"] = helper.humanize_date(df["created"])
    df["no_of_rows"] = helper.humanize_number(df["no_of_rows"])
    df["folder_size"] = helper.humanize_size(df["folder_size"])

    if quiet:
        ids = df["id"].astype(str).to_list()
        typer.echo("\n".join(ids))
    else:
        typer.echo(tabulate(df, headers="keys", tablefmt="plain", showindex=False))


# Internal Cell


@app.command()
@helper.requires_auth_token
def dtypes(
    id: str = typer.Argument(
        ...,
        help="The data ID in the airt service.",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        "-d",
        help="Set logger level to DEBUG and output everything.",
    ),
) -> None:
    """Return the dtype of each column for the given data source."""

    from ..client import DataSource

    ds = DataSource(data_id=int(id))
    dtypes = ds.dtypes.T.rename(columns={0: "dtype"})
    typer.echo(dtypes)


# Internal Cell


@app.command()
@helper.requires_auth_token
def head(
    id: str = typer.Argument(
        ...,
        help="The data ID in the airt service.",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        "-d",
        help="Set logger level to DEBUG and output everything.",
    ),
) -> None:
    """Return first few records for the given data ID."""

    from ..client import DataSource

    ds = DataSource(data_id=int(id))
    typer.echo(tabulate(ds.head(), headers="keys", tablefmt="plain", showindex=False))


# Internal Cell


@app.command()
@helper.requires_auth_token
def train(
    id: str = typer.Option(
        ...,
        "--data_id",
        "-id",
        help="The data ID in the airt service.",
    ),
    client_column: str = typer.Option(
        ...,
        "--client_column",
        help="The name of the column that uniquely identifies the users/clients as string.",
    ),
    timestamp_column: Optional[str] = typer.Option(
        None,
        "--timestamp_column",
        help="Name of the timestamp_column specifying the time of an occurred event as a string.",
    ),
    target_column: str = typer.Option(
        ...,
        "--target_column",
        help="Name of the target event for which the model needs to be trained to make predictions.",
    ),
    target: str = typer.Option(
        ...,
        "--target",
        help="Name of the target event for which the model needs to be trained to make predictions.",
    ),
    predict_after: str = typer.Option(
        ...,
        "--predict_after",
        help="Time delta in hours of the expected target event.",
    ),
    quiet: bool = typer.Option(
        False,
        "--quiet",
        "-q",
        help="Output model id only.",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        "-d",
        help="Set logger level to DEBUG and output everything.",
    ),
):
    """Train a ML model for the given data source."""

    from ..client import DataSource
    from datetime import timedelta

    ds = DataSource(data_id=int(id))

    model = ds.train(
        client_column=client_column,
        target_column=target_column,
        target=target,
        predict_after=timedelta(hours=int(predict_after)),
        timestamp_column=timestamp_column,
    )

    if quiet:
        model.wait()
        typer.echo(f"{model.model_id}")
    else:
        typer.echo(f"Training started for Model ID: {model.model_id}")
        model.progress_bar()


# Internal Cell


@app.command()
@helper.requires_auth_token
def details(
    id: int = typer.Argument(
        ...,
        help="The data ID in the airt service.",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        "-d",
        help="Set logger level to DEBUG and output everything.",
    ),
) -> None:
    """Return metadata of the given data source."""

    from ..client import DataSource

    df = DataSource.details(data_id=int(id))

    df["created"] = helper.humanize_date(df["created"])
    df["no_of_rows"] = helper.humanize_number(df["no_of_rows"])
    df["folder_size"] = helper.humanize_size(df["folder_size"])

    typer.echo(tabulate(df, headers="keys", tablefmt="plain", showindex=False))


# Internal Cell


@app.command()
@helper.requires_auth_token
def rm(
    id: int = typer.Argument(
        ...,
        help="The data ID in the airt service.",
    ),
    quiet: bool = typer.Option(
        False,
        "--quiet",
        "-q",
        help="Output the deleted data id only.",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        "-d",
        help="Set logger level to DEBUG and output everything.",
    ),
) -> None:
    """Remove the data source from airt service."""

    from ..client import DataSource

    df = DataSource.delete(data_id=int(id))

    df["created"] = helper.humanize_date(df["created"])
    df["no_of_rows"] = helper.humanize_number(df["no_of_rows"])
    df["folder_size"] = helper.humanize_size(df["folder_size"])

    if quiet:
        typer.echo(df.iloc[0]["id"])
    else:
        typer.echo(tabulate(df, headers="keys", tablefmt="plain", showindex=False))


# Internal Cell


@app.command()
@helper.requires_auth_token
def push(
    data_id: int = typer.Option(
        ...,
        "--data",
        help="The data ID in the airt service.",
    ),
    pred_id: int = typer.Option(
        ...,
        "--pred",
        help="The prediction ID in the airt service.",
    ),
    quiet: bool = typer.Option(
        False,
        "--quiet",
        "-q",
        help="Output the pushed prediction id only.",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        "-d",
        help="Set logger level to DEBUG and output everything.",
    ),
) -> None:
    """Push the prediction results into the target data source."""

    from ..client import DataSource, Prediction

    datasource = DataSource(int(data_id))
    prediction = Prediction(prediction_id=int(pred_id), datasource_id=int(data_id))

    progress = datasource.push(prediction)

    if quiet:
        progress.wait()

        typer.echo(prediction.prediction_id)
    else:
        typer.echo(
            f"Pushing the results of prediction id: {prediction.prediction_id} into the datasource id: {prediction.datasource_id}"
        )

        progress.progress_bar()


# Internal Cell


@app.command()
@helper.requires_auth_token
def tag(
    id: int = typer.Option(
        ...,
        "--data_id",
        "-id",
        help="The data ID in the airt service.",
    ),
    name: str = typer.Option(
        ...,
        "--tag_name",
        "-n",
        help="The tag name for the datasource.",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        "-d",
        help="Set logger level to DEBUG and output everything.",
    ),
) -> None:
    """Tag an existing datasource in airt service."""

    from ..client import DataSource

    df = DataSource.tag(data_id=int(id), name=name)

    df["created"] = helper.humanize_date(df["created"])
    df["no_of_rows"] = helper.humanize_number(df["no_of_rows"])
    df["folder_size"] = helper.humanize_size(df["folder_size"])

    typer.echo(tabulate(df, headers="keys", tablefmt="plain", showindex=False))


# Internal Cell


@local_ds_app.command()
@helper.requires_auth_token
def csv(
    file_path: str = typer.Option(
        ...,
        "--file_path",
        "-f",
        help="The relative or absolute path of the csv file.",
    ),
    index_column: str = typer.Option(
        ...,
        "--index_column",
        "-i",
        help="The Column name to use as the row labels.",
    ),
    sort_by: str = typer.Option(
        ...,
        "--sort_by",
        "-s",
        help="The Column name to sort the data.",
    ),
    tag_name: Optional[str] = typer.Option(
        None,
        "--tag_name",
        "-t",
        help="The tag name for the datasource.",
    ),
    deduplicate_data: bool = typer.Option(
        False,
        "--deduplicate_data",
        "-dd",
        help="A boolean flag to handle the duplicate rows in the CSV file. If set to **True** (default value **False**), "
        "the duplicate rows will be removed from the CSV file before uploading to airt server.",
    ),
    blocksize: str = typer.Option(
        "256MB",
        "--blocksize",
        "-bs",
        help="The number of bytes by which to cut up larger files. If None, the default value **256MB** will be used for each file.",
    ),
    kwargs_json: Optional[str] = typer.Option(
        None,
        "--kwargs_json",
        "-json",
        help="Any additional parameters for Pandas read_csv method as a JSON string.",
    ),
    quiet: bool = typer.Option(
        False,
        "--quiet",
        "-q",
        help="Output data id only.",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        "-d",
        help="Set logger level to DEBUG and output everything.",
    ),
) -> None:
    """Pull the data source from a local CSV file to airt server."""

    from ..client import DataSource

    ds = DataSource.csv(
        file_path=file_path,
        index_column=index_column,
        sort_by=sort_by,
        tag_name=tag_name,
        deduplicate_data=deduplicate_data,
        blocksize=blocksize,
        kwargs_json=kwargs_json,
    )

    progress = ds.pull()

    if quiet:
        progress.wait()

        typer.echo(f"{ds.id}")
    else:
        typer.echo(f"Pulling data ID: {ds.id}")

        progress.progress_bar()
