# -*- coding: utf-8 -*-
# Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Module that handles tee-ing output to a file."""

from __future__ import print_function

import errno
import fcntl
import os
import multiprocessing
import select
import signal
import sys
import traceback

from chromite.lib import cros_build_lib
from chromite.lib import cros_logging as logging


# Max amount of data we're hold in the buffer at a given time.
_BUFSIZE = 1024


# Custom signal handlers so we can catch the exception and handle it.
class ToldToDie(Exception):
  """Exception thrown via signal handlers."""

  def __init__(self, signum):
    Exception.__init__(self, 'We received signal %i' % (signum,))


def _TeeProcessSignalHandler(signum, _frame):
  """TeeProcess custom signal handler.

  This is used to decide whether or not to kill our parent.
  """
  raise ToldToDie(signum)


def _output(line, output_files, complain):
  """Print line to output_files.

  Args:
    line: Line to print.
    output_files: List of files to print to.
    complain: Print a warning if we get EAGAIN errors. Only one error
              is printed per line.
  """
  for f in output_files:
    offset = 0
    while offset < len(line):
      select.select([], [f], [])
      try:
        offset += os.write(f.fileno(), line[offset:])
      except OSError as ex:
        if ex.errno == errno.EINTR:
          continue
        elif ex.errno != errno.EAGAIN:
          raise

      if offset < len(line) and complain:
        flags = fcntl.fcntl(f.fileno(), fcntl.F_GETFL, 0)
        if flags & os.O_NONBLOCK:
          warning = '\nWarning: %s/%d is non-blocking.\n' % (f.name,
                                                             f.fileno())
          _output(warning, output_files, False)

        warning = '\nWarning: Short write for %s/%d.\n' % (f.name, f.fileno())
        _output(warning, output_files, False)


def _tee(input_fd, output_files, complain):
  """Read data from |input_fd| and write to |output_files|."""
  while True:
    # We need to use os.read() directly because it will return to us when the
    # other side has flushed its output (and is shorter than _BUFSIZE).  If we
    # use python's file object helpers (like read() and readline()), it will
    # not return until either the full buffer is filled or a newline is hit.
    data = os.read(input_fd, _BUFSIZE)
    if not data:
      return
    _output(data, output_files, complain)


class _TeeProcess(multiprocessing.Process):
  """Replicate output to multiple file handles."""

  def __init__(self, output_filenames, complain, error_fd,
               master_pid):
    """Write to stdout and supplied filenames.

    Args:
      output_filenames: List of filenames to print to.
      complain: Print a warning if we get EAGAIN errors.
      error_fd: The fd to write exceptions/errors to during
        shutdown.
      master_pid: Pid to SIGTERM if we shutdown uncleanly.
    """

    self._reader_pipe, self.writer_pipe = os.pipe()
    self._output_filenames = output_filenames
    self._complain = complain
    # Dupe the fd on the offchance it's stdout/stderr,
    # which we screw with.
    # Not passing 3 argument (0) for unbuffered output because this is not
    # supported in Python 3 and there are issues in Python 2 -- see
    # https://bugs.python.org/issue17404.
    self._error_handle = os.fdopen(os.dup(error_fd), 'w')
    self.master_pid = master_pid
    multiprocessing.Process.__init__(self)

  def _CloseUnnecessaryFds(self):
    # For python2 we were relying on subprocess.MAXFD but that does not exist
    # in python3. However, the calculation below is how it was being computed.
    try:
      max_fd_value = os.sysconf('SC_OPEN_MAX')
    except ValueError:
      max_fd_value = 256
    preserve = set([1, 2, self._error_handle.fileno(), self._reader_pipe,
                    max_fd_value])
    preserve = iter(sorted(preserve))
    fd = 0
    while fd < max_fd_value:
      current_low = next(preserve)
      if fd != current_low:
        os.closerange(fd, current_low)
        fd = current_low
      fd += 1

  def run(self):
    """Main function for tee subprocess."""
    failed = True
    input_fd = None
    try:
      signal.signal(signal.SIGINT, _TeeProcessSignalHandler)
      signal.signal(signal.SIGTERM, _TeeProcessSignalHandler)

      # Cleanup every fd except for what we use.
      self._CloseUnnecessaryFds()

      # Read from the pipe.
      input_fd = self._reader_pipe

      # Create list of files to write to.
      # Not passing 3 argument (0) for unbuffered output because this is not
      # supported in Python 3 and there are issues in Python 2 -- see
      # https://bugs.python.org/issue17404.
      output_files = [os.fdopen(sys.stdout.fileno(), 'w')]
      for filename in self._output_filenames:
        output_files.append(open(filename, 'w'))

      # Send all data from the one input to all the outputs.
      _tee(input_fd, output_files, self._complain)
      failed = False
    except ToldToDie:
      failed = False
    except Exception as e:
      tb = traceback.format_exc()