#
# 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 errno
import getpass
import os
import signal
import socket
import struct
import subprocess
import sys
import threading
import time

from mesos.interface import mesos_pb2
from twitter.common import log
from twitter.common.dirutil import safe_mkdtemp
from twitter.common.log.options import LogOptions
from twitter.common.quantity import Amount, Time

from apache.thermos.common.statuses import (
    INTERNAL_ERROR,
    INVALID_TASK,
    TERMINAL_TASK,
    UNKNOWN_ERROR,
    UNKNOWN_USER
)
from apache.thermos.config.loader import ThermosTaskWrapper
from apache.thermos.core import runner as core
from apache.thermos.monitoring.monitor import TaskMonitor

from .common.status_checker import StatusResult
from .common.task_info import mesos_task_instance_from_assigned_task, resolve_ports
from .common.task_runner import TaskError, TaskRunner, TaskRunnerProvider
from .http_lifecycle import HttpLifecycleManager

from gen.apache.thermos.ttypes import TaskState


class ThermosTaskRunner(TaskRunner):
  ESCALATION_WAIT = Amount(5, Time.SECONDS)
  EXIT_STATE_MAP = {
      TaskState.ACTIVE: StatusResult('Runner died while task was active.', mesos_pb2.TASK_LOST),
      TaskState.FAILED: StatusResult('Task failed.', mesos_pb2.TASK_FAILED),
      TaskState.KILLED: StatusResult('Task killed.', mesos_pb2.TASK_KILLED),
      TaskState.LOST: StatusResult('Task lost.', mesos_pb2.TASK_LOST),
      TaskState.SUCCESS: StatusResult('Task finished.', mesos_pb2.TASK_FINISHED),
  }
  MAX_WAIT = Amount(1, Time.MINUTES)
  PEX_NAME = 'thermos_runner.pex'
  POLL_INTERVAL = Amount(500, Time.MILLISECONDS)
  THERMOS_PREEMPTION_WAIT = Amount(1, Time.MINUTES)

  def __init__(self,
               runner_pex,
               task_id,
               task,
               role,
               portmap,
               sandbox,
               checkpoint_root,
               artifact_dir=None,
               clock=time,
               hostname=None):
    """
      runner_pex       location of the thermos_runner pex that this task runner should use
      task_id          task_id assigned by scheduler
      task             thermos pystachio Task object
      role             role to run the task under
      portmap          { name => port } dictionary
      sandbox          the sandbox object
      checkpoint_root  the checkpoint root for the thermos runner
      artifact_dir     scratch space for the thermos runner (basically cwd of thermos.pex)
      clock            clock
    """
    self._runner_pex = runner_pex
    self._task_id = task_id
    self._task = task
    self._popen, self._popen_signal, self._popen_rc = None, None, None
    self._monitor = None
    self._status = None
    self._ports = portmap
    self._root = sandbox.root
    self._checkpoint_root = checkpoint_root
    self._enable_chroot = sandbox.chrooted
    self._role = role
    self._clock = clock
    self._artifact_dir = artifact_dir or safe_mkdtemp()
    self._hostname = hostname or socket.gethostname()

    # wait events
    self._dead = threading.Event()
    self._kill_signal = threading.Event()
    self.forking = threading.Event()
    self.forked = threading.Event()

    try:
      with open(os.path.join(self._artifact_dir, 'task.json'), 'w') as fp:
        self._task_filename = fp.name
        ThermosTaskWrapper(self._task).to_file(self._task_filename)
    except ThermosTaskWrapper.InvalidTask as e:
      raise TaskError('Failed to load task: %s' % e)

  @property
  def artifact_dir(self):
    return self._artifact_dir

  def task_state(self):
    return self._monitor.task_state() if self._monitor else None

  @classmethod
  def _decode_status(cls, status):
    # Per os.waitpid documentation, status is:
    #   a 16-bit number, whose low byte is the signal number that killed the
    #   process, and whose high byte is the exit status (if the signal
    #   number is zero); the high bit of the low byte is set if a core file
    #   was produced.
    exit_signal, exit_status = struct.unpack('bb', struct.pack('H', status))
    return exit_signal & 0x7F, exit_status  # strip the signal high bit

  @property
  def is_alive(self):
    """
      Is the process underlying the Thermos task runner alive?
    """
    if not self._popen:
      return False

    if self._dead.is_set():
      return False

    # N.B. You cannot mix this code and any code that relies upon os.wait
    # mechanisms with blanket child process collection.  One example is the
    # Thermos task runner which calls os.wait4 -- without refactoring, you
    # should not mix a Thermos task runner in the same process as this
    # thread.
    try:
      pid, status = os.waitpid(self._popen.pid, os.WNOHANG)
      if pid == 0:
        return True
      else:
        self._popen_signal, self._popen_rc = self._decode_status(status)
        log.info('Detected runner termination: pid=%s, signal=%s, rc=%s' % (
            pid, self._popen_signal, self._popen_rc))
    except OSError as e:
      log.error('is_alive got OSError: %s' % e)
      if e.errno != errno.ECHILD:
        raise

    self._dead.set()
    return False

  def compute_status(self):