#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Depends on:
sane
cuneiform
hocr2pdf and optimize2bw (both part of exactimage)
stapler
python-pyparallel to use parallel processing (i.e. OCR on multiple pdf files)


Graphic formats used:
	tiff:	scanned raw data
	png:	temporary files and image type converted into the final pdf document - PROBLEMS WITH CUNEIFORM AS OF 2012/06/03
	pdf:	final output

"""

## Configuration variables

device='plustek'			# Sane device refernece, see scanimage -L for available devices
#device='genesys'			# Klaus Scanner
SaveFormatScanOnly='jpg'	# Imageformat to use in scan-only mode

# Try to catch ctrl-c in a cleaner way
try:
	import signal
	def signal_handler(signal, frame):
		print("\nInterrupted")
		try:
			sys.exit(2)
		except NameError:
			import sys
			sys.exit(2)
		signal.signal(signal.SIGINT, signal_handler)
except KeyboardInterrupt:
	print("\nInterrupted")
	try:
		sys.exit(1)
	except NameError:
		import sys
		sys.exit(1)

import argparse
from subprocess import call, Popen, PIPE
import os			# env
import random
import string
import shutil
import sys
import shlex		# execute() function
#import pp
import glob			# wildcard file remove

parser = argparse.ArgumentParser(description='Scan and start OCR.')

## General Options ##

parser.add_argument('-o',
					'--output',
					dest='output',
					help='Output filename',
					required=True)
parser.add_argument('-m',
					'--mode',
					dest='mode',
					default='scanocr',
					help='Mode to use. Values are: scanocr, scan, ocr (default: scanocr)',
					required=False)
parser.add_argument('-v',
					'--verbose',
					dest='verbose',
					action='store_true',
					default=False,
					help='Be verbose.')

## Scanner Options ##

parser.add_argument('-r',
					'--resolution',
					dest='resolution',
					default='300',
					help='Resolution to use for scan (default: 300dpi)')
parser.add_argument('-sh',
					'--scanheight',
					dest='height',
					default='297',
					help='Height of document (default: 297mm - A4)')
parser.add_argument('-sw',
					'--scanwidth',
					dest='width',
					default='211',
					help='Width of document (default: 211mm - A4)')
parser.add_argument('-p',
					'--progress',
					dest='progress',
					action='store_true',
					default=True,
					help='Show scan progress (default: yes)')

## Converting Options ##

parser.add_argument('-c',
					'--color',
					dest='color',
					action='store_true',
					default=False,
					help='Preserve colors (Do not optimize image) (default: no)')

## OCR Options ##

parser.add_argument('-mu',
					'--multi',
					dest='multi',
					action='store_true',
					default=False,
					help='Create a multi-page PDF from more than one scan (default: no)')
parser.add_argument('-l',
					'--language',
					dest='language',
					default='ger',
					help='Language to use for OCR (default: german)')





# Executes given command, returns stdout/stderr
def execute(cmd, stdin='', env=''):
	cmd = shlex.split(cmd)

	if env == '':
		p = Popen(cmd, stdout=PIPE, stdin=PIPE, stderr=PIPE)
	else:
		p = Popen(cmd, stdout=PIPE, stdin=PIPE, stderr=PIPE, env=env)

	stdout, stderr = p.communicate(stdin)
	returncode = p.wait()

	if returncode == 0:
		return (0, stdout)
	else:
		return (returncode, str(stderr))

# Executes scanimage (sane)
def scan(resolution, width, height, progress):

	print('Scan: \t\t', end='')

	progress = False

	if (progress == True):
			StringProgress = ' -p'
	else:
			StringProgress = ''
#--calibration-cache=yes
	call('scanimage --format=tiff -d ' + device + ' --resolution '+resolution+' -x '+width+' -y '+height+StringProgress+' >' +TempFile+ '.tiff', shell=True)

	print('Done.')

	"""
	p = execute('/usr/bin/scanimage --format=tiff -d ' + device + ' --calibration-cache=yes --resolution ' + resolution + ' -x ' + width + ' -y ' + height + ' ' + StringProgress)

	if p[0] != 0:
		sys.exit ('Error! ' + str(p[1]))
	else:
		print('Done.')
	
	"""

# Optimize Image to black/white, will write a new tiff file
def optimizeImage(bw):

	print('Optimization: \t', end='')

	p = execute('/usr/bin/optimize2bw -i ' + TempFile + bw +'.tiff -o ' + TempFile + '.bw.tiff')

	if p[0] != 0:
		sys.exit ('Error! ' + p[1])
	else:
		print('Done.')

# Convert tiff to SaveFormatScanOnly (used when scanning only)
def convertToImage(bw):
	
	global SaveFormatScanOnly

	print('Convert: \t', end='')

	p = execute('convert '+TempFile + bw + '.tiff '+SaveFile+'.'+SaveFormatScanOnly)

	if p[0] != 0:
		sys.exit ('Error converting tiff to ' + SaveFormatScanOnly + ': ' + p[1])
	else:
		print('Done.')

# Do ocr on a single image file
def ocr(save, language):

	print('OCR: \t', end='')

	# MALLOC_CHECK_ prevents abort due to glibc errors
	# Hilft noch nicht ganz!
	#call('MALLOC_CHECK_=0 cuneiform -l ' +language+ ' -f hocr -o ' +TempFile+ '.hOCR ' +TempFile+ '.bw.tiff '+verbosity(), shell=True)

	my_env = os.environ
	my_env["MALLOC_CHECK_"] = '0'
	#my_env.get("PATH", '')

	# To use cuneiform as OCR-engine
	#p = execute('cuneiform -l ' +language+ ' -f hocr -o ' +TempFile+ '.hOCR ' +TempFile+ '.bw.tiff', '', my_env)
	
	# To use tesseract as OCR-engine
	if language == 'ger':
		language = 'deu'

	p = execute('tesseract ' +TempFile+ '.bw.tiff ' +TempFile+ ' -l ' +language+ ' hocr', '', my_env)
	# Rename temporary file created by tesseract (.html -> .hOCR)
	os.rename(TempFile+'.html', TempFile+'.html'[:-4] + "hOCR")

	if p[0] != 0:
		sys.exit ('Error! ' + p[1])
	else:
		print('Done.')


	# Create PDF || -s: Alternative word-grouping
	call('hocr2pdf -s -i '+TempFile+'.bw.tiff -o '+save+'.pdf < ' +TempFile+ '.hOCR ', shell=True)



# Do some final optimizing on output PDF
def OptimizePDF():

	print('PDF Optimization: \t', end='')
	
	p = execute('/usr/bin/pdfopt ' + TempFile + '.pdf ' + TempFile + '.pdf')

	if p[0] != 0:
		print ('Error! ' + p[1])
	else:
		print('Done.')


# Remove garbage files, if existing
def CleanUp(filename = ''):

	if (filename == ''):
		filename = TempFile

	for filename in glob.glob(filename + '*.*') :
		os.remove(filename)


# Set Verbosity string
def verbosity():
	if (verbose == False):
		return '>/dev/null 2>&1'
	else:
		return ''


# Returns True in case the output file already exists.
def OutputFileExists(SaveFile):
	args = parser.parse_args()

	if ( (os.path.exists(SaveFile + '.pdf') == True and ( args.mode == 'scanocr' or args.mode == 'ocr' ) ) 
	or (os.path.exists(SaveFile + '.' + SaveFormatScanOnly) == True and args.mode == 'scan') ):
		return True
	else:
		return False
		

# Main program
def main():
	args = parser.parse_args()

	# Mode switcher - ask for output name if not given
	print('Current Mode: ', end='')
	if args.mode == 'scanocr':
		print('Scan and OCR')
		if args.output == None:
			args.output = str(input('Please give some output filename (without .pdf): '))
	elif args.mode == 'scan':
		print('Scan only')
		if args.output == None:
			args.output = str(input('Please give some output filename (without .pdf): '))
	elif args.mode == 'ocr':
		print('OCR only')
		print('FEHLT NOCH')
		sys.exit(0) # For now
	else:
		sys.exit('No valid mode given. Please choose one of: scanocr, scan, ocr')

	global TempFile, SaveFile, verbose

	# Process verbose-switch ...
	verbose = args.verbose

	# Path and name (random string) of temporary file without extension
	RandomString = ''.join(random.choice(string.ascii_uppercase + string.digits) for x in range(6))
	TempFile = '/tmp/scanner_'+ RandomString + args.output
	SaveFile = os.getcwd() + '/' + args.output

	# check if output file exists
	if OutputFileExists(SaveFile) == True:
		sys.exit('Output file exists.')

	# Create multi-page OCR'ed PDF
	if  (args.multi == True):
		Pages = PagesOrig = int(input('How many pages does your document contain? '))
		
		if (Pages < 1):
			sys.exit('Not a valid amount of pages.')
		
		for Pages in range(1, Pages+1):
			# Get PageNR in TempFile
			TempFile = '/tmp/scanner_'+ RandomString + args.output + '_' + str(Pages)
			print('Processing page '+str(Pages)+'...')
			Pages -= 1

			# Invoke scan, optimize and OCR
			scan(args.resolution, args.width, args.height, args.progress)

			if (args.color == True):
				convertToImage('')
			else:
				optimizeImage('')

			# Do OCR or convert to pdf without ocr'ing
			if args.mode != 'scan':
				ocr(TempFile+'_ready', args.language)
			else:
				if (args.color == True):
					call('convert '+TempFile+'.tiff '+TempFile+'_ready.pdf', shell=True)
				else:
					call('convert '+TempFile+'.bw.tiff '+TempFile+'_ready.pdf', shell=True)

			# Wait for user to continue
			if (Pages < PagesOrig-1):
				input('Hit enter to scan the next page (' + str(Pages+1) + ')...')

		# Merge PDF's and exit
		call('stapler cat /tmp/scanner_'+ RandomString + args.output + '_*_ready.pdf '+SaveFile+'.pdf', shell=True)

		# Clean up temporary files
		CleanUp('/tmp/scanner_'+ RandomString + args.output)

		# Wie wildcard ersetzen?
		#Popen(['stapler', 'cat', '/tmp/scanner_'+ RandomString + args.output + '_*_ready.pdf', SaveFile+'.pdf'])

		sys.exit(0)


	# Normal operation if not in multipage-mode, scan first.
	scan(args.resolution, args.width, args.height, args.progress)

	# optimize
	if (args.color == False):
		optimizeImage('')
		blackwhite='.bw'
	else:
		blackwhite=''

	# convert .tiff to image-format
	if (args.mode == 'scan'):
		convertToImage(blackwhite)

	# ocr
	if (args.mode == 'scanocr'):
		ocr(SaveFile, args.language)

	CleanUp()

	# End program
	sys.exit(0)


if __name__ == "__main__":
		main()
