###############################################################################
##
##  Copyright 2011,2012 Tavendo GmbH
##
##  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.
##
##
##  Modified by Mikko Ahlroth, 2013
##  Original is part of AutobahnPython and can be found at
##  https://github.com/tavendo/AutobahnPython/blob/master/examples/websocket/broadcast/server.py
##  Modifications are licensed with the MIT License, see LICENSE for details.
##
###############################################################################

import sys
import json
import psutil

from twisted.internet import reactor
from twisted.python import log

from autobahn.websocket import WebSocketServerFactory, \
                                WebSocketServerProtocol, \
                                listenWS

TICK_INTERVAL = 2
CPU_INTERVAL = 1
WEBSOCKET = 'ws://localhost:8081/'
NET_INTERFACE = 'eth0'

INITIAL_INFO = {
    'cpus': psutil.NUM_CPUS,
    'boot': psutil.get_boot_time(),
    'version': '0.1',
}


def get_sys_status():
    cpu = psutil.cpu_percent(interval=CPU_INTERVAL, percpu=True)
    mem = psutil.virtual_memory().percent
    swap = psutil.swap_memory().percent
    procs = len(psutil.get_pid_list())
    ct = psutil.cpu_times_percent()
    cpu_times = {
        'user': ct[0],
        'system': ct[1],
        'idle': ct[2],
        'nice': ct[3],
        'iowait': ct[4],
        'irq': ct[5],
        'softirq': ct[6],
        'steal': ct[7],
        'guest': ct[8],
        'guest_nice': ct[9],
    }
    net = psutil.network_io_counters(pernic=True)[NET_INTERFACE]
    net_out = {
        'bytes_out': net.bytes_sent,
        'bytes_in': net.bytes_recv,
    }
    disk = psutil.disk_io_counters()
    disk_out = {
        'read_count': disk.read_count,
        'write_count': disk.write_count,
    }

    users = len(psutil.get_users())

    return {
        'cpu': cpu,
        'mem': mem,
        'swap': swap,
        'procs': procs,
        'cpu_times': cpu_times,
        'net': net_out,
        'disk': disk_out,
        'users': users,
    }



class BroadcastServerProtocol(WebSocketServerProtocol):

    def onOpen(self):
        self.factory.register(self)

    def onMessage(self, msg, binary):
        pass

    def connectionLost(self, reason):
        WebSocketServerProtocol.connectionLost(self, reason)
        self.factory.unregister(self)


class BroadcastServerFactory(WebSocketServerFactory):

    def __init__(self, url, debug=False, debugCodePaths=False):
        WebSocketServerFactory.__init__(self, url, debug=debug, debugCodePaths=debugCodePaths)
        self.clients = []
        self.tick()

    def tick(self):
        self.broadcast(json.dumps(get_sys_status()))
        reactor.callLater(TICK_INTERVAL - CPU_INTERVAL, self.tick)

    def register(self, client):
        if not client in self.clients:
            print "registered client " + client.peerstr
            self.clients.append(client)
            client.sendMessage(json.dumps(INITIAL_INFO))

    def unregister(self, client):
        if client in self.clients:
            print "unregistered client " + client.peerstr
            self.clients.remove(client)

    def broadcast(self, msg):
        print "broadcasting message '%s' .." % msg
        for c in self.clients:
            c.sendMessage(msg)
            print "message sent to " + c.peerstr


class BroadcastPreparedServerFactory(BroadcastServerFactory):
    """
    Functionally same as above, but optimized broadcast using
    prepareMessage and sendPreparedMessage.
    """

    def broadcast(self, msg):
        print "broadcasting prepared message '%s' .." % msg
        preparedMsg = self.prepareMessage(msg)
        for c in self.clients:
            c.sendPreparedMessage(preparedMsg)
            print "prepared message sent to " + c.peerstr


if __name__ == '__main__':

    if len(sys.argv) > 1 and sys.argv[1] == 'debug':
        log.startLogging(sys.stdout)
        debug = True
    else:
        debug = False

    #ServerFactory = BroadcastServerFactory
    ServerFactory = BroadcastPreparedServerFactory

    factory = ServerFactory(WEBSOCKET,
                            debug=debug,
                            debugCodePaths=debug)

    factory.protocol = BroadcastServerProtocol
    factory.setProtocolOptions(allowHixie76=True)
    listenWS(factory)

    reactor.run()