#!/usr/bin/env python
# -*- coding: utf-8; -*-

# Copyright (c) 2022 Oracle and/or its affiliates.
# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
import os

import click

from ads.opctl.utils import publish_image as publish_image_cmd
from ads.opctl.utils import suppress_traceback
from ads.opctl.distributed.cmds import update_ini, docker_build_cmd
from ads.opctl.distributed.cmds import (
    initialize_workspace,
    show_config_info,
    cancel_distributed_run,
)
from configparser import ConfigParser


@click.group("distributed-training")
@click.help_option("--help", "-h")
def commands():
    pass


@commands.command()
@click.option(
    "--framework",
    "-f",
    help="Distributed training framework type",
    type=click.Choice(
        ["dask", "horovod-tensorflow", "horovod-pytorch", "pytorch", "tensorflow"]
    ),
    default=None,
    required=True,
)
@click.option(
    "--version",
    help="version number of the framework resource",
    default="v1",
    required=False,
)
@click.option("--debug", "-d", help="set debug mode", is_flag=True, default=False)
@click.help_option("--help", "-h")
def init(framework, version, debug):
    """
    Initialize workspace to build distributed framework compatible container images.
    Supported frameworks are - dask, horovod-tensorflow, horovod-pytorch pytorch and
    tensorflow.
    """
    initialize_workspace(
        framework,
        version,
    )


@commands.command()
@click.option(
    "--job-id",
    "-j",
    help="Job OCID of the cluster",
    default=None,
    required=False,
)
@click.option(
    "--work-dir",
    "-f",
    help="Working directory provided when the cluster was created.",
    default=None,
    required=False,
)
@click.option(
    "--file",
    "-f",
    help="Yaml file with cluster run information",
    default=None,
    required=False,
)
@click.option(
    "--worker-info", help="Show worker information", is_flag=True, default=False
)
@click.option("--debug", "-d", help="set debug mode", is_flag=True, default=False)
@click.option(
    "--oci-config",
    help="OCI Config file location",
    default=None,
    required=False,
)
@click.option(
    "--oci-profile",
    help="oci config profile",
    default=None,
    required=False,
)
def show_config(job_id, work_dir, file, worker_info, debug, **kwargs):
    """
    Displays the config json files create by the main node and the worker nodes. By default only the config generated by the Main node is displayed.
    """
    if not (job_id and work_dir) and not file:
        print("Either job ocid or yaml file with `jobId` key is required")
        return
    show_config_info(
        job_id,
        work_dir,
        file,
        worker_info,
        **kwargs,
    )


@commands.command()
@click.option(
    "--job-id",
    "-j",
    help="Job OCID of the cluster",
    default=None,
    required=False,
)
@click.option(
    "--file",
    "-f",
    help="Yaml file with cluster run information",
    default=None,
    required=False,
)
@click.option("--debug", "-d", help="set debug mode", is_flag=True, default=False)
@click.option(
    "--oci-config",
    help="OCI Config file location",
    default=None,
    required=False,
)
@click.option(
    "--oci-profile",
    help="oci config profile",
    default=None,
    required=False,
)
def cancel(job_id, file, debug, **kwargs):
    """Terminate the running cluster. Use Job ID of the main jobrun or use the file which has the yaml output displayed when cluster was created using ``ads opctl run -f <yaml>``"""
    if not (job_id or file):
        print("Either job ocid or yaml file with `jobId` key is required")
        return
    cancel_distributed_run(job_id, file, **kwargs)


@commands.command()
@click.option("--tag", "-t", help="tag of image", required=False, default=None)
@click.option(
    "--registry", "-reg", help="registry to push to", required=False, default=None
)
@click.option(
    "--dockerfile",
    "-df",
    help="relative path to Dockerfile",
    required=False,
    default=None,
)
@click.option("--debug", "-d", help="set debug mode", is_flag=True, default=None)
@click.option(
    "--push", "-p", help="push docker image to registry", is_flag=True, default=None
)
@click.option(
    "--source-folder",
    "-s",
    help="source scripts folder that will be mounted during a local run",
    required=False,
    default=None,
)
def build_image(tag, registry, dockerfile, source_folder, **kwargs):
    debug = kwargs.pop("debug")
    push = kwargs["push"]
    ini = update_ini(
        tag, registry, dockerfile, source_folder, config=None, nobuild=False
    )
    img_name = ini.get("main", "registry") + ":" + ini.get("main", "tag")
    docker_build_cmd(ini)
    if push:
        print("pushing docker image")
        suppress_traceback(debug)(publish_image_cmd)(img_name)


@commands.command()
@click.argument("image", default="None")
@click.help_option("--help", "-h")
@click.option("--debug", "-d", help="set debug mode", is_flag=True, default=False)
def publish_image(**kwargs):
    debug = kwargs.pop("debug")
    if kwargs["image"] == "None":
        config = ConfigParser()
        if os.path.isfile("config.ini"):
            config.read("config.ini")
            img_name = config.get("main", "registry") + ":" + config.get("main", "tag")
        else:
            raise ValueError("Arg image is missing")
    else:
        img_name = kwargs["image"]
    suppress_traceback(debug)(publish_image_cmd)(img_name)
