#!/usr/bin/env python
# Copyright 2003 Gentoo Technologies, Inc.
# Distributed under the terms of the GNU General Public License v2
# $Header: $
#
# Author: Thomas de Grenier de Latour <degrenier@easyconnect.fr>
# The starting point ideas were found here:
# http://forums.gentoo.org/viewtopic.php?t=3011
# Thanks to eswanson and far for there contributions, and wolf31o2 for his
# support. Thanks also to karltk, some of this code is from his "equery" tool.

###############################################################################
# Meta:
__author__ = "Thomas de Grenier de Latour"
__email__ = "degrenier@easyconnect.fr"
__version__ = "0.1"
__productname__ = "eclean"
__description__ = "A cleaning tool for Gentoo distfiles and binaries."

###############################################################################
# Python imports:
# XXX: is portage needed? (for older python packages maybe?) 
import sys
import os
import string
import re
import getopt
import fpformat
import signal
import portage
from output import *

###############################################################################
# Misc. shortcuts to some portage stuff:
port_settings = portage.settings
distdir = port_settings["DISTDIR"]
pkgdir = port_settings["PKGDIR"]
cachedir = port_settings["PORTAGE_CACHEDIR"]
dbapi = portage.db["/"]["porttree"].dbapi

###############################################################################
# printVersion:
def printVersion():
	print __productname__ + "(" + __version__ + ") - " + \
		  __description__
	print "Author(s): " + __author__

###############################################################################
# printUsage: print help message. May also print partial help to stderr if an
# error from {'options','actions'} is specified.
def printUsage(error=None):
	out = sys.stdout
	if error != None: out = sys.stderr
	print >>out, white("Usage:"), turquoise(__productname__), \
		yellow("[option] ..."), green("<action> ...")
	print >>out
	if error in (None, 'options'):
		print >>out, "Available", yellow("options")+":"
		print >>out, yellow(" -C, --nocolor")+ \
			"      - turn off colors on output"
		print >>out, yellow(" -i, --interactive")+ \
			"  - ask confirmation before deleting"
		print >>out, yellow(" -p, --pretend")+ \
			"      - only display what would be cleaned"
		print >>out, yellow(" -q, --quiet")+ \
			"        - be as quiet as possible"
		print >>out, yellow(" -h, --help")+ \
			"         - display the help screen"
		print >>out, yellow(" -V, --version")+ \
			"      - display version info"
		print >>out
	if error in (None, 'actions'):
		print >>out, "Available", green("actions")+":"
		print >>out, green(" packages")+ \
			"      - clean outdated binary packages from:"
		print >>out, "                  ",teal(pkgdir)
		print >>out, green(" distfiles")+ \
			"     - clean outdated packages sources files from:"
		print >>out, "                  ",teal(distdir)
		print >>out

###############################################################################
# einfo: display an info message depending on a color mode
def einfo(message="", nocolor=False):
	if not nocolor: prefix = " "+green('*')
	else: prefix = ">>>"
	print prefix,message

###############################################################################
# eerror: display an error depending on a color mode
def eerror(message="", nocolor=False):
	if not nocolor: prefix = " "+red('*')
	else: prefix = "!!!"
	print >>sys.stderr,prefix,message

###############################################################################
# eprompt: display a user question depending on a color mode.
def eprompt(message, nocolor=False):
	if not nocolor: prefix = " "+red('>')+" "
	else: prefix = "??? "
	sys.stdout.write(prefix+message)
	sys.stdout.flush()

###############################################################################
# prettySize: integer -> byte/kilo/mega/giga converter. Optionnally justify the
# result. Output is a string.
def prettySize(size,justify=False):
	units = [" G"," M"," K"]
	size = size / 1024.
	while len(units) and size > 1024:
		size = size / 1024.
		units.pop()
	sizestr = fpformat.fix(size,1)+units[-1]
	if justify: 
		sizestr = " " + blue("[ ") + " "*(7-len(sizestr)) \
		          + green(sizestr) + blue(" ]")
	return sizestr

###############################################################################
# yesNoAllPrompt: print a prompt until user answer in yes/no/all. Return a 
# boolean for answer, and also may affect the 'accept_all' option.
# Note: i gave up with getch-like functions, to much bugs in case of escape
# sequences. Back to raw_input.
def yesNoAllPrompt(myoptions,message="Do you want to proceed?"):
	user_string="xxx"
	while not user_string.lower() in ["","y","n","a","yes","no","all"]:
		eprompt(message+" [Y/n/a]: ", myoptions['nocolor'])
		user_string = raw_input()
	if user_string.lower() in ["a","all"]:
		myoptions['accept_all'] = True
	myanswer = user_string.lower() in ["","y","a","yes","all"]
	return myanswer

###############################################################################
# ParseArgsException: for parseArgs() -> main() communication
class ParseArgsException(Exception):
	def __init__(self, value):
		self.value = value
	def __str__(self):
		return repr(self.value)

###############################################################################
# parseArgs: parse options and then actions. Raise exceptions on errors or
# non-action modes (help/version). Returns an actions list and an options dict.
def parseArgs(myoptions={}):
	# apply getopts to command line, show partial help on failure
	try: opts, args = getopt.getopt(sys.argv[1:], "CipqhV", \
		["nocolor", "interactive", "pretend", "quiet", "help", "version"])
	except: raise ParseArgsException('options')
	# set default options. 'nocolor' is set in main()
	myoptions['interactive'] = False
	myoptions['pretend'] = False
	myoptions['quiet'] = False
	myoptions['accept_all'] = False
	# modify this options depending on command line args
	for o, a in opts:
		if o in ("-h", "--help"):
			raise ParseArgsException('help')
		elif o in ("-V", "--version"):
			raise ParseArgsException('version')
		elif o in ("-C", "--nocolor"):
			myoptions['nocolor'] = True
			nocolor()
		elif o in ("-i", "--interactive") and not myoptions['pretend']:
			myoptions['interactive'] = True
		elif o in ("-p", "--pretend"):
			myoptions['pretend'] = True
			myoptions['interactive'] = False
		elif o in ("-q", "--quiet"):
			myoptions['quiet'] = True
	# only two actions are allowed: 'packages' and 'distfiles'.
	myactions = []
	if not len(args):
		raise ParseArgsException('actions')
	for a in args:
		if a in ['packages','distfiles'] and not a in myactions:
			myactions.append(a)
		else: raise ParseArgsException('actions')
	# return both actions and new option dictionnary
	return (myactions, myoptions)

###############################################################################
# findDistfiles: find all obsolete distfiles.
# XXX: what about cvs ebuilds? i should install some to see where it goes...
def findDistfiles():
	# this regexp extracts files names from SRC_URI. It is not very precise,
	# but we don't care (may return empty strings, etc.), since it is fast.
	file_regexp = re.compile('([a-zA-Z0-9_,\.\-\+]*)[\s\)]')
	clean_dict = {}
	keep = []
	# create a big list of files that appear in some SRC_URI.
	# Work on dep cache since it is much faster.
	for root, dirs, files in os.walk(cachedir):
		for file in files:
			if not file[-8:] == '.cpickle':
				# if not a .cpickle, then it is metadata cache
				path = os.path.join(root, file)
				try: 
					f = open(path)
					# SRC_URI is 4th line of metadata cache
					src_uri = f.readlines()[3]
				except: pass
				else:
					f.close()
					# split SRC_URI keeping at least all 
					# files names.
					keep += file_regexp.findall(src_uri)
	# create a dictionary of files from distdir that are no more in
	# any SRC_URI (ie not in 'keep').
	for file in os.listdir(distdir):
		if file not in keep:
			filepath = os.path.join(distdir, file)
			if os.path.isfile(filepath): 
				clean_dict[file]=[filepath]
	return clean_dict

###############################################################################
# findPackages: find all binary packages, and return a dictionnary of those
# that have no more corresponding ebuild in portage.
# XXX: packages are found only by symlinks. Maybe i should also return .tbz2
# files from All/ that have no corresponding symlinks.
def findPackages():
	clean_dict = {}
	# create a full package dictionnary
	for root, dirs, files in os.walk(pkgdir):
		if root[-3:] == 'All': continue
		for file in files:
			if not file[-5:] == ".tbz2":
				continue
			path = os.path.join(root, file)
			category = os.path.split(root)[-1]
			cpv = category+"/"+file[:-5]
			clean_dict[cpv] = [path]
			if os.path.islink(path):
				clean_dict[cpv].append(os.path.realpath(path))
	# keep only obsolete ones
	for mycpv in clean_dict.keys():
		if dbapi.cpv_exists(mycpv): 
			del clean_dict[mycpv]
	return clean_dict

###############################################################################
# doCleanup: takes a dictionnary {'display name':[list of files]}. Calculate
# size of each entry for display, prompt user if needed, delete files if needed
# and return the total size of files that [have been / would be] deleted.
def doCleanup(clean_dict,action,myoptions):
	# define vocabulary of this action
	if action == 'distfiles': file_type = 'file'
	else: file_type = 'binary package'
	# sorting helps reading
	clean_keys = clean_dict.keys()
	clean_keys.sort()
	clean_size = 0
	# clean all entries one by one
	for mykey in clean_keys:
		key_size = 0
		for file in clean_dict[mykey]:
			# get total size for an entry (may be several files, and
			# symlinks count zero)
			if os.path.islink(file): continue
			try: key_size += os.path.getsize(file)
			except: eerror("Could not read size of "+file, \
			               myoptions['nocolor'])
		if not myoptions['quiet']:
			# pretty print mode
			print prettySize(key_size,True),teal(mykey)
		elif myoptions['pretend'] or myoptions['interactive']:
			# file list mode
			for file in clean_dict[mykey]: print file
		#else: actually delete stuff, but don't print anything
		if myoptions['pretend']: clean_size += key_size
		elif not myoptions['interactive'] \
		     or myoptions['accept_all'] \
		     or yesNoAllPrompt(myoptions, \
		                       "Do you want to delete this " \
				       + file_type+"?"):
			# non-interactive mode or positive answer. 
			# For each file,...
			for file in clean_dict[mykey]:
				# ...get its size...
				filesize = 0
				if not os.path.exists(file): continue
				if not os.path.islink(file):
					try: filesize = os.path.getsize(file)
					except: eerror("Could not read size of "\
					               +file, myoptions['nocolor'])
				# ...and try to delete it.
				try: os.unlink(file)
				except: eerror("Could not delete "+file, \
				               myoptions['nocolor'])
				# only count size if successfully deleted
				else: clean_size += filesize
	# return total size of deleted or to delete files
	return clean_size

###############################################################################
# doAction: execute one action, ie display a few message, call the right find*
# function, and then call doCleanup with its result.
def doAction(action,myoptions):
	# define vocabulary for the output
	if action == 'packages': files_type = "binary packages"
	else: files_type = "distfiles"
	# find files to delete, depending on the action
	if not myoptions['quiet']:
		einfo("Building file list for "+action+" cleaning...", \
		      myoptions['nocolor'])
	if action == 'packages':
		clean_dict = findPackages()
	else:
		clean_dict = findDistfiles()
	# actually clean files if something was found
	if len(clean_dict.keys()):
		# verbose pretend message
		if myoptions['pretend'] and not myoptions['quiet']:
			einfo("Here are "+files_type+" that would be deleted:", \
			      myoptions['nocolor'])
		# verbose non-pretend message
		elif not myoptions['quiet']:
			einfo("Cleaning "+files_type+"...",myoptions['nocolor'])
		# do the cleanup, and get size of deleted files
		clean_size = doCleanup(clean_dict,action,myoptions)
		# vocabulary for final message
		if myoptions['pretend']: verb = "would be"
		else: verb = "has been"
		# display freed space
		if not myoptions['quiet']:
			einfo("Total space that "+verb+" freed in " \
			      + action + " directory: " \
			      + red(prettySize(clean_size)), \
			      myoptions['nocolor'])
	# nothing was found, return
	elif not myoptions['quiet']:
		einfo("Your "+action+" directory was already clean.", \
		      myoptions['nocolor'])


###############################################################################
# main: parse command line and execute all actions
def main():
	# set default options
	myoptions = {}
	myoptions['nocolor'] = port_settings["NOCOLOR"] in ('yes','true') \
	                       and sys.stdout.isatty()
	if myoptions['nocolor']: nocolor()
	# parse command line options and actions
	try: myactions, options = parseArgs(myoptions)
	# filter exception to know what message to display
	except ParseArgsException, e:
		if e.value == 'help':
			printUsage()
			sys.exit(0)
		elif e.value == 'version':
			printVersion()
			sys.exit(0)
		elif e.value in ('options','actions'):
			printUsage(e.value)
			sys.exit(2)
	# security check for non-pretend mode
	if not myoptions['pretend'] and portage.secpass != 2:
		eerror("Permission denied: you must be root.", \
		       myoptions['nocolor'])
		sys.exit(1)
	# execute actions
	for action in myactions:
		doAction(action, myoptions.copy())

###############################################################################
# actually call main() if launched as a script
if __name__ == "__main__":
	try: main()
	except KeyboardInterrupt:
		print "Aborted."
		sys.exit(130)
	sys.exit(0)

###############################################################################