#!/usr/bin/python
# vi: ts=4 noexpandtab
#
#    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", "virtualization_type",
)

# in query data, if a field is not set, use this value
#   (when virtualization-type was added to data, the default for
#    an empty value in that field is 'paravirtual')
defaults = { }
defaults["virtualization_type"]="paravirtual"

fields = { }
for i in range(len(result_fields)):
	fields[result_fields[i]]=i

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

vocabularies = [
	("suite", {
		"hardy": "hardy",
		"karmic": "karmic",
		"latest": "maverick",
		"lucid": "lucid",
		"maverick": "maverick",
		"natty": "natty",
	}),
	("arch", {
		"i386": "i386",
		"amd64": "amd64",
		"x86_64": "x86_64",
		"64": "amd64",
		"32": "i386"}),
	("img_type", {
		"instance": "instance-store",
		"instance-store": "instance-store",
		"ebs": "ebs"}
	),
	("build_name", {"server": "server", "desktop": "desktop"}),
	("stream", {"daily": "daily", "released": "released"}),
	("ec2_region", {
		"us-east-1": "us-east-1",
		"east": "us-east-1",
		"us": "us-east-1",
		"us-west-1": "us-west-1",
		"west": "us-west-1",
		"eu-west-1": "eu-west-1",
		"eu": "eu-west-1",
		"ap-southeast-1": "ap-southeast-1",
		"southeast": "ap-southeast-1",
		"ap": "ap-southeast-1"}),
	("virtualization_type", {
		"pvm": "paravirtual",
		"hvm": "hvm",
		"paravirtual" : "paravirtual" }),
	]


class MissingArgumentException(Exception):
	pass

class DuplicateArgumentException(Exception):
	pass

class UnknownArgumentException(Exception):
	pass


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

def parse_args(args, options):
	d = {}
	unknown = [ ]
	for arg in args:
		found = False
		for name, values in vocabularies:
			if arg in values:
				if name in d:
					raise DuplicateArgumentException(
						"%s specified multiple times (%s, %s)"%(
							name, d[name], arg))
				d[name] = values[arg]
				found = True
				break
		if not found:
			print "adding %s" % arg
			unknown.append(arg)

	if len(unknown) > 1:
		raise UnknownArgumentException("unrecognized strings: %s" % unknown)
 	elif len(unknown) == 1:
		if "suite" in d or options.suite:
			raise UnknownArgumentException("unrecognized string: %s" % unknown[0])
		d["suite"]=unknown[0]
		

	for k, v in d.items():
		setattr(options, k, v)

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
			ind=result_fields.index(name)
			try:
				rowval=row[ind]
			except:
				try:
					rowval=defaults[name]
				except:
					rowval=None
			#print "val=%s name=%s rowval=%s" % (val,name,rowval)
			if val is not None and val != rowval:
				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", default="server",
		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("--virtualization-type", dest="virtualization_type",
		metavar="VTYPE", default="pvm",
		help="the virtualization type (one of 'pvm' or 'hvm')")

	#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 opts.virtualization_type == "pvm":
		opts.virtualization_type = "paravirtual"

	parse_args(args[1:], opts)

	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, 
			"virtualization_type" : opts.virtualization_type }
		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()
