#!/usr/bin/python
# snake-tree - program to manipulate the snake tree db

import os, sys, time
import optparse
import snake.uri, snake.client, snake.log
from snake.constants import *

try:
    import snake.treedb
    import snake.config
except ImportError, e:
    # will check sys.modules later
    pass

snake_commands = ["add","check","update","remove","list","info"]

# initialize data methods
snake_tree_list   = None
snake_uri_add     = None
snake_uri_remove  = None
snake_tree_remove = None

def setup_option_parser():
    parser=optparse.OptionParser(usage="snake-tree [options] < %s >" % (", ".join(snake_commands)))
    parser.add_option("-s", "--server",
                action="store", dest="server", default=os.environ.get("SNAKE_SERVER",""),
                help="Specify the SNAKE server")
    parser.add_option("-p", "--port",
                action="store", dest="port", default=os.environ.get("SNAKE_PORT","2903"),
                help="The SNAKE xmlrpc port")
    parser.add_option("-v", "--verbose",
                action="store_true", dest="verbose", default=False,
                help="Output extra information")

    # if we are accessing a local treedb, offer a conffile cmdline parameter
    if sys.modules.has_key("snake.treedb"):
        parser.add_option("-c", "--config",
            action="store", dest="conffile", default=snake.config.DEFAULT_CONFIGFILE,
            help="Override default configuration file (only for local system)")

    return parser

def parse_tree_args(args):
    '''examines the argument list and returns (urilist,idlist), where
    urilist is all the given URIs and idlist is all of the given arguments
    that are known tree IDs. As a special case, 'all' will expand to a list
    of all known IDs. Other arguments are ignored.'''
    urilist = list()
    idlist = list()
    known_ids = [hasattr(t,'id') and t.id or t.get('id') for t in snake_tree_list(dict())]

    # special case: for "all", return a list of all known tree ids.
    if 'all' in args:
        idlist = known_ids
        # remove it from the arg list
        while 'all' in args:
            args.remove('all')

    # Parse the rest of the args given.
    for a in args:
        if '/' in a:
            urilist.append(a)
        elif a in known_ids and a not in idlist:
            idlist.append(a)

    return (urilist,idlist)

def print_table(trees, action=None, sort=True, reverse=False):
    '''Display a formatted table of trees.
       Accepts arguments:
         - a list of tree objects or list of tree dictionaries
         - action: optional string representing an action
         - sort: optional boolean indicating whether sorting is needed
         - reverse: True|False boolean argument passed to .sort() method
    '''

    # sort the list for display
    # FIXME: should this be done by the server instead?
    if sort:
        undecorated = [isinstance(t,dict) and t or t.__dict__ for t in trees]
        decorated = [(dict_['id'], dict_) for dict_ in undecorated]
        decorated.sort(reverse=reverse)
        trees = [dict_ for (key, dict_) in decorated]

    print " %-15s %-8s %-40s" % ("ID", "Arch", "Name")
    print "=%-15s=%-8s=%-40s" % ("="*15, "="*10, "="*40)

    if action: print action

    for d in trees:
        if isinstance(d, dict):
            t = snake.tree.Tree(dict=d)
        else:
            t = d
        print " %-15s %-8s %-40s" % (t.id, t.arch, str(t))

def add(uri):
    r = True
    try:
        r = snake_uri_add(uri)
    except Exception, e:
        print "Could not add %s: %s" % (uri,e)
        r = False
    return r

# FIXME - perhaps move into snake/tree.py as `print t.info()` ?
# FIXME - yes, some of this is ugly ... but it works
def print_tree_info(t):
    '''Utility function to display formatted information about a tree'''
    for (k,v) in [["Name",      t.name],
                  ["Arch",      t.arch],
                  ["Id",        t.id],
                  ["Version",   t.version],
                  ["Family",    t.family],
                  ["Variant",   t.variant],
                  ["Time",      time.strftime("%Y-%m-%d %H:%M %Z",time.localtime(float(t.time)))],
                  ["URI's",     ", ".join(t.uris)],
                  ["Images",    ", ".join(map(lambda (k,v): "%s (%s)" % (k,", ".join(v.keys())), t.images.items()))],
                 ]:
        print "%-12s : %s" % (k,v)
    if hasattr(t,"kernel_args") and t.kernel_args:
        print "%-12s : %s" % ("KernelArgs", t.kernel_args)
    print ""

def main():
    parser = setup_option_parser()
    (opt,args) = parser.parse_args()

    global log, snake_tree_list, snake_uri_add, snake_uri_remove, snake_tree_remove

    # Set up logging
    log = snake.log.setup_logger(opt.verbose and snake.log.DEBUG or snake.log.INFO)

    # Use XMLRPC?
    if opt.server:
        if not snake.client.check_server(opt.server, opt.port):
            '''the server is not responding'''
            log.error("The snake server '%s:%s' is not responding. \
 Please check your server URL and try again." % (opt.server,opt.port))
            sys.exit(1)
        server = snake.client.connect(opt.server, opt.port)
        snake_tree_list   = server.tree.list
        snake_uri_add     = server.uri.add
        snake_uri_remove  = server.uri.remove
        snake_tree_remove = server.tree.remove
        log.debug("Using treedb on %s:%s as data source" % (opt.server, opt.port))

    # Use local db?
    elif sys.modules.has_key("snake.treedb"):
        if not sys.modules.has_key("snake.treedb"):
            log.error("Unable to use local treedb, is snake-server installed?")
            sys.exit(1)

        # load server configuration file
        try:
            snake.config.getconfig(filename=opt.conffile)
        except Exception, e:
            log.error("Error reading configuration from '%s'" % opt.conffile)
            log.debug(e)
            sys.exit(1)

        snake_tree_list   = snake.treedb.dicttrees
        snake_uri_add     = snake.treedb.add
        snake_uri_remove  = snake.treedb.removeuri
        snake_tree_remove = snake.treedb.removetree
        log.debug("Using local treedb as data source")

    # no db used ... only useful commands will be 'info' and 'check' using URI's
    else:
        def raise_exception(x):
            raise Exception, "No data source defined"
        snake_tree_list   = lambda x: x
        snake_uri_add     = lambda x: raise_exception(x)
        snake_uri_remove  = lambda x: False
        snake_tree_remove = lambda x: False

        log.info("No data source used")

    # no command given
    if len(args) == 0:
        parser.error("You need to give some command")
        sys.exit(1)

    cmd = args[0]
    cmdopts = args[1:]

    # Handle commands
    if cmd == "add":
        for o in cmdopts:
            if add(o):
                t = snake_tree_list({'uri':o}).pop()
                if isinstance(t,dict): t = snake.tree.Tree(dict=t)
                print "%s: %s" % (t.id,t)
                print "  [  OK  ] " + o 
            else:
                print "  [FAILED] " + o 

    elif cmd == "info":
        (urilist,idlist) = parse_tree_args(cmdopts)

        if len(idlist) + len(urilist) == 0:
            log.error("No matching Trees to list")
            sys.exit(1)

        for id in idlist:
            local_t = snake_tree_list({'id':id}).pop()
            t = snake.tree.Tree(dict=local_t)
            print_tree_info(t)

        for u in urilist:
            status = snake.client.check_tree(u)
            if status == CHECK_TREE_SUCCESS:
                t = snake.uri.uri_to_tree(u)
                print_tree_info(t)
            else:
                print "Not recognized as a valid tree: %s" % u

    elif cmd == "list":
        arg = dict()
        for o in cmdopts:
            if '=' in o:
                k,v = o.split('=')
                arg[k]=v
        trees = snake_tree_list(arg)

        if len(trees) == 0:
            log.error("No matching Trees to list")
            sys.exit(1)

        print_table(trees)

    elif cmd in ("check","update"):
        (urilist,idlist) = parse_tree_args(cmdopts)
        arg={}
        for o in cmdopts:
            if '=' in o:
                k,v = o.split('=')
                arg[k]=v

        # build list of trees
        if arg:
            trees = snake_tree_list(arg)
        else:
            trees = list()

        for t in trees:
            idlist.append(t['id'])

        if len(urilist) + len(idlist) == 0:
            log.error("No trees found matching supplied arguments '%s'" % " ".join(cmdopts))
            sys.exit(1)

        if cmd == 'update':
            update = True
        else:
            update = False

        # Check locally-cached trees
        for id in idlist:
            local_t = snake_tree_list({'id':id}).pop()
            if isinstance(local_t,dict): local_t = snake.tree.Tree(dict=local_t)
            print "%s: %s" % (local_t.id, local_t)
            delete_tree = False
            for u in local_t.uris:
                # Check to see if there's still a tree there
                status = snake.client.check_tree(u)
                if status == CHECK_TREE_GONE: # tree is permanently gone
                    print "  [ GONE ] " + u
                    if update: snake_uri_remove(u)
                elif status == CHECK_TREE_SUCCESS:
                    # OK, there's a tree. Is it the same tree?
                    remote_t = snake.uri.uri_to_tree(u)
                    if remote_t.id != id:
                        # That tree's been updated or changed.
                        print "  [UPDATE] " + u
                        if update:
                            # Nuke the out-of-date info
                            snake_uri_remove(u)
                            # Add the new info
                            add(u)
                            # And add the new info to the list of ids
                            if remote_t.id not in idlist:
                                idlist.append(remote_t.id)
                            print "    %s->%s" % (local_t,remote_t)
                    else:
                        # The URL is reachable, it contains a tree, and the
                        # tree has the same ID. Sounds good!
                        print "  [  OK  ] " + u
                else: # missing metadata, etc.
                    print "  [BROKEN] " + u
                    # leave it alone - it might be fixed later
            print

        # Check remote trees listed on the commandline
        if not update:
            for u in urilist:
                status = snake.client.check_tree(u)
                if status == CHECK_TREE_SUCCESS:
                    t = snake.uri.uri_to_tree(u)
                    print "[  OK  ] " + u
                    print "         " + str(t)
                elif status == CHECK_TREE_GONE:
                    print "[ GONE ] " + u
                else:
                    print "[BROKEN] " + u

    elif cmd == "remove":
        (urilist,idlist) = parse_tree_args(cmdopts)

        for u in urilist:
            snake_uri_remove(u)

        for i in idlist:
            snake_tree_remove(i)
    else:
        parser.error("Unrecognized command: %s" % cmd)
        sys.exit(1)

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