#!/usr/bin/python

import gc
import os
import re
import sys
import shutil
import logging
import getopt
import marshal
import binascii
import tempfile
from subprocess import Popen, PIPE, STDOUT

#### MODIFY THIS SECTION AS APPROPRIATE 

P4CMD="/usr/local/bin/p4"
P4PORT="localhost:1666"
# P4USER should have super access on the system, as it needs to update counters, etc.
P4USER="p4super"
# PASSFILE is not required, as long as the user above has a valid ticket (e.g. unlimited)
PASSFILE="superpass"
P4CLIENT="thumb_client"
CONVERTCMD="/opt/local/bin/convert"
THUMBNAIL_GENERATOR_COUNTER="thumb-generator"
THUMBNAIL_FORMAT=".png"
THUMBNAIL_SIZE="x160"
# TEMP_DIR should be a temporary location where you can write. The directory will be
# created, if necessary. it will be cleaned out, so ensure that you are not pointing
# to a directory with existing information
TEMP_DIR = "/tmp/thumbs"
EXPLICIT_GC = True
LOCK_FILE = "/tmp/thumbnail-generator.lock"

#### END MODIFICATION SECTION

#### GLOBAL VARIABLES

forceUnlock = False

####

# usage
#     provides usage information for the utility
def usage():
	PROGRAM_USAGE = """\
Usage:	ThumbnailGenerator.py [-h] [-f depotFile] [-c changelistNumber] [-l logfile] [-v] [-g] [-i] 
           -h                   show help (this message)
           -v                   verbose (debug mode)
           -f depotFile         specify file (depot path and version //path/to/file#nn)
           -c changelistNumber  process only the specified changelist number and exit
           -p pageNum           specify page number of file (relevant for PDF files; default is 0 or first page)
           -l logfile           specify logfile (default is STDOUT)
           -g                   generate sample client specification
           -i                   initialize the system, creating the logger counter and processing all changelists
           -x                   ignore the lock file and run anyway
"""
	print(PROGRAM_USAGE)
	normalExit()

# generateSampleClient
#    this routine generates a sample client specification. The output should be sanity
#    checked, but can make creating the client spec a little easier.
def generateSampleClient():
	CLIENT_PREAMBLE = """\
Client:	{0}
Owner:	{1}
Description:
	Created by p4admin.
Root:	{2}
Options:	noallwrite noclobber nocompress unlocked nomodtime normdir
SubmitOptions:	submitunchanged
LineEnd:	local""".format(P4CLIENT, P4USER, TEMP_DIR)
	CLIENT_VIEW = """\
	//{0}/....ppm //{1}/{0}/....ppm
	//{0}/....PPM //{1}/{0}/....PPM
	//{0}/....bmp //{1}/{0}/....bmp
	//{0}/....BMP //{1}/{0}/....BMP
	//{0}/....jpg //{1}/{0}/....jpg
	//{0}/....JPG //{1}/{0}/....JPG
	//{0}/....bpm //{1}/{0}/....bpm
	//{0}/....BPM //{1}/{0}/....BPM
	//{0}/....gif //{1}/{0}/....gif
	//{0}/....GIF //{1}/{0}/....GIF
	//{0}/....pgm //{1}/{0}/....pgm
	//{0}/....PGM //{1}/{0}/....PGM
	//{0}/....png //{1}/{0}/....png
	//{0}/....PNG //{1}/{0}/....PNG
	//{0}/....xbm //{1}/{0}/....xbm
	//{0}/....XBM //{1}/{0}/....XBM
	//{0}/....xpm //{1}/{0}/....xpm
	//{0}/....XPM //{1}/{0}/....XPM
	//{0}/....tga //{1}/{0}/....tga
	//{0}/....TGA //{1}/{0}/....TGA
	//{0}/....psd //{1}/{0}/....psd
	//{0}/....PSD //{1}/{0}/....PSD
	//{0}/....PDF //{1}/{0}/....PDF
	//{0}/....pdf //{1}/{0}/....pdf
	//{0}/....AI //{1}/{0}/....AI
	//{0}/....ai //{1}/{0}/....ai
"""
	depotList = []
	cmd = ["depots"]
	results = p4MarshalCmd(cmd)
	for r in results:
		if 'type' in r and (r['type'] == "local" or r['type'] == "stream"):
			depotList.append(r['name'])
	print(CLIENT_PREAMBLE)
	print("View:")
	for d in depotList:
		print(CLIENT_VIEW.format(d, P4CLIENT))
	normalExit()

# normalExit
#     exits 
def normalExit(code = 0, msg = None):
	if(msg != None):
		print msg;
	if not forceUnlock and os.path.exists(LOCK_FILE):
		# remove the lock file		
		os.unlink(LOCK_FILE)
	sys.exit(code)
	
# errorExit
#     exits after logging the error message -- does not remove the lock file
def errorExit(msg):
	print("### ERROR: " + msg)
	sys.exit(9)

# cleanTmpDir
#     cleans up the temporary directory 
def cleanTempDir():
	for file_object in os.listdir(TEMP_DIR):
		file_object_path = os.path.join(TEMP_DIR, file_object)
		if os.path.isfile(file_object_path):
			os.unlink(file_object_path)
		else:
			shutil.rmtree(file_object_path)

# p4MarshalCmd
#     executes the p4 command, results sent to a list
def p4MarshalCmd(cmd,quiet=False):
	if not quiet:
		logging.debug("p4 {0}".format(" ".join(cmd)))
	list = []
	pipe = Popen([P4CMD, "-p", P4PORT, "-u", P4USER, "-c", P4CLIENT, "-G"] + cmd, stdout=PIPE).stdout
	try:
		while 1:
			record = marshal.load(pipe)
			list.append(record)
	except EOFError:
		pass
	pipe.close()
	return list

# p4InputCmd
#     executes the p4 command with input
def p4InputCmd(data,cmd,quiet=False):
	if not quiet:
		logging.debug("p4 {0}".format(" ".join(cmd)))
	list = []
	proc = Popen([P4CMD, "-p", P4PORT, "-u", P4USER, "-c", P4CLIENT, "-G"] + cmd, stdout=PIPE, stdin=PIPE, stderr=PIPE)
	outPipe = proc.stdout
	proc.stdin.write(data)
	return proc.communicate()

# p4Cmd
#     executes a p4 command, returns results
def p4Cmd(cmd,quiet=False):
	if not quiet:
		logging.debug("p4 {0}".format(" ".join(cmd)))
	proc = Popen([P4CMD, "-p", P4PORT, "-u", P4USER, "-c", P4CLIENT] + cmd, stdout=PIPE, stderr=PIPE)
	return proc.communicate()

# p4Cmd
#     executes a p4 command, returns results
def convert(original,thumb,page=0):
	logging.debug("convert -thumbnail {0} {1}[{2}] {3}".format(THUMBNAIL_SIZE, original, page, thumb))
	proc = Popen([CONVERTCMD, "-thumbnail", THUMBNAIL_SIZE, original + "[{0}]".format(page), thumb], stdout=PIPE, stderr=PIPE)
	return proc.communicate()

# containsError
#     utility function to check for any error code in the results array
def containsError(results=[],logError=True):
	foundError = False
	for r in results:
		if 'code' in r:
			if r['code'] == 'error':
				foundError = True
				if logError:
					logging.error(r['data'].strip())
			elif r['code'] == 'info':
				#code info output can be important in troubleshooting
				logging.debug(r['data'])
	return foundError

# checkLogin 
#     check the login ticket on the server
def checkLogin(username=""):
	cmd = ["login","-s"]
	result = p4MarshalCmd(cmd,quiet=True)
	if containsError(result, False):
		return False
	else:
		return True

# login
#    logs the user in using the password in the password file
def login():
	if PASSFILE is not None and os.path.isfile(PASSFILE):
		f = open(PASSFILE)
		lines = f.readlines()
		f.close()
		adminpass = lines[0].strip()
		cmd = ["login"]
		result = p4InputCmd(adminpass, cmd)
		if len(result[1]) > 0:
			return False
		else:
			return True
	else:
		return False

# isSubmitted
#    checks to see if the changelist exists and that it is a 'submitted' state
def isSubmitted(change):
	cmd = ["describe", "-s", change]
	result=p4MarshalCmd(cmd)
	if containsError(result, False):
		return False
	if result[0]['status'] == "submitted":
		return True
	else:
		return False

# setLoggerCounter
def resetCounters():
	cmd = ["counter", "logger", "0"]
	p4Cmd(cmd)
	cmd = ["counter", "-d", THUMBNAIL_GENERATOR_COUNTER]
	p4Cmd(cmd)

# setThumbAttribute
#    sets the 'thumb' attribute on the revision to the hex value specified
def setThumbAttribute(revision, hex):
	cmd = ["attribute", "-e", "-f", "-n", "thumb", "-i", revision]
	result = p4InputCmd(hex, cmd)

# generateThumbnail
#    generates a thumbnail for the specified depot file
def generateThumbnail(depotFile, page=0):
	logging.debug("generating thumbnail for " + depotFile)
	# create TEMP_DIR, if it does not exist
	if not os.path.exists(TEMP_DIR):
		os.makedirs(TEMP_DIR)
		
	file = None
	rev = 0
	matchObj = re.match( r'(.*)#(\d+)$', depotFile, re.M|re.I)
	if matchObj:
		file = matchObj.group(1)
		rev = int(matchObj.group(2))
	else:
		file = depotFile
	extension = os.path.splitext(file)[1]
	p4print = tempfile.NamedTemporaryFile(dir=TEMP_DIR, prefix="p4_", suffix=extension)
	p4print.close()
	thumb = tempfile.NamedTemporaryFile(dir=TEMP_DIR, prefix="thumb_", suffix=THUMBNAIL_FORMAT)
	thumb.close()
	logging.debug("getting file from depot...")
	cmd = ["print", "-q", "-o", p4print.name, depotFile]
	result = p4Cmd(cmd)
	logging.debug("trying to generate thumbnail...")
	convert(p4print.name, thumb.name, page)
	if os.path.exists(thumb.name):
		file = open(thumb.name, "rb")
		bytes = file.read()
		file.close()
		hex = binascii.hexlify(bytes)
		logging.debug("storing thumbnail in metadata...")
		setThumbAttribute(depotFile, hex)
	# clean up the temporary directory
	cleanTempDir()

# processChange
#    processes the files in the specified changelist
def processChange(change):
	path = "//...@{0},@{0}".format(change)
	cmd = ["sync", "-p", "-n", path]
	result=p4MarshalCmd(cmd)
	if not containsError(result, False):
		for r in result:
			if r['action'] == "deleted":
				continue
			depotFile = "{0}#{1}".format(r['depotFile'], r['rev'])
			generateThumbnail(depotFile, 0)
			if(EXPLICIT_GC):
				# explicitly invoke garbage collection because thumbnails are being
				# read into memory and we want to ensure that memory is freed as
				# quickly as possible
				gc.collect()
			
# initialConfiguration
#    routine to initialize the system by (re)setting the logger counter, deleting
#    the thumbnail generator counter, and processing all of the submitted changelists
#    starting from the beginning. Needless to say, this can potentially take a long
#    time to run. if you simply want to turn on the functionality, you just have
#    to create the logger counter (p4 counter logger 0) and then schedule the thumbnail
#    generator to run. Only new changelists will be processed, however (i.e. only new
#    versions of files will have thumbnails created)
def initialConfiguration():
	logging.info("performing initial configuration")
	resetCounters()
	cmd = ["changes"]
	result = p4MarshalCmd(cmd)
	if(containsError(result, False)):
		errorExit("Error finding updates")
	for r in reversed(result):
		if r['status'] == "submitted":
			change = r['change']
			processChange(change)
	
# checkForChanges
#    uses the logger command to find any changelists since the last run
def checkForChanges():
	cmd=["logger", "-t", THUMBNAIL_GENERATOR_COUNTER]
	result=p4MarshalCmd(cmd)
	if(containsError(result, False)):
		errorExit("Error finding updates")
	last = 0
	for r in result:
		sequence = r['sequence']
		change = r['attr']
		if (int(change) > last) and isSubmitted(change):
			logging.info("Processing changelist @{0}".format(change))
			processChange(change)
			last = int(change)
		logging.debug("updating logger counter...")
		cmd=["logger", "-c", sequence, "-t", THUMBNAIL_GENERATOR_COUNTER]
		p4Cmd(cmd)

###########################################################################
##### MAIN PROGRAM STARTS HERE
#####

def main(argv=None):

	global forceUnlock
	
	verbose = False
	logFile = None
	depotFile = None
	changeId = 0
	pageNum = 0
	generateSample = False
	initialize = False
	
	try:
		opts, args = getopt.getopt(argv, "hl:f:c:p:vgix")
	
		for opt, arg in opts:
			if opt == "-v":
				verbose = True
			elif opt == "-h":
				usage()
			elif opt == "-l":
				logFile = arg
			elif opt == "-f":
				depotFile = arg
			elif opt == "-c":
				changeId = int(arg)
			elif opt == "-p":
				pageNum = int(arg)
			elif opt == "-g":
				generateSample = True
			elif opt == "-i":
				initialize = True
			elif opt == "-x":
				forceUnlock = True
	
		logLevel = logging.WARN
		if verbose:
			logLevel = logging.DEBUG
			
		if logFile is not None:
			logging.basicConfig(filename=logFile, format='%(asctime)s [%(levelname)s] %(message)s', datefmt='%m/%d/%Y %I:%M:%S %p', level=logLevel)
		else:
			logging.basicConfig(format='[%(levelname)s] %(message)s', datefmt='%m/%d/%Y %I:%M:%S %p', level=logLevel)

		# check to see if the lockfile is present
		if os.path.exists(LOCK_FILE):
			if not forceUnlock:
				errorExit("Program cannot run... lockfile exists -- " + LOCK_FILE)
			
		#  if we get to this point, create the lock file as long as we're not in a force condition
		if not forceUnlock:
			open(LOCK_FILE, 'a').close()
		
		# check to see if the user is logged in, and if not then log in
		if not checkLogin():
			if not login():
				errorExit("Not logged in")
				
		if generateSample:
			generateSampleClient()
		elif depotFile is not None:
			generateThumbnail(depotFile, pageNum)
		elif changeId != 0:
			processChange(changeId)
		else:
			if initialize:
				initialConfiguration()
			else:
				checkForChanges()
		
		normalExit()

	except getopt.GetoptError as e:
		print(e)
		print("ERROR: unknown argument\n")
		usage()
		normalExit(2)

if __name__ == '__main__':
	main(sys.argv[1:])