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