#!/usr/bin/env python
#
# gource - converts juju status yaml output into gource log
# Copyright (C) 2011 Canonical Ltd. All Rights Reserved.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

from hashlib import sha256
import subprocess
import yaml
import time
import argparse
import sys, os

"""
 status -> enact changes
 if cmd then
   save command
   execute command
 status -> enact changes
 loop
"""


class Thing(object):
    def __init__(self, thing_id, thing='', parent=None):
        self.thing_id = thing_id
        self.thing = thing
        self.parent = parent

    def __hash__(self):
        return sha256(self.thing_id)

    def __eq__(self, other):
        return self.thing_id == other.thing_id


class Bootstrap(Thing):
    def __init__(self, thing_id, thing):
        super(Bootstrap, self).__init__(thing_id, thing)

        self.color = '00FF00'

    @property
    def path(self):
        return ""

    @property
    def filename(self):
        return "bootstrap%s" % self.thing_id


class Service(Thing):
    def __init__(self, thing_id, thing=''):
        super(Service, self).__init__(thing_id, thing)

        self.color = 'FFFF00'

    @property
    def path(self):
        return "services/%s" % self.thing_id

    @property
    def filename(self):
        return "%s/service" % self.path


class Machine(Thing):
    def __init__(self, thing_id, thing=''):
        super(Machine, self).__init__(thing_id, thing)

        self.color = '0000FF'

    @property
    def path(self):
        return "machines/%s" % self.thing_id

    @property
    def filename(self):
        return "%s/instance" % self.path


class Relation(Thing):
    def __init__(self, thing_id, thing):
        super(Relation,self).__init__(thing_id, thing)
        (self.a, self.b) = self.thing_id.split('-',1)
        """ This makes it so we only get one relation per pair """
        if self.a > self.b:
            save = self.a
            self.a = self.b
            self.b = save
            self.thing_id = "%s-%s" % (self.a, self.b)
        self.service_a = ServiceRelation(self.a, self.b, Service(self.b))
        self.service_b = ServiceRelation(self.b, self.a, Service(self.a))

        self.color = '772953'

    @property
    def path(self):
        return "relations/%s" % self.thing_id

    @property
    def filename(self):
        return ["%s/%s/%s" % (self.path, self.a, self.a),
                "%s/%s/%s" % (self.path, self.b, self.b),
                "%s" % self.service_a.filename,
                "%s" % self.service_b.filename]


class ServiceRelation(Thing):
    def __init__(self, thing_id, thing, parent):
        super(ServiceRelation, self).__init__(thing_id, thing, parent)

        self.color = '772953'

    @property
    def path(self):
        return "%s/relations" % self.parent.path

    @property
    def filename(self):
        return "%s/%s" % (self.path, self.thing_id)


class ServiceMachine(Thing):

    def __init__(self, thing_id, thing, parent):
        super(ServiceMachine, self).__init__(thing_id, thing, parent)

        self.color = '0000FF'

    @property
    def path(self):
        return "%s/%s" % (self.parent.path, self.thing_id)

    @property
    def filename(self):
        return "%s/machine%s" % (self.path, self.thing)


class MachineService(Thing):

    def __init__(self, thing_id, thing, parent, state=None):
        super(MachineService, self).__init__(thing_id, thing, parent)

        if state == 'error':
            self.color = 'FF0000'
        elif state == 'started' or state == 'installed':
            self.color = '00FF00'
        else:
            self.color = 'FFFF00'

    @property
    def path(self):
        return "%s/%s" % (self.parent.path, self.thing_id)

    @property
    def filename(self):
        return "%s/%s-%s" % (self.path, self.thing_id, self.thing)


def calc_relations(services):
    relations = set()
    for srvname, service in services.iteritems():
        if 'relations' in service:
            for rel, relserv in service['relations'].iteritems():
               relations.add('%s-%s' % (srvname, relserv))
    return relations


def calc_units(services):
    new_units = {}
    for service_name, service in services.iteritems():
        if 'units' in service:
            for unit_id, unit in service['units'].iteritems():
                new_units[unit_id] = unit
    return new_units

class Status(object):

    def __init__(self, status={}):
        self._last_status = status
        self.last_command = 'juju'

    def bootstrap(self):
        save_last_command = self.last_command
        self.last_command = 'bootstrap'
        self.add(Bootstrap('0',None))
        self.last_command = save_last_command

    def unbootstrap(self):
        save_last_command = self.last_command
        self.last_command = 'destroy-environment'
        self.delete(Bootstrap('0',None))
        self.last_command = save_last_command

    def update_status(self, new_status):

        if type(new_status) is not dict:
            return

        if self._last_status == {} and new_status != {}:
            self.bootstrap()

        if 'machines' in new_status:
            self.update(new_status['machines'], Machine, 'machines')
        else:
            self.update({}, Machine, 'machines')

        if 'services' in new_status:
            self.update(new_status['services'], Service, 'services')
            new_relations = calc_relations(new_status['services'])
        else:
            self.update({}, Service, 'services')
            new_relations = set()

        if 'services' in self._last_status:
            old_relations = calc_relations(self._last_status['services'])
        else:
            old_relations = set()

        self.update_relations(old_relations, new_relations)

        self.update_machine_assignments(new_status)

        if new_status == {} and self._last_status != {}:
            self.unbootstrap()

        self._last_status = new_status

    def update_machine_assignments(self, new):
        if 'services' in new:
            new_units = calc_units(new['services'])
        else:
            new_units = {}

        if 'services' in self._last_status:
            old_units = calc_units(self._last_status['services'])
        else:
            old_units = {}

        for unit_id, unit in new_units.iteritems():
            (service,idonly) = unit_id.split('/')
            if unit_id not in old_units:
                self.add(ServiceMachine(unit_id,
                            unit['machine'],
                            Service(service)))
                self.add(MachineService(service,
                            idonly,
                            Machine(unit['machine']), state=unit['agent-state']))
            else:
                if unit['agent-state'] != old_units[unit_id]['agent-state']:
                    self.modify(MachineService(service,
                                idonly,
                                Machine(unit['machine']), state=unit['agent-state']))

        for unit_id, unit in old_units.iteritems():
            if unit_id not in new_units:
                (service,idonly) = unit_id.split('/')
                self.delete(ServiceMachine(unit_id,
                            unit['machine'],
                            Service(service)))
                self.delete(MachineService(service,
                            idonly,
                            Machine(unit['machine'])))
            else:
                pass

    def update_relations(self, old, new):
        for rel_id in new:
            if rel_id not in old:
                self.add(Relation(rel_id, rel_id))
            else:
                pass
        for rel_id in old:
            if rel_id not in new:
                self.delete(Relation(rel_id, rel_id))

    def update(self, new_things, thing_class, thing_key):
        for thing_id, thing in new_things.iteritems():
            if thing_key not in self._last_status or thing_id not in self._last_status[thing_key]:
                self.add(thing_class(thing_id, thing))
            else:
                pass

        if thing_key in self._last_status:
            for thing_id, thing in self._last_status[thing_key].iteritems():
                if thing_id not in new_things:
                    self.delete(thing_class(thing_id, thing))

    def statLog(self, user, filename, operation, color):
        logtime = int(time.time())
        if type(filename) is str:
            filename = [filename]
        for f in filename:
            print "%d|%s|%s|%s|%s" % (logtime, user, operation, f, color)

    def add(self, thing):
        self.statLog(self.last_command, thing.filename, 'A', thing.color)

    def delete(self, thing):
        self.statLog(self.last_command, thing.filename, 'D', thing.color)

    def modify(self, thing):
        self.statLog(self.last_command, thing.filename, 'M', thing.color)

GOURCE_ARGS=['--highlight-dirs',
             '--file-idle-time','1000000',
             '--log-format','custom',
             '--user-friction','0.5',
             '-']

def main():
    parser = argparse.ArgumentParser(description="""
            poll juju status and turn it into a stream of logs compatible with gource""")
    parser.add_argument('--run-gource', default=False, action='store_true',
            help='run gource and pipe stdout directly to it')
    parser.add_argument('--only-gource', default=False, action='store_true',
            help='just run gource the way --run-gource would')
    parser.add_argument('--print-gource-args', default=False, action='store_true',
            help='just prints gource arguments and exits')
    parser.add_argument('juju_args', nargs='*', default=[], 
            help='Arguments are passed along to juju status')
    args = parser.parse_args()

    if args.print_gource_args:
        print ' '.join(GOURCE_ARGS)
        return

    if args.only_gource:
        os.execvp('gource',GOURCE_ARGS)

    if args.run_gource:
        devnull = open('/dev/null', 'a')
        run_gource = GOURCE_ARGS
        run_gource.insert(0, 'gource')
        gource = subprocess.Popen(run_gource, stdin=subprocess.PIPE, stdout=devnull)
        sys.stdout = gource.stdin

    stat = Status()

    while(True):
        run_juju=['juju','status']
        run_juju.extend(args.juju_args)
        s = subprocess.Popen(run_juju, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        (output, errors) = s.communicate()
        if s.returncode == 0:
            newstat = yaml.load(output)
            stat.update_status(newstat)
        elif s.returncode == 1 and errors is not None and errors.find('environment not found') != -1:
            stat.update_status({})

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        pass
