# @Time    : 2021/7/20 14:11
# @Author  : Boyang
# @Site    : 
# @File    : waveletai_backend.py
# @Software: PyCharm
import itertools
import logging
import os
import posixpath
import re
import subprocess
import time

import docker
import mlflow
from docker.errors import BuildError
from docker.utils.json_stream import json_stream
from mlflow import tracking
from mlflow.projects import _wait_for
from mlflow.projects.backend.abstract_backend import AbstractBackend
from mlflow.projects.backend.local import _get_docker_command
from mlflow.projects.docker import _get_docker_image_uri, _create_docker_build_ctx
from mlflow.projects.submitted_run import LocalSubmittedRun
from mlflow.projects.utils import (
    fetch_and_validate_project,
    get_or_create_run,
    load_project,
    get_run_env_vars,
    get_databricks_env_vars,
    get_entry_point_command,
    MLFLOW_LOCAL_BACKEND_RUN_ID_CONFIG,
    MLFLOW_DOCKER_WORKDIR_PATH,
    PROJECT_USE_CONDA,
    PROJECT_DOCKER_ARGS,
    PROJECT_STORAGE_DIR,
)
from mlflow.utils.conda import get_conda_command, get_or_create_conda_env
from mlflow.utils.mlflow_tags import MLFLOW_PROJECT_ENV, MLFLOW_DOCKER_IMAGE_URI, MLFLOW_DOCKER_IMAGE_ID

from waveletai import envs

_logger = logging.getLogger(__name__)

_GENERATED_DOCKERFILE_NAME = "Dockerfile.mlflow-autogenerated"
_MLFLOW_DOCKER_TRACKING_DIR_PATH = "/mlflow/tmp/mlruns"
_PROJECT_TAR_ARCHIVE_NAME = "mlflow-project-docker-build-context"
_BACKEND_DEPENDENCE = "DEPENDENCE"
_REQUIREMENT_FILE_NAME = 'requirements.txt'


class WaveletAIBackend(AbstractBackend):
    def run(self, project_uri, entry_point, params, version, backend_config, tracking_uri, experiment_id):
        work_dir = fetch_and_validate_project(project_uri, version, entry_point, params)
        project = load_project(work_dir)
        if MLFLOW_LOCAL_BACKEND_RUN_ID_CONFIG in backend_config:
            run_id = backend_config[MLFLOW_LOCAL_BACKEND_RUN_ID_CONFIG]
        else:
            run_id = None
        active_run = get_or_create_run(
            run_id, project_uri, experiment_id, work_dir, version, entry_point, params
        )
        command_args = []
        command_separator = " "
        use_conda = backend_config[PROJECT_USE_CONDA]
        docker_args = backend_config[PROJECT_DOCKER_ARGS]
        storage_dir = backend_config[PROJECT_STORAGE_DIR]
        dependence = backend_config.get(_BACKEND_DEPENDENCE, {})
        # If a docker_env attribute is defined in MLproject then it takes precedence over conda yaml
        # environments, so the project will be executed inside a docker container.
        if project.docker_env:
            from mlflow.projects.docker import (
                validate_docker_env,
                validate_docker_installation,
            )

            tracking.MlflowClient().set_tag(active_run.info.run_id, MLFLOW_PROJECT_ENV, "docker")
            validate_docker_env(project)
            validate_docker_installation()
            image = _build_docker_image(
                work_dir=work_dir,
                repository_uri=project.name,
                base_image=project.docker_env.get("image"),
                run_id=active_run.info.run_id,
                dependence=dependence,
            )
            user_env_vars = project.docker_env.get("environment", [])
            if isinstance(user_env_vars, list):
                user_env_vars.append(envs.API_TOKEN)
            else:
                user_env_vars = [user_env_vars, envs.API_TOKEN]
            command_args += _get_docker_command(
                image=image,
                active_run=active_run,
                docker_args=docker_args,
                volumes=project.docker_env.get("volumes"),
                user_env_vars=user_env_vars,
            )
        # Synchronously create a conda environment (even though this may take some time)
        # to avoid failures due to multiple concurrent attempts to create the same conda env.
        elif use_conda:
            tracking.MlflowClient().set_tag(active_run.info.run_id, MLFLOW_PROJECT_ENV, "conda")
            command_separator = " && "
            conda_env_name = get_or_create_conda_env(project.conda_env_path)
            command_args += get_conda_command(conda_env_name)
        # In synchronous mode, run the entry point command in a blocking fashion, sending status
        # updates to the tracking server when finished. Note that the run state may not be
        # persisted to the tracking server if interrupted
        command_args += get_entry_point_command(project, entry_point, params, storage_dir)
        command_str = command_separator.join(command_args)

        return _run_entry_point(
            command_str, work_dir, experiment_id, run_id=active_run.info.run_id
        )


def _find_requirement_file(work_dir, requirements_filename=None):
    """
    在根目录下查询requirements.txt文件
    :return: 返回requirements文件的路径
    """
    requirements_filename = requirements_filename if requirements_filename else _REQUIREMENT_FILE_NAME
    for filename in os.listdir(work_dir):
        if filename.lower() == requirements_filename:
            return filename

    return None


def _build_docker_image(work_dir, repository_uri, base_image, run_id, dependence):
    """
    Build a docker image containing the project in `work_dir`, using the base image.
    """
    image_uri = _get_docker_image_uri(repository_uri=repository_uri, work_dir=work_dir)
    requirements_filename = dependence.get('requirements_filename')
    requirements_file = _find_requirement_file(work_dir, requirements_filename)
    if requirements_file:
        dockerfile = (
            "FROM {imagename}\n"
            "ADD {build_context_path}/{requirements_file} {workdir}\n"
            "RUN pip config set global.index-url https://mirrors.cloud.tencent.com/pypi/simple\n"
            "RUN if [ -f \"{workdir}{requirements_file}\" ] ; then pip install -r {workdir}{requirements_file}; fi\n"
            "COPY {build_context_path}/ {workdir}\n"
            "RUN pip install waveletai --upgrade\n"
            "WORKDIR {workdir}\n"
        ).format(
            imagename=base_image,
            build_context_path=_PROJECT_TAR_ARCHIVE_NAME,
            workdir=MLFLOW_DOCKER_WORKDIR_PATH,
            requirements_file=requirements_file,
        )
        # TODO support conda
    else:
        dockerfile = (
            "FROM {imagename}\n"
            "COPY {build_context_path}/ {workdir}\n"
            "WORKDIR {workdir}\n"
        ).format(
            imagename=base_image,
            build_context_path=_PROJECT_TAR_ARCHIVE_NAME,
            workdir=MLFLOW_DOCKER_WORKDIR_PATH,
        )

    build_ctx_path = _create_docker_build_ctx(work_dir, dockerfile)
    with open(build_ctx_path, "rb") as docker_build_ctx:
        _logger.info("=== Building docker image %s ===", image_uri)
        client = docker.APIClient(base_url='unix://var/run/docker.sock')
        resp = client.build(
            tag=image_uri,
            forcerm=True,
            dockerfile=posixpath.join(_PROJECT_TAR_ARCHIVE_NAME, _GENERATED_DOCKERFILE_NAME),
            fileobj=docker_build_ctx,
            custom_context=True,
            encoding="gzip",
        )
        last_event = None
        image_id = None
        result_stream, internal_stream = itertools.tee(json_stream(resp))
        last_output_time = None
        for chunk in internal_stream:
            if 'error' in chunk:
                _logger.info(chunk['error'])
                raise BuildError(chunk['error'], result_stream)
            if 'stream' in chunk:
                if chunk['stream'][-2:] == '[K':
                    continue
                match = re.search(r'[0-9]+([.]{1}[0-9]+){0,1} [k,M]B [0-9]+([.]{1}[0-9]+){0,1} [k,M]B/s',
                                  chunk['stream'])
                if match:
                    # 匹配到最开始下载
                    if last_output_time is None:
                        _logger.info(chunk['stream'])
                        last_output_time = int(time.time())
                        continue
                    # 匹配到结束
                    elif 'eta' not in chunk['stream']:
                        _logger.info(chunk['stream'])
                        last_output_time = None
                        continue
                    # 匹配距离上次输出已经超过三秒
                    elif int(time.time()) - last_output_time > 3:
                        _logger.info(chunk['stream'])
                        last_output_time = int(time.time())

                    continue

                _logger.info(chunk['stream'])
                match = re.search(r'(^Successfully built |sha256:)([0-9a-f]+)$', chunk['stream'])
                if match:
                    image_id = match.group(2)
                    break
            last_event = chunk
    try:
        os.remove(build_ctx_path)
    except Exception:
        _logger.info("Temporary docker context file %s was not deleted.", build_ctx_path)
    if image_id:
        _cli = docker.from_env()
        image = _cli.images.get(image_id)
        tracking.MlflowClient().set_tag(run_id, MLFLOW_DOCKER_IMAGE_URI, image_uri)
        tracking.MlflowClient().set_tag(run_id, MLFLOW_DOCKER_IMAGE_ID, image.id)
        return image

    raise BuildError(last_event or 'Unknown', result_stream)


def _get_pure_command(command):
    """
    清除系统敏感信息
    :param command:
    :return:
    """
    for c in re.findall(r'(-e .*? )', command):
        command = command.replace(c, '')

    return command


def _run_entry_point(command, work_dir, experiment_id, run_id):
    """
    Run an entry point command in a subprocess, returning a SubmittedRun that can be used to
    query the run's status.
    :param command: Entry point command to run
    :param work_dir: Working directory in which to run the command
    :param run_id: MLflow run ID associated with the entry point execution.
    """
    env = os.environ.copy()
    env.update(get_run_env_vars(run_id, experiment_id))
    env.update(get_databricks_env_vars(tracking_uri=mlflow.get_tracking_uri()))
    pure_command = _get_pure_command(command)
    _logger.info("=== Running command '%s' in run with ID '%s' === ", pure_command, run_id)
    # in case os name is not 'nt', we are not running on windows. It introduces
    # bash command otherwise.
    if os.name != "nt":
        process = subprocess.Popen(["bash", "-c", command], close_fds=True, cwd=work_dir, env=env)
    else:
        # process = subprocess.Popen(command, close_fds=True, cwd=work_dir, env=env)
        process = subprocess.Popen(["cmd", "/c", command], close_fds=True, cwd=work_dir, env=env)
    lsr = LocalSubmittedRun(run_id, process)
    _wait_for(lsr)

    return lsr
