#!/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

try:
    import snake.treedb
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="URL of SNAKE server")
    parser.add_option("-v", "--verbose",
                action="store_true", dest="verbose", default=False,
                help="Output extra information")
    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 add(uri):
    r = True
    try:
        r = snake_uri_add(uri)
    except snake.uri.URIError, e:
        print "Could not add %s: %s" % (uri,e.args[0])
        r = False
    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],
                  ["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)
    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):
            '''the server is not responding'''
            log.error("The snake server '%s' is not responding. \
 Please check your server URL and try again." % opt.server)
            sys.exit(1)
        server = snake.client.connect(opt.server)
        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.info("Using treedb on %s as data source" % opt.server)
    else:
        if not sys.modules.has_key("snake.treedb"):
            log.error("Unable to use local treedb, is snake-server installed?")
            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.info("Using local treedb as data source")

    # 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 == 0:
                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 "%-15s %-8s %-40s" % ("ID", "Arch", "Name")
        print "%-15s=%-8s=%-40s" % ("="*15, "="*10, "="*40)

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

    elif cmd in ('check','update'):
        (urilist,idlist) = parse_tree_args(cmdopts)

        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 == 1: # tree is permanently gone
                    print "  [ GONE ] " + u
                    if update: snake_uri_remove(u)
                elif status > 1: # missing metadata, etc.
                    print "  [BROKEN] " + u
                    # leave it alone - it might be fixed later
                elif status == 0:
                    # 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
            print

        # Check remote trees listed on the commandline
        if not update:
            for u in urilist:
                status = snake.client.check_tree(u)
                if status == 0:
                    t = snake.uri.uri_to_tree(u)
                    print "[  OK  ] " + u
                    print "         " + str(t)
                elif status == 1:
                    print "[ GONE ] " + u
                elif status > 1:
                    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
