#
# # Copyright © 2023 Peak AI Limited. or its affiliates. All Rights Reserved.
# #
# # Licensed under the Apache License, Version 2.0 (the "License"). You
# # may not use this file except in compliance with the License. A copy of
# # the License is located at:
# #
# # https://github.com/PeakBI/peak-sdk/blob/main/LICENSE
# #
# # or in the "license" file accompanying this file. This file is
# # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# # ANY KIND, either express or implied. See the License for the specific
# # language governing permissions and limitations under the License.
# #
# # This file is part of the peak-sdk.
# # see (https://github.com/PeakBI/peak-sdk)
# #
# # You should have received a copy of the APACHE LICENSE, VERSION 2.0
# # along with this program. If not, see <https://apache.org/licenses/LICENSE-2.0>
#
"""Peak Workflow management service commands."""

from typing import Any, Dict, List, Optional

import typer
from peak.cli import args, helpers
from peak.helpers import parse_list_of_strings
from peak.resources.workflows import Workflow
from rich.console import Console
from typing_extensions import Annotated

app = typer.Typer(
    help="Create Machine Learning Pipelines using Workflows. Please note that Workflows with only standard steps are supported.",
    short_help="Create and manage Workflows.",
)
console = Console()

_WORKFLOW_STATUS = typer.Option(None, help="List of status to filter workflows.")

_LAST_EXECUTION_STATUS = typer.Option(
    None,
    help="List of execution status of the latest run of the workflows to filter them by.",
)

_LAST_MODIFIED_BY = typer.Option(None, help="List of users who last modified the workflow, used to filter workflows.")

_WORKFLOW_ID = typer.Argument(..., help="ID of the workflow to be used in this operation.")

_LIST_NAME = typer.Option(None, help="Workflow name to search for.")

_NAME = typer.Option(None, help="Name of the workflow")

_REPOSITORY = typer.Option(None, help="URL of the repository containing the required files.")

_BRANCH = typer.Option(None, help="The branch of the repository to use.")

_TOKEN = typer.Option(None, help="The token to use to access the repository.")

_COMMAND = typer.Option(None, help="The command to run when workflow step is executed.")

_IMAGE_ID = typer.Option(None, help="The ID of the image to use for the workflow step.")

_IMAGE_VERSION_ID = typer.Option(None, help="The ID of the image version to use for the workflow step.")

_INSTANCE_TYPE_ID = typer.Option(None, help="The ID of the instance type to use for the workflow step.")

_STORAGE = typer.Option(None, help="The storage to use for the workflow step.")

_STEP_TIMEOUT = typer.Option(None, help="The time after which the step timeouts.")

_CLEAR_IMAGE_CACHE = typer.Option(None, help="Whether to clear image cache on workflow execution.")

_STEP_NAMES = typer.Option(None, help="List of step names to be updated.")


@app.command(short_help="Create a new workflow.")
def create(
    ctx: typer.Context,
    file: str = args.TEMPLATE_PATH,
    params_file: str = args.TEMPLATE_PARAMS_FILE,
    params: List[str] = args.TEMPLATE_PARAMS,
) -> None:
    """***Create*** a new workflow. Workflows with only standard steps are supported.

    \b
    🧩 ***Input file schema(yaml):***<br/>
    ```yaml
      body (map):
        name (string): Name of the worklfow.
        tags (list(map) | required: false):
            - name (string): Name of the tag.
        triggers (list(map) | required: false):
            - cron (string | required: false): A valid cron expression.
              webhook (boolean | required: false): Should be true if webhook type trigger is to be used.
              webhookPolicy (string | required: false): Policy of the webhook to be used. Should be "generate" when creating new worklfow.
        watchers (list(map) | required: false):
            - events (map):
                success (boolean | required: false): Whether to call event on success.
                fail (boolean | required: false): Whether to call event on success.
                runtimeExceeded (int | required: false): The runtime after which event is called.
              user (string | required: false): User to be notified.
              webhook (map | required: false):
                name (string): Name of the webhook.
                url (string): URL of the webhook.
                payload (string): Webhook payload.
        steps(map):
            <stepName> (map): Dictionary containing the step configuration. Here the key is name of the step.
                imageId (int): ID of the existing image.
                imageVersionId: (int | required: false): ID of the existing image version.
                type (string | required: false): Type of workflow step. Currently only standard type is supported.
                command (string): Command to run when step is executed.
                clearImageCache (boolean | required: false): Whether to clear image cache on workflow execution.
                stepTimeout (int | required: false): Time after which the step timeouts.
                parameters (map | required: false):
                    env (map | required: false): Key-Value pair where key is the name of the env.
                    secrets (list(str) | required: false): List of secret names to be passed.
                parents (list(str) | required: false): List containing names of steps on which this step is dependent on.
                repository (map | required: false):
                    branch (string): Branch of the repository containing the required files.
                    token (string | required: false): The token to be used to clone the repository.
                    url (string): URL of the repository.
                resources (map | required: false):
                    instanceTypeId (int): ID of the instance type to be used in the step.
                    storage (string): Storage in GB. For example, "10GB".
    ```

    \b
    📝 ***Example usage:***
    ```bash
    peak workflows create /path/to/file.yml -v /path/to/params.yml
    ```

    \b
    🆗 ***Response:***
    ```json
    {
        "id": 123,
        "triggers": [
            {
                "webhook": "db88c21d-1add-45dd-a72e-8c6b83b68dee"
            }
        ]
    }
    ```

    🔗 [**API Documentation**](https://service.peak.ai/workflows/api-docs/index.htm#/Workflows/create-workflow)
    """
    workflow_client: Workflow = ctx.obj["client"]

    body: Dict[str, Any] = helpers.template_handler(file=file, params_file=params_file, params=params)
    body = helpers.remove_unknown_args(body, workflow_client.create_workflow)
    response: Dict[str, int] = workflow_client.create_workflow(**body)
    console.print(response)


@app.command(short_help="Update an existing workflow.")
def update(
    ctx: typer.Context,
    workflow_id: int = _WORKFLOW_ID,
    file: str = args.TEMPLATE_PATH,
    params_file: str = args.TEMPLATE_PARAMS_FILE,
    params: List[str] = args.TEMPLATE_PARAMS,
) -> None:
    """***Update*** an existing workflow. Workflows with only standard steps are supported.

    \b
    🧩 ***Input file schema(yaml):***<br/>
    ```yaml
      body (map):
        name (string): Name of the worklfow.
        tags (list(map) | required: false):
            - name (string): Name of the tag.
        triggers (list(map) | required: false):
            - cron (string | required: false): A valid cron expression.
              webhook (boolean | required: false): Should be true if webhook type trigger is to be used.
              webhookPolicy (string | required: false): Policy of the webhook to be used. Should be "generate" or "preserve".
        watchers (list(map) | required: false):
            - events (map):
                success (boolean | required: false): Whether to call event on success.
                fail (boolean | required: false): Whether to call event on success.
                runtimeExceeded (int | required: false): The runtime after which event is called.
              user (string | required: false): User to be notified.
              webhook (map | required: false):
                name (string): Name of the webhook.
                url (string): URL of the webhook.
                payload (string): Webhook payload.
        steps(map):
            <stepName> (map): Dictionary containing the step configuration. Here the key is name of the step.
                imageId (int): ID of the existing image.
                imageVersionId: (int | required: false): ID of the existing image version.
                type (string | required: false): Type of workflow step. Currently only standard type is supported.
                command (string): Command to run when step is executed.
                clearImageCache (boolean | required: false): Whether to clear image cache on workflow execution.
                stepTimeout (int | required: false): Time after which the step timeouts.
                parameters (map | required: false):
                    env (map | required: false): Key-Value pair where key is the name of the env.
                    secrets (list(str) | required: false): List of secret names to be passed.
                parents (list(str) | required: false): List containing names of steps on which this step is dependent on.
                repository (map | required: false):
                    branch (string): Branch of the repository containing the required files.
                    token (string | required: false): The token to be used to clone the repository.
                    url (string): URL of the repository.
                resources (map | required: false):
                    instanceTypeId (int): ID of the instance type to be used in the step.
                    storage (string): Storage in GB. For example, "10GB".
    ```

    \b
    📝 ***Example usage:***
    ```bash
    peak workflows update 9999 /path/to/file.yml -v /path/to/params.yml
    ```

    \b
    🆗 ***Response:***
    ```json
    {
        "id": 123,
        "triggers": [
            {
                "webhook": "db88c21d-1add-45dd-a72e-8c6b83b68dee"
            }
        ]
    }
    ```

    🔗 [**API Documentation**](https://service.peak.ai/workflows/api-docs/index.htm#/Workflows/update-workflow)
    """
    workflow_client: Workflow = ctx.obj["client"]

    body: Dict[str, Any] = helpers.template_handler(file=file, params_file=params_file, params=params)
    body = helpers.remove_unknown_args(body, workflow_client.update_workflow)
    response: Dict[str, Any] = workflow_client.update_workflow(workflow_id=workflow_id, **body)
    console.print(response)


@app.command(short_help="Create a new workflow or Update an existing workflow.")
def create_or_update(
    ctx: typer.Context,
    file: str = args.TEMPLATE_PATH,
    params_file: str = args.TEMPLATE_PARAMS_FILE,
    params: List[str] = args.TEMPLATE_PARAMS,
) -> None:
    """***Create*** a new workflow or ***Update*** an existing workflow based on workflow name.

    \b
    🧩 ***Input file schema(yaml):***<br/>
    ```yaml
      body (map):
        name (string): Name of the worklfow.
        tags (list(map) | required: false):
            - name (string): Name of the tag.
        triggers (list(map) | required: false):
            - cron (string | required: false): A valid cron expression.
              webhook (boolean | required: false): Should be true if webhook type trigger is to be used.
              webhookPolicy (string | required: false): Policy of the webhook to be used. Should be "generate" when creating new worklfow.
        watchers (list(map) | required: false):
            - events (map):
                success (boolean | required: false): Whether to call event on success.
                fail (boolean | required: false): Whether to call event on success.
                runtimeExceeded (int | required: false): The runtime after which event is called.
              user (string | required: false): User to be notified.
              webhook (map | required: false):
                name (string): Name of the webhook.
                url (string): URL of the webhook.
                payload (string): Webhook payload.
        steps(map):
            <stepName> (map): Dictionary containing the step configuration. Here the key is name of the step.
                imageId (int): ID of the existing image.
                imageVersionId: (int | required: false): ID of the existing image version.
                type (string | required: false): Type of workflow step. Currently only standard type is supported.
                command (string): Command to run when step is executed.
                clearImageCache (boolean | required: false): Whether to clear image cache on workflow execution.
                stepTimeout (int | required: false): Time after which the step timeouts.
                parameters (map | required: false):
                    env (map | required: false): Key-Value pair where key is the name of the env.
                    secrets (list(str) | required: false): List of secret names to be passed.
                parents (list(str) | required: false): List containing names of steps on which this step is dependent on.
                repository (map | required: false):
                    branch (string): Branch of the repository containing the required files.
                    token (string | required: false): The token to be used to clone the repository.
                    url (string): URL of the repository.
                resources (map | required: false):
                    instanceTypeId (int): ID of the instance type to be used in the step.
                    storage (string): Storage in GB. For example, "10GB".
    ```

    \b
    📝 ***Example usage:***
    ```bash
    peak workflows create-or-update /path/to/file.yml -v /path/to/params.yml
    ```

    \b
    🆗 ***Response:***
    ```json
    {
        "id": 123,
        "triggers": [
            {
                "webhook": "db88c21d-1add-45dd-a72e-8c6b83b68dee"
            }
        ]
    }
    ```

    🔗 [**API Documentation**](https://service.peak.ai/workflows/api-docs/index.htm#/Workflows/create-workflow)
    """
    workflow_client: Workflow = ctx.obj["client"]

    body: Dict[str, Any] = helpers.template_handler(file=file, params_file=params_file, params=params)
    body = helpers.remove_unknown_args(body, workflow_client.create_or_update_workflow)
    response: Dict[str, int] = workflow_client.create_or_update_workflow(**body)
    console.print(response)


@app.command(short_help="Update required fields of an existing workflow.")
def patch(
    ctx: typer.Context,
    workflow_id: int = _WORKFLOW_ID,
    file: Annotated[
        Optional[str],
        typer.Argument(
            ...,
            help="Path to the file that defines the body for this operation, supports both `yaml` file or a `jinja` template.",
        ),
    ] = None,
    params_file: str = args.TEMPLATE_PARAMS_FILE,
    params: List[str] = args.TEMPLATE_PARAMS,
    name: str = _NAME,
    repository: str = _REPOSITORY,
    branch: str = _BRANCH,
    token: str = _TOKEN,
    image_id: int = _IMAGE_ID,
    image_version_id: int = _IMAGE_VERSION_ID,
    command: str = _COMMAND,
    clear_image_cache: bool = _CLEAR_IMAGE_CACHE,
    step_timeout: int = _STEP_TIMEOUT,
    instance_type_id: int = _INSTANCE_TYPE_ID,
    storage: str = _STORAGE,
    step_names: List[str] = _STEP_NAMES,
) -> None:
    """***Update*** an existing workflow. Workflows with only standard steps are supported.

    \b
    This command allows to efficiently modify trigger details, watchers, workflow name, and specific step attributes such as repository URL, branch, token, image ID, version ID etc.

    \b
    By specifying step_names, we can globally update specified steps with provided parameters, streamlining the update process. If step_names is not provided, all the steps for that workflow would be updated.

    \b
    Alternatively, we can selectively modify individual step attributes across different steps by providing the details in the yaml file. With this, we can also add new steps to the workflow by providing the parameters required by the step.

    \b
    If both body and specific parameters are used, the latter takes precedence.

    \b
    🧩 ***Input file schema(yaml):***<br/>
    ```yaml
    body (map | required: false):
        name (string | required: false): Name of the worklfow.
        tags (list(map) | required: false):
            - name (string): Name of the tag.
        triggers (list(map) | required: false):
            - cron (string | required: false): A valid cron expression.
              webhook (boolean | required: false): Should be true if webhook type trigger is to be used.
              webhookPolicy (string | required: false): Policy of the webhook to be used. Should be "generate" or "preserve".
        watchers (list(map) | required: false):
            - events (map):
                success (boolean | required: false): Whether to call event on success.
                fail (boolean | required: false): Whether to call event on success.
                runtimeExceeded (int | required: false): The runtime after which event is called.
              user (string | required: false): User to be notified.
              webhook (map | required: false):
                name (string): Name of the webhook.
                url (string): URL of the webhook.
                payload (string): Webhook payload.
        steps(map):
            <stepName> (map): Dictionary containing the step configuration. Here the key is name of the step.
                imageId (int): ID of the existing image.
                imageVersionId: (int | required: false): ID of the existing image version.
                type (string | required: false): Type of workflow step. Currently only standard type is supported.
                command (string): Command to run when step is executed.
                clearImageCache (boolean | required: false): Whether to clear image cache on workflow execution.
                stepTimeout (int | required: false): Time after which the step timeouts.
                parameters (map | required: false):
                    env (map | required: false): Key-Value pair where key is the name of the env.
                    secrets (list(str) | required: false): List of secret names to be passed.
                parents (list(str) | required: false): List containing names of steps on which this step is dependent on.
                repository (map | required: false):
                    branch (string): Branch of the repository containing the required files.
                    token (string | required: false): The token to be used to clone the repository.
                    url (string): URL of the repository.
                resources (map | required: false):
                    instanceTypeId (int): ID of the instance type to be used in the step.
                    storage (string): Storage in GB. For example, "10GB".
    ```

    \b
    📝 ***Example usage:***
    ```bash
    peak workflows patch 9999 /path/to/file.yml -v /path/to/params.yml
    ```

    \b
    📝 ***Example usage for updating workflow by passing only required parameters:***
    ```bash
    peak workflows patch 9999 --name <workflow_name> --image-id <image-id> --version-id <image-version-id> --step-names <step-name-1>,<step-name-2>
    ```

    \b
    🆗 ***Response:***
    ```js
    {
        "triggers": [
            {
                "webhook": "db88c21d-1add-45dd-a72e-8c6b83b68dee"
            }
        ]
        "id": 9999,
    }
    ```
    """
    workflow_client: Workflow = ctx.obj["client"]
    body: Dict[str, Any] = {}

    if file:
        body = helpers.template_handler(file=file, params_file=params_file, params=params)
        body = helpers.remove_unknown_args(body, workflow_client.patch_workflow)
    response: Dict[str, Any] = workflow_client.patch_workflow(
        workflow_id=workflow_id,
        body=body["body"] if body else {},
        step_names=step_names,
        name=name,
        repository=repository,
        branch=branch,
        token=token,
        image_id=image_id,
        image_version_id=image_version_id,
        command=command,
        clear_image_cache=clear_image_cache,
        step_timeout=step_timeout,
        instance_type_id=instance_type_id,
        storage=storage,
    )
    console.print(response)


@app.command("list", short_help="List workflows.")
def list_workflows(
    ctx: typer.Context,
    page_size: Optional[int] = args.PAGE_SIZE,
    page_number: Optional[int] = args.PAGE_NUMBER,
    workflow_status: Optional[List[str]] = _WORKFLOW_STATUS,
    last_execution_status: Optional[List[str]] = _LAST_EXECUTION_STATUS,
    last_modified_by: Optional[List[str]] = _LAST_MODIFIED_BY,
    name: Optional[str] = _LIST_NAME,
) -> None:
    """***List*** all workflows that exist for the tenant.

    \b
    📝 ***Example usage:***<br/>
    ```bash
    peak workflows list --page-size 10 --page-number 1 --workflow_status "Draft" --last_execution_status "Success" --last_modified_by "abc@peak.ai" --name "test"
    ```

    \b
    🆗 ***Response:***
    ```
    {
        "pageCount": 1,
        "pageNumber": 1,
        "pageSize": 25,
        "workflows": [...],
        "workflowsCount": 1
    }
    ```

    🔗 [**API Documentation**](https://service.peak.ai/workflows/api-docs/index.htm#/Workflows/get-workflows)
    """
    workflow_client: Workflow = ctx.obj["client"]
    response = workflow_client.list_workflows(
        page_size=page_size,
        page_number=page_number,
        workflow_status=parse_list_of_strings(workflow_status),
        last_execution_status=parse_list_of_strings(last_execution_status),
        last_modified_by=parse_list_of_strings(last_modified_by),
        name=name,
        return_iterator=False,
    )
    console.print(response)


@app.command(short_help="Describe details of a workflow.")
def describe(ctx: typer.Context, workflow_id: int = _WORKFLOW_ID) -> None:
    """***Describe*** details of a specific workflow. Workflows with only standard steps are supported.

    \b
    📝 ***Example usage:***<br/>
    ```bash
    peak workflows describe 9999
    ```

    \b
    🆗 ***Response:***
    ```
    {
        "id": 9999,
        "name": "workflow-name",
        "createdAt": "2020-01-01T18:00:00.000Z",
        "updatedAt": "2020-01-01T18:00:00.000Z",
        "status": "Available",
        "steps": {
            "step-name": {...}
        },
        "lastExecution": {...},
        "triggers": [...],
        "watchers": [...],
        "tags": [...]
    }
    ```

    🔗 [**API Documentation**](https://service.peak.ai/workflows/api-docs/index.htm#/Workflows/describe-workflow)
    """
    workflow_client: Workflow = ctx.obj["client"]
    response = workflow_client.describe_workflow(workflow_id=workflow_id)
    console.print(response)


@app.command(short_help="Delete a workflow.")
def delete(ctx: typer.Context, workflow_id: int = _WORKFLOW_ID) -> None:
    """***Delete*** a workflow.

    \b
    📝 ***Example usage:***<br/>
    ```bash
    peak workflows delete 9999
    ```

    \b
    🆗 ***Response:***
    ```json
    {}
    ```

    🔗 [**API Documentation**](https://service.peak.ai/workflows/api-docs/index.htm#/Workflows/delete-workflow)
    """
    workflow_client: Workflow = ctx.obj["client"]
    response = workflow_client.delete_workflow(workflow_id=workflow_id)
    console.print(response)


@app.command(short_help="Start a workflow run.")
def execute(
    ctx: typer.Context,
    workflow_id: int = _WORKFLOW_ID,
    file: Annotated[Optional[str], args.TEMPLATE_PATH] = None,
    params_file: str = args.TEMPLATE_PARAMS_FILE,
    params: List[str] = args.TEMPLATE_PARAMS,
) -> None:
    """***Start*** a workflow run. This allows you to pass dynamic parameters to the workflow while running it.

    \b
    🧩 ***Input file schema(yaml):***<br/>
    ```yaml
      body (map):
        params (map | required: false):
            global (map | required: false):
                <keyName> (string | required: false): Key-Value pair. Key is the name of the param which and value is the value of the param.
            stepWise (map | required: false):
                <stepName> (map | required: false): Parameters to be passed to the step with name `stepName`.
                    <keyName> (string | required: false): Key-Value pair. Key is the name of the param which and value is the value of the param.
    ```

    \b
    📝 ***Example usage:***
    ```bash
    peak workflows execute 9999 /path/to/file.yml -v /path/to/params.yml
    ```

    \b
    🆗 ***Response:***
    ```json
    {
        "executionId": "d6116a56-6b1d-41b4-a599-fb949f08863f"
    }
    ```

    🔗 [**API Documentation**](https://service.peak.ai/workflows/api-docs/index.htm#/Workflows/execute-workflow)
    """
    workflow_client: Workflow = ctx.obj["client"]

    response: Dict[str, str] = {}

    if file:
        body: Dict[str, Any] = helpers.template_handler(file=file, params_file=params_file, params=params)
        body = helpers.remove_unknown_args(body, workflow_client.execute_workflow)

        response = workflow_client.execute_workflow(workflow_id=workflow_id, **body)
    else:
        response = workflow_client.execute_workflow(workflow_id=workflow_id)

    console.print(response)


@app.command(short_help="List all available resources.")
def list_resources(ctx: typer.Context) -> None:
    """***List*** all available resources that can be used in workflow steps.

    \b
    📝 ***Example usage:***<br/>
    ```bash
    peak workflows list-resources
    ```

    \b
    🆗 ***Response:***
    ```json
    [
        {
            "cpu": 0.25,
            "gpu": null,
            "gpuMemory": null,
            "id": 21,
            "memory": 0.5
        }
    ]
    ```

    🔗 [**API Documentation**](https://service.peak.ai/workflows/api-docs/index.htm#/Resources/get-resources)
    """
    workflow_client: Workflow = ctx.obj["client"]
    response = workflow_client.list_resources()
    console.print(response)


@app.command(short_help="List default resources.")
def list_default_resources(ctx: typer.Context) -> None:
    """***List*** default resources for the workflows.

    \b
    📝 ***Example usage:***<br/>
    ```bash
    peak workflows list-default-resources
    ```

    \b
    🆗 ***Response:***
    ```json
    {
        "instanceTypeId": 21,
        "storage": "10GB"
    }
    ```

    🔗 [**API Documentation**](https://service.peak.ai/workflows/api-docs/index.htm#/Resources/get-default-resources)
    """
    workflow_client: Workflow = ctx.obj["client"]
    response = workflow_client.get_default_resource()
    console.print(response)


@app.command(short_help="List executions for the given workflow.")
def list_executions(
    ctx: typer.Context,
    workflow_id: int = _WORKFLOW_ID,
    date_from: Optional[str] = args.DATE_FROM,
    date_to: Optional[str] = args.DATE_TO,
    page_size: Optional[int] = args.PAGE_SIZE,
    page_number: Optional[int] = args.PAGE_NUMBER,
) -> None:
    """***List*** executions for the given workflow.

    \b
    📝 ***Example usage:***<br/>
    ```bash
    peak workflows list-executions 9999 --page-size 10 --page-number 1
    ```

    \b
    🆗 ***Response:***
    ```
    {
        "executions": [...],
        "pageSize": 25,
        "pageNumber": 1,
        "pageCount": 1,
        "workflowId": 11876
    }
    ```

    🔗 [**API Documentation**](https://service.peak.ai/workflows/api-docs/index.htm#/Executions/get-workflow-executions)
    """
    workflow_client: Workflow = ctx.obj["client"]
    response = workflow_client.list_executions(
        workflow_id=workflow_id,
        date_from=date_from,
        date_to=date_to,
        page_size=page_size,
        page_number=page_number,
        return_iterator=False,
    )
    console.print(response)
