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

import cPickle as pickle
import errno
import json
import logging
import os
import socket
import SocketServer
import sys
import tempfile
import threading
import time
import traceback
import types
from Queue import Queue

import factory_common
from autotest_lib.client.cros import factory
from autotest_lib.client.cros.factory.unicode_to_string import UnicodeToString


# Environment variable storing the path to the endpoint.
CROS_FACTORY_EVENT = 'CROS_FACTORY_EVENT'

# Maximum allowed size for messages.  If messages are bigger than this, they
# will be truncated by the seqpacket sockets.
_MAX_MESSAGE_SIZE = 65535


def json_default_repr(obj):
    '''Converts an object into a suitable representation for
    JSON-ification.

    If obj is an object, this returns a dict with all properties
    not beginning in '_'.  Otherwise, the original object is
    returned.
    '''
    if isinstance(obj, object):
        return dict([(k,v) for k, v in obj.__dict__.iteritems()
                     if k[0] != "_"])
    else:
        return obj


class Event(object):
    '''
    An event object that may be written to the event server.

    E.g.:

        event = Event(Event.Type.STATE_CHANGE,
                      test='foo.bar',
                      state=TestState(...))
    '''
    Type = type('Event.Type', (), {
            # The state of a test has changed.
            'STATE_CHANGE':          'goofy:state_change',
            # The UI has come up.
            'UI_READY':              'goofy:ui_ready',
            # Tells goofy to switch to a new test.
            'SWITCH_TEST':           'goofy:switch_test',
            # Tells goofy to rotate visibility to the next active test.
            'SHOW_NEXT_ACTIVE_TEST': 'goofy:show_next_active_test',
            # Tells goofy to show a particular test.
            'SET_VISIBLE_TEST':      'goofy:set_visible_test',
            # Tells goofy to clear all state and restart testing.
            'RESTART_TESTS':  'goofy:restart_tests',
            # Tells goofy to run all tests that haven't been run yet.
            'AUTO_RUN': 'goofy:auto_run',
            # Tells goofy to set all failed tests' state to untested and re-run.
            'RE_RUN_FAILED': 'goofy:re_run_failed',
            # Tells goofy to go to the review screen.
            'REVIEW': 'goofy:review',
            # Tells the UI about a single new line in the log.
            'LOG': 'goofy:log',
            # A hello message to a new WebSocket.  Contains a 'uuid' parameter
            # identification the particular invocation of the server.
            'HELLO': 'goofy:hello',
            # A keepalive message from the UI.  Contains a 'uuid' parameter
            # containing the same 'uuid' value received when the client received
            # its HELLO.
            'KEEPALIVE': 'goofy:keepalive',
            # Sets the UI in the test pane.
            'SET_HTML': 'goofy:set_html',
            # Runs JavaScript in the test pane.
            'RUN_JS': 'goofy:run_js',
            # Calls a JavaScript function in the test pane.
            'CALL_JS_FUNCTION': 'goofy:call_js_function',
            # Event from a test UI.
            'TEST_UI_EVENT': 'goofy:test_ui_event',
            # Message from the test UI that it has finished.
            'END_TEST': 'goofy:end_test',
            # Message to tell the test UI to destroy itself.
            'DESTROY_TEST': 'goofy:destroy_test',
            # Message telling Goofy should re-read system info.
            'UPDATE_SYSTEM_INFO': 'goofy:update_system_info',
            # Message containing new system info from Goofy.
            'SYSTEM_INFO': 'goofy:system_info',
            # Tells Goofy to stop all tests and update factory
            # software.
            'UPDATE_FACTORY': 'goofy:update_factory',
            # Tells Goofy to stop all tests.
            'STOP': 'goofy:stop',
            })

    def __init__(self, type, **kw):  # pylint: disable=W0622
        self.type = type
        self.timestamp = time.time()
        for k, v in kw.iteritems():
            setattr(self, k, v)

    def __repr__(self):
        return factory.std_repr(
            self,
            extra=[
                'type=%s' % self.type,
                'timestamp=%s' % time.ctime(self.timestamp)],
            excluded_keys=['type', 'timestamp'])

    def to_json(self):
        return json.dumps(self, default=json_default_repr)

    @staticmethod
    def from_json(encoded_event):
        kw = UnicodeToString(json.loads(encoded_event))
        type = kw.pop('type')
        return Event(type=type, **kw)

_unique_id_lock = threading.Lock()
_unique_id = 1
def get_unique_id():
    global _unique_id
    with _unique_id_lock:
        ret = _unique_id
        _unique_id += 1
    return ret


class EventServerRequestHandler(SocketServer.BaseRequestHandler):
    '''
    Request handler for the event server.

    This class is agnostic to message format (except for logging).
    '''
    # pylint: disable=W0201,W0212
    def setup(self):
        SocketServer.BaseRequestHandler.setup(self)
        threading.current_thread().name = (
            'EventServerRequestHandler-%d' % get_unique_id())
        # A thread to be used to send messages that are posted to the queue.
        self.send_thread = None
        # A queue containing messages.
        self.queue = Queue()

    def handle(self):
        # The handle() methods is run in a separate thread per client
        # (since EventServer has ThreadingMixIn).
        logging.debug('Event server: handling new client')
        try:
            self.server._subscribe(self.queue)

            self.send_thread = threading.Thread(
                target=self._run_send_thread,
                name='EventServerSendThread-%d' % get_unique_id())
            self.send_thread.daemon = True
            self.send_thread.start()

            # Process events: continuously read message and broadcast to all
            # clients' queues.
            while True:
                msg = self.request.recv(_MAX_MESSAGE_SIZE + 1)
                if len(msg) > _MAX_MESSAGE_SIZE:
                    logging.error('Event server: message too large')
                if len(msg) == 0:
                    break  # EOF
                self.server._post_message(msg)
        except socket.error, e:
            if e.errno in [errno.ECONNRESET, errno.ESHUTDOWN]:
                pass  # Client just quit
            else:
                raise e
        finally:
            logging.debug('Event server: client disconnected')
            self.queue.put(None)  # End of stream; make writer quit
            self.server._unsubscribe(self.queue)

    def _run_send_thread(self):
        while True:
            message = self.queue.get()
            if message is None:
                return
            try:
                self.request.send(message)
            except:  # pylint: disable=W0702
                return


class EventServer(SocketServer.ThreadingUnixStreamServer):
    '''
    An event server that broadcasts messages to all clients.

    This class is agnostic to message format (except for logging).
    '''
    allow_reuse_address = True
    socket_type = socket.SOCK_SEQPACKET
    daemon_threads = True

    def __init__(self, path=None):
        '''
        Constructor.

        @param path: Path at which to create a UNIX stream socket.
            If None, uses a temporary path and sets the CROS_FACTORY_EVENT
            environment variable for future clients to use.
        '''
        # A set of queues listening to messages.
        self._queues = set()
        # A lock guarding the _queues variable.
        self._lock = threading.Lock()
        if not path:
            path = tempfile.mktemp(prefix='cros_factory_event.')
            os.environ[CROS_FACTORY_EVENT] = path
            logging.info('Setting %s=%s', CROS_FACTORY_EVENT, path)
        SocketServer.UnixStreamServer.__init__(  # pylint: disable=W0233
            self, path, EventServerRequestHandler)

    def _subscribe(self, queue):
        '''
        Subscribes a queue to receive events.

        Invoked only from the request handler.
        '''
        with self._lock:
            self._queues.add(queue)

    def _unsubscribe(self, queue):
        '''
        Unsubscribes a queue to receive events.

        Invoked only from the request handler.
        '''
        with self._lock:
            self._queues.discard(queue)

    def _post_message(self, message):
        '''
        Posts a message to all clients.

        Invoked only from the request handler.
        '''
        try:
            if logging.getLogger().isEnabledFor(logging.DEBUG):
                logging.debug('Event server: dispatching object %s',
                              pickle.loads(message))
        except:  # pylint: disable=W0702
            # Message isn't parseable as a pickled object; weird!
            logging.info(
                'Event server: dispatching message %r', message)

        with self._lock:
            for q in self._queues:
                # Note that this is nonblocking (even if one of the
                # clients is dead).
                q.put(message)


class EventClient(object):
    EVENT_LOOP_GOBJECT_IDLE = 'EVENT_LOOP_GOBJECT_IDLE'
    EVENT_LOOP_GOBJECT_IO = 'EVENT_LOOP_GOBJECT_IO'

    '''
    A client used to post and receive messages from an event server.

    All events sent through this class must be subclasses of Event.  It
    marshals Event classes through the server by pickling them.
    '''
    def __init__(self, path=None, callback=None, event_loop=None, name=None):
        '''
        Constructor.

        @param path: The UNIX seqpacket socket endpoint path.  If None, uses
            the CROS_FACTORY_EVENT environment variable.
        @param callback: A callback to call when events occur.  The callback
            takes one argument: the received event.
        @param event_loop: An event loop to use to post the events.  May be one
            of:

            - A Queue object, in which case a lambda invoking the callback is
              written to the queue.
            - EVENT_LOOP_GOBJECT_IDLE, in which case the callback will be
              invoked in the gobject event loop using idle_add.
            - EVENT_LOOP_GOBJECT_IO, in which case the callback will be
              invoked from an async IO handler.
        @param name: An optional name for the client
        '''
        self.socket = socket.socket(socket.AF_UNIX, socket.SOCK_SEQPACKET)
        self.callbacks = set()
        logging.debug('Initializing event client')

        should_start_thread = True

        path = path or os.environ[CROS_FACTORY_EVENT]
        self.socket.connect(path)
        self._lock = threading.Lock()

        if callback:
            if isinstance(event_loop, Queue):
                self.callbacks.add(
                    lambda event: event_loop.put(
                        lambda: callback(event)))
            elif event_loop == self.EVENT_LOOP_GOBJECT_IDLE:
                import gobject
                self.callbacks.add(
                    lambda event: gobject.idle_add(callback, event))
            elif event_loop == self.EVENT_LOOP_GOBJECT_IO:
                import gobject
                should_start_thread = False
                gobject.io_add_watch(
                    self.socket, gobject.IO_IN,
                    lambda source, condition: self._read_one_message())
                self.callbacks.add(callback)
            else:
                self.callbacks.add(callback)

        if should_start_thread:
            self.recv_thread = threading.Thread(
                target=self._run_recv_thread,
                name='EventServerRecvThread-%s' % (name or get_unique_id()))
            self.recv_thread.daemon = True
            self.recv_thread.start()
        else:
            self.recv_thread = None

    def close(self):
        '''Closes the client, waiting for any threads to terminate.'''
        if not self.socket:
            return

        # Shutdown the socket to cause recv_thread to terminate.
        self.socket.shutdown(socket.SHUT_RDWR)
        if self.recv_thread:
            self.recv_thread.join()
        self.socket.close()
        self.socket = None

    def __del__(self):
        self.close()

    def post_event(self, event):
        '''
        Posts an event to the server.
        '''
        logging.debug('Event client: sending event %s', event)
        message = pickle.dumps(event, protocol=2)
        if len(message) > _MAX_MESSAGE_SIZE:
            # Log it first so we know what event caused the problem.
            logging.error("Message too large (%d bytes): event is %s" %
                          (len(message), event))
            raise IOError("Message too large (%d bytes)" % len(message))
        self.socket.sendall(message)

    def wait(self, condition):
        '''
        Waits for an event matching a condition.

        @param condition: A function to evaluate.  The function takes one
            argument (an event to evaluate) and returns whether the condition
            applies.
        '''
        queue = Queue()

        def check_condition(event):
            if condition(event):
                queue.put(event)

        try:
            with self._lock:
                self.callbacks.add(check_condition)
            return queue.get()
        finally:
            with self._lock:
                self.callbacks.remove(check_condition)

    def _run_recv_thread(self):
        '''
        Thread to receive messages and broadcast them to callbacks.
        '''
        while self._read_one_message():
            pass

    def _read_one_message(self):
        '''
        Handles one incomming message from the socket.

        @return: True if event processing should continue (i.e., not EOF).
        '''
        bytes = self.socket.recv(_MAX_MESSAGE_SIZE + 1)
        if len(bytes) > _MAX_MESSAGE_SIZE:
            # The message may have been truncated - ignore it
            logging.error('Event client: message too large')
            return True

        if len(bytes) == 0:
            return False

        try:
            event = pickle.loads(bytes)
            logging.debug('Event client: dispatching event %s', event)
        except:
            logging.warn('Event client: bad message %r', bytes)
            traceback.print_exc(sys.stderr)
            return True

        with self._lock:
            callbacks = list(self.callbacks)
        for callback in callbacks:
            try:
                callback(event)
            except:
                logging.warn('Event client: error in callback')
                traceback.print_exc(sys.stderr)
                # Keep going

        return True