# Copyright (c) YugaByte, 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.

"""
Utilities for running external commands.
"""

import os
import sys
import logging
import subprocess
import threading

from yugabyte_pycommon.text_manipulation import cmd_line_args_to_str, decode_utf8, trim_long_text
from yugabyte_pycommon.logging_util import is_verbose_mode


# Default number of lines to shorten long stdout/stderr to.
DEFAULT_MAX_LINES_TO_SHOW = 1000


class ProgramResult:
    """
    This represents the result of executing an external program.
    """
    def __init__(self, cmd_line, cmd_line_str, returncode, stdout, stderr, program_path,
                 invocation_details_str, max_lines_to_show, output_captured):
        self.cmd_line = cmd_line
        self.cmd_line_str = cmd_line_str
        self.returncode = returncode
        self.stdout = stdout
        self.stderr = stderr
        self.program_path = program_path
        self.invocation_details_str = invocation_details_str
        self.max_lines_to_show = max_lines_to_show
        self.output_captured = output_captured

        self._set_error_msg()

    def success(self):
        """
        :return: whether the external program exited with a success
        """
        return self.returncode == 0

    def failure(self):
        """
        :return: whether the external program exited with a failure
        """
        return self.returncode != 0

    def get_stdout_and_stderr_together(self):
        """
        :return: a string with user-friendly versions of stdout and stderr of the external program,
                 concatenated together.
        """
        stdout_and_stderr = (
            self.get_user_friendly_stdout_msg() + self.get_user_friendly_stderr_msg())
        if not stdout_and_stderr:
            stdout_and_stderr = "No stdout or stderr from command: " + self.invocation_details_str
        return stdout_and_stderr

    def print_output_to_stdout(self):
        """
        Print both stdout and stderr of the external program to the stdout.
        """
        sys.stdout.write(self.get_stdout_and_stderr_together())
        sys.stdout.flush()

    def _set_error_msg(self):
        if self.returncode == 0:
            self.error_msg = None
            return

        self.error_msg = "Non-zero exit code {} from {}.".format(
            self.returncode,
            self.invocation_details_str,
            self.returncode, cmd_line_args_to_str)

        self.error_msg += self.get_user_friendly_stdout_msg()
        self.error_msg += self.get_user_friendly_stderr_msg()
        self.error_msg = self.error_msg.rstrip()

    def _wrap_for_error_msg(self, stream_type):
        assert stream_type in ['output', 'error']
        if stream_type == 'output':
            value = self.stdout
        else:
            value = self.stderr
        if value is None or not value.strip():
            return ""
        value = value.rstrip()
        return "\nStandard {} from {}:\n{}\n(end of standard {})\n".format(
            stream_type, self.invocation_details_str,
            trim_long_text(value, self.max_lines_to_show),
            stream_type)

    def get_user_friendly_stdout_msg(self):
        """
        :return: a user-friendly version of the external program's standard output
        """
        return self._wrap_for_error_msg("output")

    def get_user_friendly_stderr_msg(self):
        """
        :return: a user-friendly version of the external program's standard error
        """
        return self._wrap_for_error_msg("error")

    def raise_error_if_failed(self):
        """
        This is useful for delayed handling of external program errors. Raises an error if the
        external program failed. Otherwise does nothing.
        """
        if self.failure():
            raise ExternalProgramError(self.error_msg, self)


class ExternalProgramError(Exception):
    def __init__(self, message, result):
        self.message = message
        self.result = result


class WorkDirContext:
    """
    Allows setting a working directory context for running external programs. The directory will
    be changed to the given directory on entering the block, and will be restored to the old
    directory on exit.

    .. code-block:: python

       with WorkDirContext('/tmp'):
           run_program('ls')
    """
    def __init__(self, work_dir):
        self.thread_local = threading.local()
        self.work_dir = work_dir

    def __enter__(self):
        self.thread_local.old_dir = os.getcwd()
        os.chdir(self.work_dir)

    def __exit__(self, exception_type, exception_value, traceback):
        os.chdir(self.thread_local.old_dir)


def run_program(args, error_ok=False, report_errors=None, capture_output=True,
                max_lines_to_show=DEFAULT_MAX_LINES_TO_SHOW, cwd=None, shell=None, **kwargs):
    """
    Run the given program identified by its argument list, and return a :py:class:`ProgramResult`
    object.

    :param error_ok: if this is true, we won't raise an exception in case the external program
                     fails.
    """
    if isinstance(args, str) and shell is None:
        shell = True

    if isinstance(args, str):
        # This is a special case, but very common.
        cmd_line_str = args
        args = [args]
    else:
        if isinstance(args, tuple):
            args = list(args)

        if isinstance(args, str):
            args = [args]

        def normalize_arg(arg):
            if isinstance(arg, int):
                return str(arg)
            return arg

        args = [normalize_arg(arg) for arg in args]

        cmd_line_str = cmd_line_args_to_str(args)

    invocation_details_str = "external program {{ %s }} running in '%s'" % (
            cmd_line_str, cwd or os.getcwd())

    if is_verbose_mode():
        logging.info("Running %s", invocation_details_str)

    try:
        output_redirection = subprocess.PIPE if capture_output else None
        program_subprocess = subprocess.Popen(
            args,
            stdout=output_redirection,
            stderr=output_redirection,
            shell=shell,
            cwd=cwd,
            **kwargs)
        program_stdout, program_stderr = program_subprocess.communicate()

    except OSError:
        logging.error("Failed to run %s", invocation_details_str)
        raise

    def cleanup_output(out_str):
        if out_str is None:
            return None
        return decode_utf8(out_str)

    clean_stdout = cleanup_output(program_stdout)
    clean_stderr = cleanup_output(program_stderr)

    result = ProgramResult(
        cmd_line=args,
        cmd_line_str=cmd_line_str,
        program_path=os.path.realpath(args[0]),
        returncode=program_subprocess.returncode,
        stdout=clean_stdout,
        stderr=clean_stderr,
        invocation_details_str=invocation_details_str,
        max_lines_to_show=max_lines_to_show,
        output_captured=capture_output)

    if program_subprocess.returncode != 0:
        if report_errors is None:
            report_errors = not error_ok
        if report_errors:
            logging.error(result.error_msg)
        if not error_ok:
            result.raise_error_if_failed()

    return result


def check_run_program(*args, **kwargs):
    """
    Similar to subprocess.check_call but using our run_program facility.
    """
    kwargs['capture_output'] = False
    kwargs['report_errors'] = True
    run_program(*args, **kwargs)
    return 0


def program_fails_no_log(args, **kwargs):
    """
    Run the given program, and returns if it failed. Does not log anything in case of success
    or failure.

    :param args: command line arguments or a single string to run as a shell command
    :param kwargs: additional keyword arguments for subprocess.Popen
    :return: ``True`` if the program succeeded
    """
    return run_program(args, error_ok=True, report_errors=False, **kwargs).failure()


def program_succeeds_no_log(args, **kwargs):
    """
    Run the given program, and returns True if it succeeded. Does not log anything in case of
    success or failure.

    :param args: command line arguments or a single string to run as a shell command
    :param kwargs: additional keyword arguments for subprocess.Popen
    :return: ``True`` if the program failed
    """
    return run_program(args, error_ok=True, report_errors=False, **kwargs).success()


def program_succeeds_empty_output(args, **kwargs):
    """
    Runs a program that is not expected to produce any output.

    :param args: command line arguments or a single string to run as a shell command
    :param kwargs: additional keyword arguments for subprocess.Popen
    :raises ExternalProgramError: if the program succeeds but produces extra output
    :return: ``True`` if the program succeeds and does not produce any output
    """
    result = run_program(args, error_ok=True, report_errors=False, **kwargs)
    if result.failure():
        return False

    if result.stdout.strip():
        error_msg = "Unexpected output in case of success. " + result.get_user_friendly_stdout_msg()
        logging.error(error_msg)
        raise ExternalProgramError(error_msg, result)

    return True
