# Copyright 2021-2023 Kolena Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License 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.
import dataclasses
import json
import os
import sys
from typing import Optional

import click

import kolena.workflow
from kolena._utils import repository
from kolena._utils.consts import KOLENA_TOKEN_ENV
from kolena._utils.state import _client_state
from kolena.errors import NotFoundError

KOLENA_COMMAND = "kolena"
CONTEXT_SETTINGS = dict(help_option_names=["--help", "-h"])

_shared_options = [
    click.option(
        "--api-token",
        "-t",
        required=False,
        type=str,
        default=os.environ.get(KOLENA_TOKEN_ENV),
        help="api token for client to authorize to kolena server; if not specified, the command will use value of "
        f"environment variable '{KOLENA_TOKEN_ENV}'",
    ),
]

_shared_evaluator_options = [click.option("--workflow", "-w", required=True)]


def add_options(*args):
    def _add_options(func):
        options = [x for n in args for x in n]
        for option in reversed(options):
            func = option(func)
        return func

    return _add_options


@click.group(name=KOLENA_COMMAND, context_settings=CONTEXT_SETTINGS)
@click.option("--base-url", "-b", required=False, type=str, hidden=True, help="Use alternate sdk-server url")
def base_command(base_url: Optional[str] = None) -> None:
    if base_url:
        _client_state.base_url = base_url


@base_command.group(name="evaluator")
def evaluator_command() -> None:
    ...


@base_command.group(name="repository")
def repository_command() -> None:
    ...


@evaluator_command.command(name="register", help="Attach a remote evaluator to a particular workflow")
@add_options(_shared_options)
@add_options(_shared_evaluator_options)
@click.option(
    "--evaluator-name",
    "-n",
    required=True,
    help="name of the evaluator; must match evaluator name generated by image",
)
@click.option(
    "--image",
    "-i",
    required=True,
    help=(
        "name of the fully qualified image that has already been uploaded "
        "(eg: quay.io/kolena/pedestrian_crossing:0.1.0);"
        "NOTE: this is not validated at registration time but will cause issues at runtime "
        "if set to an invalid value"
    ),
)
@click.option(
    "--secret",
    "-s",
    required=False,
    help=(
        "sensitive data required by evaluator as string; "
        "NOTE: content would be store securely in kolena platform "
        "and accessible to the evaluator through environment variable"
    ),
)
@click.option(
    "--aws-assume-role",
    "-r",
    required=False,
    help="the AWS role arn the evaluator would assume to access AWS APIs",
)
def evaluator_register(
    workflow: str,
    evaluator_name: str,
    image: str,
    secret: Optional[str] = None,
    aws_assume_role: Optional[str] = None,
    api_token: Optional[str] = None,
) -> None:
    kolena.initialize(api_token=api_token)
    evaluator = kolena.workflow.workflow.register_evaluator(workflow, evaluator_name, image, secret, aws_assume_role)
    print(f"Image {image} successfully registered for workflow {workflow} evaluator {evaluator_name}.")
    print(json.dumps(dataclasses.asdict(evaluator), indent=2))


@evaluator_command.command(name="list", help="List registered remote evaluators for a particular workflow")
@add_options(_shared_options)
@add_options(_shared_evaluator_options)
def evaluator_list(workflow: str, api_token: Optional[str] = None) -> None:
    kolena.initialize(api_token=api_token)
    evaluators = kolena.workflow.workflow.list_evaluators(workflow)
    for evaluator in evaluators:
        print(f"evaluator_name='{evaluator.name}', image='{evaluator.image}', created='{evaluator.created}'")


@evaluator_command.command(name="get", help="Get the latest remote evaluator of a workflow")
@add_options(_shared_options)
@add_options(_shared_evaluator_options)
@click.option(
    "--evaluator-name",
    "-n",
    required=True,
    help="name of the evaluator",
)
@click.option(
    "--include-secret/--no-include-secret",
    "-s",
    help="include sensitive data registered with the evaluator or not",
    default=False,
)
def evaluator_get(workflow: str, evaluator_name: str, include_secret: bool, api_token: Optional[str] = None) -> None:
    kolena.initialize(api_token=api_token)
    try:
        evaluator = kolena.workflow.workflow.get_evaluator(workflow, evaluator_name, include_secret)
    except NotFoundError:
        print("Error: workflow not found", file=sys.stderr)
        sys.exit(1)

    print("Evaluator:")
    print(json.dumps(dataclasses.asdict(evaluator), indent=2))


@repository_command.command(name="create", help="Create a docker repository on kolena docker registry")
@add_options(_shared_options)
@click.option(
    "--name",
    "-n",
    required=True,
    help="repository name",
)
def repository_create(name: str, api_token: Optional[str] = None) -> None:
    kolena.initialize(api_token=api_token)
    try:
        repository.create(name)
    except Exception as e:
        print(f"Failed to create repository: {e}", file=sys.stderr)
        sys.exit(0)

    print("repository created")


def run():
    base_command()
