import logging
import os
import re
import signal
import socket

from datetime import datetime

from psutil import Process

from fixtures import (
    Fixture,
    TempDir,
)

from twisted.internet.protocol import (
    Factory,
    Protocol,
    ProcessProtocol,
)
from twisted.internet.defer import (
    Deferred,
    inlineCallbacks,
)
from twisted.internet.task import LoopingCall
from twisted.internet.endpoints import TCP4ClientEndpoint
from twisted.internet.error import (
    ConnectionRefusedError,
    ConnectingCancelledError,
)
from twisted.protocols.basic import LineOnlyReceiver

from txfixtures._twisted.backports.defer import addTimeout


TIMEOUT = 15

# Some processes (like mongodb) use an abbreviated code for level names. We
# keep a mapping for transparently convert between them and standard Python
# level names.
SHORT_LEVELS = {
    "C": "CRITICAL",
    "E": "ERROR",
    "W": "WARNING",
    "I": "INFO",
    "D": "DEBUG",
}


class Service(Fixture):
    """Spawn, control and monitor a background service."""

    def __init__(self, reactor, command, args=None, env=None, timeout=None):
        super(Service, self).__init__()

        self.reactor = reactor
        self.command = command
        self.args = args or []

        if env is None:
            env = os.environ
        self.env = _encodeDictValues(env)

        self._reactor = reactor.reactor
        self._eventTriggerID = None
        self._data_dirs = []

        if timeout is None:
            timeout = TIMEOUT
        self.protocol = ServiceProtocol(reactor=self._reactor, timeout=timeout)

    def reset(self):
        if self.protocol.terminated.called:
            raise RuntimeError("Service died")

    def addDataDir(self):
        data_dir = self.useFixture(TempDir())
        self._data_dirs.append(data_dir.path)

    def expectOutput(self, data):
        self.protocol.expectedOutput = data

    def expectPort(self, port):
        self.protocol.expectedPort = port

    def setOutputFormat(self, outFormat):
        self.protocol.parser.pattern = outFormat

    def allocatePort(self):
        """Allocate an unused port.

        This method can be used by subclasses to allocate a random ports for
        the service they spawn.

        There is a small race condition here (between the time we allocate the
        port, and the time it actually gets used), but for the purposes for
        which this method gets used it isn't a problem in practice.
        """
        sock = socket.socket()
        try:
            sock.bind(("localhost", 0))
            _, port = sock.getsockname()
            return port
        finally:
            sock.close()

    def _setUp(self):
        logging.info("Spawning service process %s", self.command)
        self.addCleanup(self._stop)
        self._callFromThread(self._start)

    def _extraArgs(self):
        return []

    @property
    def _name(self):
        return os.path.basename(self.command)

    @inlineCallbacks
    def _start(self):
        self.protocol.parser.setServiceName(self._name)
        args = [self.command] + self.args + self._extraArgs()
        self._reactor.spawnProcess(
            self.protocol, self.command, args=args, env=self.env)

        # This cleanup handler will be triggered in case of SIGTERM and SIGINT,
        # when the reactor will initiate an unexpected shutdown sequence.
        self._eventTriggerID = self._reactor.addSystemEventTrigger(
            "before", "shutdown", self._terminateProcess)

        yield self.protocol.ready

    def _stop(self):
        logging.info("Stopping service process %s", self.command)

        try:
            self._callFromThread(self._terminateProcess)
        except:
            if self.protocol.transport.pid:
                # In case something goes wrong let's try our best to not leave
                # running processes around.
                logging.info(
                    "Service process didn't terminate, trying to kill it")
                process = Process(self.protocol.transport.pid)
                process.kill()
                process.wait(timeout=1)

    def _callFromThread(self, f):
        # Set an additional timeout for the callFromThread call itself. We
        # want this timeout to be greater than the 'ready' deferred timeout
        # set in _start(), so if the reactor thread is hung or dies we still
        # properly timeout.
        self.reactor.call(self.protocol.timeout + 1, f)

    @inlineCallbacks
    def _terminateProcess(self):
        if self._eventTriggerID:
            # Clear the shutdown event trigger, since we're going to cleanup
            # normally.
            self._reactor.removeSystemEventTrigger(self._eventTriggerID)
        if self.protocol.transport.pid:
            logging.info("Sending SIGTERM to service process '%s'", self._name)
            self.protocol.transport.signalProcess(signal.SIGTERM)
            logging.info("Waiting for service process to terminate")
            yield self.protocol.terminated


class ServiceProtocol(ProcessProtocol):
    """Start and stop a background service process.

    This :class:`~twisted.internet.protocol.ProcessProtocol` manages the start
    up and termination phases of a background service process. The process is
    considered 'running' when it has stayed up for at least 0.1 seconds (or any
    other non default value which `minUptime` is set too), and optionally when
    it has emitted a certain string and/or it has started listening to a
    certain port.
    """

    #: The service process must stay up at least this amount of seconds, before
    #: it's considered running. This allows to catch common issues like the
    #: service process executable not being in PATH or not being executable.
    minUptime = 0.2

    def __init__(self, reactor, parser=None, timeout=TIMEOUT):
        self.reactor = reactor
        self.parser = parser or ServiceOutputParser()

        #: Maximum amount of seconds to wait for the service to be ready. After
        #: that, the 'ready' deferred will errback with a TimeoutError.
        self.timeout = timeout

        #: Optional text that we expect the process to emit in standard output
        #: before we consider it ready.
        self.expectedOutput = None

        #: Optional port number that we expect the service process to listen,
        #: before we consider it ready.
        self.expectedPort = None

        #: Deferred that will fire when the service is considered ready, i.e.
        #: it has stayed up for at least minUptime seconds, has produced the
        #: expected output (if any), and is listening to the expected port (if
        #: any). Upon cancellation, any waiting activity will be stopped.
        self.ready = Deferred(lambda _: self._stopWaitingForReady())

        #: Deferred that will fire when the service has fully terminated, i.e.
        #: it has exited and we parent process have read any outstanding data
        #: in the pipes and have closed them.
        self.terminated = Deferred()

        # Delayed call that gets started right after the process has been
        # spawned. Its purpose is to make the protocol "sleep" for a minUptime
        # seconds (typically 0.1 seconds): if the process exits before this
        # little time has elapsed, an error gets raised.
        self._minUptimeCall = None

        # Deferred that will be fired when the process emits the expected
        # output (if any).
        self._expectedOutputReady = Deferred()

        # A LoopingCall instance that will periodically try to open the port
        # that the process is supposed to start listening to.
        self._probePortLoop = None

        # A connector as returned by TCP4ClientEndpoint.connect() that can be
        # used to abort an ongoing connection attempt as performed by the
        # port probe loop.
        self._probePortAttempt = None

    def connectionMade(self):
        # Called (indirectly) by `spawnProcess` after the `os.fork` call has
        # succeeded.

        logging.info("Service process spawned")

        # XXX Replace with self.ready.addTimeout once support for Twisted<16.5
        #     gets dropped.
        addTimeout(self.ready, self.timeout, self.reactor)

        # The twisted.protocols.basic.LineOnlyReceiver class expects to know
        # when the transport is disconnecting.
        self.disconnecting = False

        # Let's see if the process stays running for at least
        self._minUptimeCall = self.reactor.callLater(
            self.minUptime, self._minUptimeElapsed)

        if self.expectedOutput:
            # From this point on, be prepared to receive the expected output at
            # any time.
            self.parser.whenLineContains(
                self.expectedOutput, self._expectedOutputReceived)
        else:
            # There's no output we expect, so we fire this Deferred right away.
            # When _minUptimeElapsed will be called, the callback that gets
            # attached to this Deferred will fire synchronously.
            self._expectedOutputReady.callback(None)

        self.parser.makeConnection(self)

    def outReceived(self, data):
        # Called when we receive data from the standard output of the service.
        self.parser.dataReceived(data)

    errReceived = outReceived

    def processExited(self, reason):
        # Called when the service process exited.

        logging.info("Service process exited: %s", reason.getErrorMessage())

        # If we did not reach the 'ready' state yet, the let's fire the 'ready'
        # Deferred with an error.
        if not self.ready.called:
            self._stopWaitingForReady(reason)

    def processEnded(self, reason):
        # Called when the process has been reaped.
        logging.info("Service process reaped")
        self.terminated.callback(None)

    def _minUptimeElapsed(self):
        """
        Called if the process didn't exit in the first `minUptime` seconds
        after having been spawned.
        """
        logging.info("Service process alive for %.1f seconds", self.minUptime)

        # Now wait for the expected output and then start polling the port
        # we expect the service to listen to (if there's no expected output
        # and/or no expected port, these deferreds will fire synchronously).
        if self.expectedPort:
            self._expectedOutputReady.addCallback(self._startProbePortLoop)
        self._expectedOutputReady.addCallback(self._maybeFireReady)

    def _expectedOutputReceived(self):
        """
        Called after `_minUptimeElapsed` and the service process has emitted
        the expected output string.
        """
        # Let's fire the relevant deferred, so we can move forward to polling
        # the expected port, or declaring the service as ready (if there's no
        # expected port).
        logging.info("Service process emitted '%s'", self.expectedOutput)
        self._expectedOutputReady.callback(None)

    @inlineCallbacks
    def _startProbePortLoop(self, _):
        """
        Called when the service process has stayed up for at least `minUptime`
        seconds and it has emitted the expected output string (or there was no
        expected output string at all).
        """

        self._probePortLoop = LoopingCall(self._probePort)
        self._probePortLoop.clock = self.reactor

        # The LoopingCall.start() method returns a deferred that will fire
        # when the loop stops, i.e. when we successfully probe the port.
        yield self._probePortLoop.start(0.1)

    @inlineCallbacks
    def _probePort(self):
        """Perform a single attempt to connect to the expected port.

        If the probe succeeds the probe loop will be stoped.

        If the probe fails with a connection error, we'll just return
        gracefully (we'll be invoked again at the next loop iteration).
        """
        logging.info("Polling service port '%s'", self.expectedPort)

        endpoint = TCP4ClientEndpoint(
            self.reactor, "localhost", self.expectedPort)

        try:
            factory = Factory()
            factory.protocol = Protocol
            self._probePortAttempt = endpoint.connect(factory)
            yield self._probePortAttempt
        except ConnectionRefusedError as error:
            logging.info("Service port probe failed: %s", error)
        except ConnectingCancelledError as error:
            # This happens if _stopWaitingForReady gets called while we are
            # waiting for the enpoint connect() to succeed or fail.
            logging.info("Service port probe cancelled: %s", error)
        else:
            if self._probePortLoop.running:
                self._probePortLoop.stop()
                logging.info("Service opened port %d", self.expectedPort)
        finally:
            self._probePortAttempt = None

    def _maybeFireReady(self, result):
        """Fire the 'ready' deferred, unless we're aborting the startup.

        If the startup sequence is aborting (either because the `ready`
        deferred was cancelled by user code, or because the process died and
        `processExited` was called), this will just be a no-op, as we rely
        on the aborting code to errback the `ready` deferred.
        """
        if not self.disconnecting:
            logging.info("Service process ready")
            self.ready.callback(result)

    def _stopWaitingForReady(self, reason=None):
        """
        Stop any delayed call or activity associated with the initial waiting
        for the service to be ready.

        If `reason` is passed, the `ready` deferred will errback with the given
        failure.
        """
        # This will prevent the ServiceOutputParser protocol from firing any
        # further lineReceived event, so we don't fire _expectedOutputReady.
        #
        # It will also prevent _maybeFireReady from firing the 'ready'
        # deferred, since we want to do it ourselves with the given reason (if
        # any).
        self.disconnecting = True

        message = None

        if self._minUptimeCall.active():
            self._minUptimeCall.cancel()
            message = "minimum uptime not yet elapsed"

        elif self.expectedOutput and not self._expectedOutputReady.called:
            message = "expected output not yet received"

        elif self.expectedPort:
            if self._probePortAttempt:
                self._probePortAttempt.cancel()

            self._probePortLoop.stop()
            message = "expected port not yet open"

        # We can safely assume that one of the conditions above is holding,
        # because otherwise the 'ready' deferred would have already fired. In
        # any case let's put an explicit assertion here for good measure.
        assert message, "Unexpected protocol state while cancelling wait"

        logging.info(
            "Give up waiting for the service to be ready: %s", message)

        if reason:
            self.ready.callback(reason)


class ServiceOutputParser(LineOnlyReceiver):
    """
    Parse the standard output stream of a service and forward it to the Python
    logging system.

    The stream is assumed to be a UTF-8 sequence of lines each delimited by
    a (configurable) delimiter character.

    Each received line is tested against the RegEx pattern provided in the
    constructor. If a match is found, a :class:`~logging.LogRecord` is built
    using the information from the groups of the match object, otherwise
    default values will be used.

    The record is then passed to the :class:`~logging.Logger` provided in the
    constructor.

    Match objects that result from the RegEx pattern are supposed to provide
    groups named after the substitutions below.
    """

    #: The delimiter character identifying the end of a line.
    delimiter = b"\n"

    #: Substitutions for commonly used groups in line match patterns. For
    #: example, this allows you to use "{Y}-{m}-{S}" as pattern snippet, as
    #  opposed to an explicit "(?P<Y>\d{4})-(?P<m>\d{2})-(?P<d>\d{2})".
    substitutions = {
        "Y": "(?P<Y>\d{4})",
        "m": "(?P<m>\d{2})",
        "d": "(?P<d>\d{2})",
        "H": "(?P<H>\d{2})",
        "M": "(?P<M>\d{2})",
        "S": "(?P<S>\d{2})",
        "msecs": "(?P<msecs>\d{3})",
        "levelname": "(?P<levelname>[a-zA-Z]+)",
        "name": "(?P<name>.+)",
        "message": "(?P<message>.+)",
    }

    #: A string identifying the service whose output is being
    #: parsed. It will be attached as 'service' attribute to all log
    #: records emitted.
    service = ""

    def __init__(self, logger=None, pattern=None):
        self.pattern = pattern or "{message}"
        self.logger = logger or logging.getLogger("")
        self._callbacks = {}

    def setServiceName(self, name):
        self.service = name

    def whenLineContains(self, text, callback):
        """Fire the given callback when a line contains the given text.

        The callback will be fired only once when and if a match is found.
        """
        self._callbacks[text] = callback

    def lineReceived(self, line):
        """Foward the received line to the Python logging system."""
        message = line.decode("utf-8")
        params = {
            "levelname": "NOTSET",
            "levelno": 0,
            "msg": message,
            "processName": self.service,
        }

        match = re.match(self.pattern.format(**self.substitutions), message)
        if match:
            params.update(self._getLogRecordParamsForMatch(match))

        record = logging.makeLogRecord(params)
        self.logger.handle(record)

        for text in list(self._callbacks.keys()):
            if text in record.msg:
                self._callbacks.pop(text)()

    def lineLengthExceeded(self, line):
        """Simply truncate the line."""
        self.lineReceived(line[:self.MAX_LENGTH])

    def _getLogRecordParamsForMatch(self, match):
        """
        Use the given `match` regex object to create a dict of parameters
        to be passed to `logging.makeLogRecord`.

        This method will try to use all the information extracted by the
        match. If some of it is missing or incomplete, it will be discarded.
        """
        groups = _filterNoneValues(match.groupdict())
        params = {
            "name": groups.get("name"),
            "msg": groups.get("message"),
        }

        if "levelname" in groups:
            levelname = groups["levelname"].upper()
            if len(levelname) == 1:
                levelname = SHORT_LEVELS.get(levelname, "INFO")
            params["levelname"] = levelname
            params["levelno"] = logging.getLevelName(params["levelname"])

        # Only set creation time if all date-related groups are there.
        if set(groups.keys()).issuperset({"Y", "m", "d", "H", "M", "S"}):
            params["created"] = float(datetime(
                int(groups["Y"]),
                int(groups["m"]),
                int(groups["d"]),
                int(groups["H"]),
                int(groups["M"]),
                int(groups["S"]),
            ).strftime("%s"))

        if "msecs" in groups:
            params["msecs"] = float(groups["msecs"])

        return params


def _filterNoneValues(d):
    """
    Return a dict which is the same as `d`, except for keys with None values,
    which get discarded.
    """
    return dict([(k, v) for k, v in d.items() if v is not None])


def _encodeDictValues(d):
    """
    Return a dict whose unicode values get UTF-8 encoded to bytes.
    """
    return dict(
        [(_maybeEncode(k), _maybeEncode(v))
         for k, v in d.items() if v is not None])


def _maybeEncode(x):
    """
    If x is a string, encode it to bytes using UTF-8.
    """
    if isinstance(x, str):
        x = x.encode("utf-8")
    return x
