#!/usr/bin/python
# vi: ts=4 expandtab
#
#    Copyright (C) 2009-2010 Canonical Ltd.
#
#    Author: Scott Moser <scott.moser@canonical.com>
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License version 3, as
#    published by the Free Software Foundation.
#
#    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 General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.

import optparse
import urllib2
import yaml
import sys
import subprocess

# Usage: uec-query-builds [ options ] mode [ arguments ]
#   --suite
#   --build
#   --stream
#
#   --config
#   --base-url
#
#   --serial
#   --cloud
#   --arch
#   --type (instance-store, ebs)
#
# mode
#   latest
#     suite, build, stream=default, serial
#   is-update-available

BASE_URL="http://uec-images.ubuntu.com/query/"
result_fields = (
    "suite", "build_name", "name", "serial", "img_type",
    "arch", "region", "ami", "aki", "ari"
)
fields = { }
for i in range(len(result_fields)):
    fields[result_fields[i]]=i

qtype_summary="summary"
qtype_all="all"
qtype_ec2_current="ec2_current"

class MissingArgumentException(Exception):
    pass

def exitMissingArgument(parser,missingArgumentException):
    parser.print_help()
    sys.stderr.write("%s\n" % str(missingArgumentException))
    sys.exit(1)

def limiter(data,limits):
    ret=[ ]
    for row in data:
        match = True
        for name, val in limits.iteritems():
            #print "name = %s, val = %s" % (name,val)
            if name not in fields: continue
            if val is not None and val != row[fields[name]]:
                match = False
                break
        if match is True: ret.append(row)
    return(ret)

def checkopts(options, req):
    for f in req:
        if getattr(options, f, None) is None:
            raise MissingArgumentException("must provide argument for %s" % f)

def get_data(opts, type):
    if type == qtype_summary:
        checkopts(opts, ( "stream", "base_url" ))
        url = "%s/%s.latest.txt" % (opts.base_url, opts.stream )
    elif type == qtype_ec2_current:
        checkopts(opts, ( "stream", "suite", "build_name", "base_url" ))
        url = "%s/%s/%s/%s.current.txt" % \
            (opts.base_url, opts.suite, opts.build_name, opts.stream )
    else:
        checkopts(opts, ( "stream", "suite", "build_name", "base_url" ))
        url = "%s/%s/%s/%s.txt" % \
            (opts.base_url, opts.suite, opts.build_name, opts.stream )
    try:
        request = urllib2.urlopen(url)
    except Exception as e:
        raise Exception("Unable to load %s\n\t%s\n" % (url,e))

    lines=request.read().split('\n')
    data = [ ]
    for l in lines:
        data.append(l.split('\t'))
        if(len(data[len(data)-1]) < 2): data.pop()
    return(data)

def serial_gt(ser1, ser2):
    if fix_serial(str(ser1)) > fix_serial(str(ser2)): return True
    return False

def fix_serial(ser):
    if len(ser.split('.')) > 1: return("%s.0" % ser)
    return(ser)

def handle_config(option, opt_str, value, parser):
    try:
        f=open(value)
        cfg = yaml.load(f)
        f.close()
    except:
        raise optparse.OptionValueError("unable to open %s" % value)
    for n in result_fields + ( "stream", ):
        if n in cfg:
            setattr(parser.values,n,cfg[n])
    return

def main():
    parser = optparse.OptionParser()
    modes = ( "is-update-available", "latest", "latest-ec2" )

    parser.add_option("--suite", dest="suite", metavar="SUITE",
        help="suite to query ('hardy', 'karmic', 'lucid')")
    parser.add_option("--build-name", dest="build_name",
        metavar="BUILD_NAME",
        help="build name ('server', 'desktop' ..)")
    parser.add_option("--stream", dest="stream", metavar="STREAM",
        default="released",
        help="stream query ('released', 'daily')")
    parser.add_option("--base-url", dest="base_url", metavar="BASE_URL",
        default=BASE_URL,
        help="the base url to query")
    parser.add_option("--output", dest="output_fname", metavar="FILE",
        default="-", help="write output to file, default is stdout")
    parser.add_option("--serial", dest="serial", metavar="SERIAL",
        help="build serial serial to use (YYYYMMDD)")
    parser.add_option("--system-suite", dest="system_suite",
        action="store_true", default=False,
        help="use output of 'lsb_release --codename --short' for suite")

    parser.add_option("--config", metavar="CONFIG", dest="config",
        action="callback", callback=handle_config, type="string",
        help="yaml config file to read")

    parser.add_option("--region", dest="ec2_region", metavar="REGION",
        help="the ec2 region to query")

    parser.add_option("--img-type", dest="img_type", metavar="TYPE",
        help="the ec2 image type (one of: ebs, instance)")

    parser.add_option("--arch", dest="arch", metavar="ARCH",
        help="the architecture. (one of: i386, amd64)")

    #parser.add_option("-v","--verbose", dest="verbose",
    #    action="store_true", default=False,
    #    help="increase verbosity")

    (opts, args) = parser.parse_args()

    if (len(args)) < 1 or args[0] not in modes:
        parser.error("Must give a mode (%s)" % ','.join(modes))
        sys.exit(1)

    if opts.output_fname == "-":
        output = sys.stdout
    else:
        output=open(opts.output_fname,"w")

    if opts.arch == "x86_64": opts.arch="amd64"

    if opts.system_suite:
        cmd=['lsb_release', '--codename', '--short' ]
        sp = subprocess.Popen(cmd, stdout=subprocess.PIPE)
        stdout, stderr = sp.communicate()
        if sp.returncode != 0:
            sys.error("Failed to get suite from system");
            sys.exit(1)
        opts.suite = stdout.strip()

    if args[0] == "latest":
        try:
            data = get_data(opts,qtype_summary)
            limits = { "build_name": opts.build_name, "suite" : opts.suite }
        except MissingArgumentException as e:
            exitMissingArgument(parser,e)

        for row in limiter(data,limits):
            output.write("%s\n" % '\t'.join(row))

    elif args[0] == "is-update-available":
        try:
            checkopts(opts,
                ( "stream", "base_url", "build_name", "suite", "serial" ))
            data = get_data(opts,qtype_summary)
        except MissingArgumentException as e:
            exitMissingArgument(parser,e)
        limits = { "build_name": opts.build_name, "suite" : opts.suite }
        result = limiter(data,limits)
        if len(result) > 1:
            sys.stderr.write("Received multiple matching results for %s:%s\n" \
                % ( opts.build_name, opts.suite ))
            sys.exit(1)
        elif len(result) == 1:
            result=result[0]
            if serial_gt(result[fields["serial"]],opts.serial):
                output.write("%s\n" % '\t'.join(result))
            else:
                output.write("")
        else:
            sys.stderr.write("Received no matching results for %s:%s\n" \
                % ( opts.build_name, opts.suite ))
            sys.exit(1)
    elif args[0] == "latest-ec2":
        try:
            checkopts(opts, ("stream", "base_url", "build_name", "suite" ))
            data = get_data(opts,qtype_ec2_current)
        except MissingArgumentException as e:
            exitMissingArgument(parser,e)
        limits = { "region": opts.ec2_region,
            "img_type" : opts.img_type , "arch": opts.arch }
        result = limiter(data,limits)
        for row in limiter(data,limits):
            output.write("%s\n" % '\t'.join(row))
    else:
        parser.error("Unknown mode %s" % args[0])

if __name__ == '__main__':
    main()
