#!/usr/bin/python
# -*- coding: utf-8 -*-

import sys
import optparse
import os
import pg
import types
import urllib
import urlparse
import datetime
import re
import json
import time
from pianos import parseConnectionString


	
class Error( Exception ):
	pass
	
	
class Database:

	def __init__( self, params ):
		self.db = None
		self.params = params
		
	def connect( self ):
		self.db = pg.connect( self.params[2], host = self.params[0], port = self.params[1], user = self.params[3], passwd = self.params[4] )
		
class sqlverbatim:
	def __init__( self, value ):
		self.value = value
	def __str__( self ):
		return str( self.value )

def sqlarg( args ):
	all = []
	for a in args:
		if type( a ) is types.InstanceType:
			t = 'verbatim'
		else:
			if a is None:
				t = 'null'
			elif type( a ) is types.StringType:
				t = 'string'
			elif type( a ) is types.BooleanType:
				t = 'bool'
			else:
				t = 'verbatim'
			
		if t == 'null':
			v = 'NULL'
		elif t == 'string':
			v = "E'" + pg.escape_string( a ) + "'"
		elif t == 'bool':
			v = 'true' if a else 'false'
		else:
			v = str( a )

		all.append( v )
		
	return ','.join( all )

class Pianos3Database( Database ):
	
	def __init__( self, params ):
		Database.__init__( self, params );
		self.imageTypes = {
			'TRACE': 'MARK',
			'PRINT': 'PRINT',
		}

	def init( self ):
		pass
	
	def getUsers( self ):
		return self.db.query( "SELECT trustee.*,person.* FROM trustee,person WHERE trustee.id = person.id;" ).dictresult()
	
	def getGroups( self ):
		groups = self.db.query( "SELECT trustee.*,persongroup.* FROM trustee,persongroup WHERE trustee.id = persongroup.id;" ).dictresult()
		for g in groups:
			g['members'] = [x[0] for x in self.db.query( "SELECT person FROM groupbinding WHERE persongroup = " + str( g['id'] ) + ";" ).getresult()]
		return groups

	def getImages( self ):
		return filter( lambda x: x['typeshort'] in self.imageTypes, self.db.query( "SELECT object.*,imageview.*,imagetype.short AS typeshort FROM imageview,object,imagetype WHERE imageview.id = object.id AND imagetype.id=imageview.type;" ).dictresult() )
	def getImageData( self, id ):
		return pg.unescape_bytea( self.db.query( "SELECT data FROM image WHERE id = " + str( id ) + ";" ).dictresult()[0]['data'] )
		
	def getFolders( self ):
		folders = self.db.query( "SELECT object.*,exercisegroup.* FROM object,exercisegroup WHERE object.id = exercisegroup.id;" ).dictresult()
		for f in folders:
			f['security'] = self.db.query( "SELECT security.trustee AS pgroupid,persongroup.name AS pgroupname FROM security,persongroup WHERE security.object = " + str( f['id'] ) + " AND security.trustee = persongroup.id" ).dictresult()
		return folders
		
	def getExercises( self ):
		return self.db.query( "SELECT object.*,exercise.* FROM object,exercise WHERE object.id = exercise.id;" ).dictresult()
		
	def getAssignments( self ):
		return self.db.query( "SELECT object.*,assignment.* FROM object,assignment WHERE object.id = assignment.id;" ).dictresult()
		
	def getAssignmentResults( self, aid ):
		results = self.db.query( "SELECT object.*,result.*,resulttype.short AS typeshort FROM object,result,resulttype WHERE object.id = result.id AND result.type = resulttype.id AND result.iscomputelr IS false AND result.assignment = " + str( aid ) + ";" ).dictresult()
		ret = {}
		for result in results:
			result['validated'] = result['validated'] == 't'
			result['minutiae'] = self.db.query( "SELECT annotation.*,minutia.*,minutiatype.short AS typeshort FROM annotationtype,annotation,minutia,minutiatype WHERE annotationtype.short = 'MINUTIA' AND annotationtype.id = annotation.annotationtype AND annotation.id = minutia.id AND minutia.type = minutiatype.id AND annotation.result = " + str( result['id'] ) + ";" ).dictresult()
			result['ridges'] = self.db.query( "SELECT annotation.*,polyline.* FROM annotationtype,annotation,polyline WHERE annotationtype.short = 'POLYLINE' AND annotationtype.id = annotation.annotationtype AND annotation.id = polyline.id AND annotation.result = " + str( result['id'] ) + ";" ).dictresult()
			ret[result['typeshort']] = result
		return ret
		
			
class Pianos4Interface:
	
	def __init__( self ):
		self.source = None
		self.nextID = None
		self.idmap = {}
		self.tidmap = {}
		self.imageTypes = {
			'TRACE': 'MARK',
			'PRINT': 'PRINT',
		}
		self.minutiaeTypes = {
			'BIFURCATION': 'bifurcation',
			'RIDGE_ENDING': 'ridge_ending',
			'CENTER': 'center',
			'UNKNOWN': 'unknown',
		}
		self.ridgeRE = re.compile( r'\(x=([\d\.]+), y=([\d\.]+)\)' )
		self.eraseExisting = False
		
	def setSource( self, source ):
		self.source = source
	
	def emit( self, sql ):
		sys.stdout.write( sql + '\n' )
	
	def getUsername( self, username ):
		return username
		
	def init( self ):
		self.emit( "BEGIN;" )
		
		if self.eraseExisting:
			print "-- Erasing existing data"
			self.emit( "SET client_min_messages TO 'WARNING';" )
			self.emit( "TRUNCATE object CASCADE;" )
			self.emit( "SET client_min_messages TO 'NOTICE';" )
		
		self.nextID = 1
		
	def close( self ):
		self.emit( "SELECT setval('object_id_seq', " + str( self.nextID ) + ", false);" )
		self.emit( "COMMIT;" )
		print "-- Performing maintenance"
		print "--  Vacuum"
		self.emit( "SET client_min_messages TO 'ERROR';" )
		self.emit( "VACUUM FULL ANALYZE;" )
		self.emit( "SET client_min_messages TO 'NOTICE';" )

	def getID( self, id ):
		if id in self.idmap:
			return self.idmap[id]
		newid = self.nextID
		self.nextID += 1
		self.idmap[id] = newid
		return newid
		
	def getTrusteeID( self, id ):
		if id in self.tidmap:
			return self.tidmap[id]
		newid = self.nextID
		self.nextID += 1
		self.tidmap[id] = newid
		return newid
		
	def copyUsers( self ):
		users = self.source.getUsers()
		print "-- Found", len( users ), "users"

		self.emit( "ALTER TABLE object ALTER COLUMN creator DROP NOT NULL;" )
		
		for user in users:
			self.emit( "INSERT INTO object(id,createdon,updatedon) VALUES (" + sqlarg( [self.getTrusteeID( user['id'] ), user['createdon'], sqlverbatim( "'now'" )] ) + ");" )
			self.emit( "INSERT INTO trustee(id,isgroup) VALUES (" + sqlarg( [self.getTrusteeID( user['id'] ), False] ) + ");" )
			self.emit( "INSERT INTO person(id,username,name,password,active,email,language) VALUES (" + sqlarg( [self.getTrusteeID( user['id'] ), self.getUsername( user['username'] ), user['displayname'], user['password'], user['active'] == 't', '', sqlverbatim( "(SELECT id FROM language ORDER BY id ASC LIMIT 1)" )] ) + ");" )

		for user in users:
			self.emit( "UPDATE object SET creator = " + str( self.getTrusteeID( user['creator'] ) ) + " WHERE id = " + str( self.getTrusteeID( user['id'] ) ) + ";" )
			
		self.emit( "ALTER TABLE object ALTER COLUMN creator SET NOT NULL;" )

	def copyGroups( self ):
		groups = self.source.getGroups()
		print "-- Found", len( groups ), "groups"
		
		for group in groups:
			self.emit( "INSERT INTO object(id,createdon,updatedon,creator) VALUES (" + sqlarg( [self.getTrusteeID( group['id'] ), group['createdon'], sqlverbatim( "'now'" ), self.getTrusteeID( group['creator'] )] ) + ");" )
			self.emit( "INSERT INTO trustee(id,isgroup) VALUES (" + sqlarg( [self.getTrusteeID( group['id'] ), True] ) + ");" )
			self.emit( "INSERT INTO pgroup(id,name,tutors) VALUES (" + sqlarg( [self.getTrusteeID( group['id'] ), group['displayname'], group['tutor'] == 't'] ) + ");" )
			for uid in group['members']:
				self.emit( "INSERT INTO membership(person,pgroup) VALUES (" + sqlarg( [self.getTrusteeID( uid ), self.getTrusteeID( group['id'] )] ) + ");" )

		
	def copyImagesDef( self ):
		images = self.source.getImages()
		print "-- Found", len( images ), "images"
		
		for image in images:
			self.emit( "INSERT INTO object(id,createdon,updatedon,creator) VALUES (" + sqlarg( [self.getID( image['id'] ), image['createdon'], sqlverbatim( "'now'" ), self.getTrusteeID( image['creator'] )] ) + ");" )
			self.emit( "INSERT INTO image(id,type,name,mimetype,info,resolution,width,height,data) VALUES (" + sqlarg( [self.getID( image['id'] ), sqlverbatim( "(SELECT id FROM imagetype WHERE short = '" + self.imageTypes[image['typeshort']] + "')" ), image['name'], image['mimetype'], image['info'], image['resolution'], image['width'], image['height']] ) + ",E'0');" )
	
	def copyFolders( self ):
		folders = self.source.getFolders()
		print "-- Found", len( folders ), "folders"
		
		for folder in folders:
			self.emit( "INSERT INTO object(id,createdon,updatedon,creator) VALUES (" + sqlarg( [self.getID( folder['id'] ), folder['createdon'], sqlverbatim( "'now'" ), self.getTrusteeID( folder['creator'] )] ) + ");" )
			self.emit( "INSERT INTO folder(id,name) VALUES (" + sqlarg( [self.getID( folder['id'] ), folder['name']] ) + ");" )
			for s in folder['security']:
				self.emit( "INSERT INTO security(object,trustee) VALUES (" + sqlarg( [self.getID( folder['id'] ), self.getTrusteeID( s['pgroupid'] )] ) + ");" )

	def copyExercises( self ):
		exercises = self.source.getExercises()
		print "-- Found", len( exercises ), "exercises"
		
		for ex in exercises:
			self.emit( "INSERT INTO object(id,createdon,updatedon,creator) VALUES (" + sqlarg( [self.getID( ex['id'] ), ex['createdon'], sqlverbatim( "'now'" ), self.getTrusteeID( ex['creator'] )] ) + ");" )
			self.emit( "INSERT INTO exercise(id,name,description,folder,mark,print) VALUES (" + sqlarg( [self.getID( ex['id'] ), ex['name'], ex['description'], self.getID( ex['exercisegroup'] ), self.getID( ex['traceimage'] ), self.getID( ex['printimage'] )] ) + ");" )

	def copyImagesData( self ):
		images = self.source.getImages()
		print "-- Updating", len( images ), "images"
		
		icount = 0
		icountshow = 10
		
		for image in images:
			data = self.source.getImageData( image['id'] )
			self.emit( "UPDATE image SET data = E'" + pg.escape_bytea( data ) + "' WHERE id = " + str( self.getID( image['id'] ) ) + ";" )
			
			icount += 1
			if icount % icountshow == 0:
				print "--  ", icount, "images done"
	
		print "--  ", icount, "images done"
			
			
	def copyAssignments( self ):
		assignments = self.source.getAssignments()
		print "-- Found", len( assignments ), "assignments"

		timestamps_keys = [
			'analysis_annotations',
			'current_analysis',
			'validated_analysis',
			'comparison_annotations',
			'current_comparison',
			'validated_comparison',
		]
		timestamps = {}
		now = datetime.datetime.utcnow()
		i = 0
		for key in timestamps_keys:
			offset = datetime.timedelta( minutes = i )
			timestamps[key] = ( now + offset ).isoformat() + 'Z'
			i += 1
			
		time.sleep( 1 )
		timestampValidated = datetime.datetime.utcnow().isoformat() + 'Z'
		
		acount = 0
		acountshow = 100
		
		for a in assignments:
			self.emit( "INSERT INTO object(id,createdon,updatedon,creator) VALUES (" + sqlarg( [self.getID( a['id'] ), a['createdon'], sqlverbatim( "'now'" ), self.getTrusteeID( a['creator'] )] ) + ");" )
			self.emit( "INSERT INTO assignment(id,exercise,ispreview) VALUES (" + sqlarg( [self.getID( a['id'] ), self.getID( a['exercise'] ), False] ) + ");" )
	
			results = self.source.getAssignmentResults( a['id'] )
			# print "-- Found",len(results),"results for assignment",a['id']

			if len( results ):
			
				newres = {}
				
				if 'ANALYSIS_TRACE' in results:
					jsdata = {}
					r = results['ANALYSIS_TRACE']
					jsdata['mark'], nextID = self.convertResult( r, timestamp = timestamps['analysis_annotations'] )
					newres['ANALYSIS'] = ( jsdata, r, 'analysis' )
				
				if 'COMPARISON_TRACE' in results or 'COMPARISON_PRINT' in results:
					jsdata = {}
					r1, r2 = None, None
					if 'COMPARISON_TRACE' in results:
						r1 = results['COMPARISON_TRACE']
						jsdata['mark'], nextID = self.convertResult( r1, timestamp = timestamps['analysis_annotations'], timestamp_minutia_compadded = timestamps['comparison_annotations'] )
					else:
						jsdata['mark'] = {}
						nextID = 1
					if 'COMPARISON_PRINT' in results:
						r2 = results['COMPARISON_PRINT']
						jsdata['print'], nextID = self.convertResult( r2, nextID = nextID, timestamp = timestamps['comparison_annotations'] )
					else:
						jsdata['print'] = {}
					
					r = r2 if r2 else r1
					
					newres['COMPARISON'] = ( jsdata, r, 'comparison' )
					
				toAdd = []
				for t, ( jsdata, r, tskey ) in newres.iteritems():
					if len( r['conclusion'] ):
						jsdata['observations'] = {'conclusion': r['conclusion']}
					toAdd.append( [t, 'CURRENT', jsdata, timestamps['current_' + tskey]] )
					if r['validated']:
						toAdd.append( [t, 'VALIDATED', jsdata, timestamps['validated_' + tskey]] )

				toAdd.sort( cmp = lambda x, y: cmp( x[0], y[0] ) )  # Make sure ANALYSIS comes before COMPARISON (alphabetic sort works)
				
				for res in toAdd:
					self.emit( "INSERT INTO result(assignment,type,step,data,updatedon) VALUES (" + sqlarg( [self.getID( a['id'] ), sqlverbatim( "(SELECT id FROM resulttype WHERE short = '" + res[1] + "')" ), sqlverbatim( "(SELECT id FROM resultstep WHERE short = '" + res[0] + "')" ), json.dumps( res[2] ), res[3]] ) + ");" )
				
			acount += 1
			if acount % acountshow == 0:
				print "--  ", acount, "assignments processed"
		
		print "--  ", acount, "assignments processed"
		
	def convertResult( self, result, nextID = 1, timestamp = None, timestamp_minutia_compadded = None ):
		id = nextID
		if timestamp is None:
			timestamp = datetime.datetime.utcnow().isoformat() + 'Z' 
		if timestamp_minutia_compadded is None:
			timestamp_minutia_compadded = timestamp
		mres = {}
		for m in result['minutiae']:
			json = {
				'angle': m['theta'],
				'created': timestamp,
				'id': id,
				'modified': timestamp if ( 'comparisontraceadded' in m and m['comparisontraceadded'] == 'f' ) else timestamp_minutia_compadded,
				'pos': [m['xpos'], m['ypos']],
				'type': self.minutiaeTypes[m['typeshort']],
				'userid': self.getTrusteeID( result['creator'] ),
			}
			mres[id] = json
			id += 1
		rres = {}
		for r in result['ridges']:
			json = {
				'color': r['color'],
				'created': timestamp,
				'id': id,
				'modified': timestamp,
				'opacity': 0.85,
				'points': [( float( x[0] ), float( x[1] ) ) for x in self.ridgeRE.findall( r['points'] )],
				'userid': self.getTrusteeID( result['creator'] ),
				'width': r['width']
			}
			rres[id] = json
			id += 1
		res = {
			'minutiae': mres,
			'ridges': rres,
		}
		return res, id	
			
	
class Pianos4Database( Database, Pianos4Interface ):
	
	def __init__( self, params ):
		Database.__init__( self, params );
		Pianos4Interface.__init__( self )
		self.usernames = None

	def emit( self, sql ):
		# print sql
		self.db.query( sql )

	def init( self ):
		Pianos4Interface.init( self )
		self.nextID = self.db.query( "SELECT max(id)+1 AS nextid FROM object;" ).dictresult()[0]['nextid']
		print "-- Using next object ID", self.nextID
		self.usernames = {}
		for usr in self.db.query( "SELECT id,username FROM person;" ).dictresult():
			self.usernames[usr['username']] = usr['id']
		
	def close( self ):
		Pianos4Interface.close( self )
		print "--  Reindex"
		self.emit( "SET client_min_messages TO 'WARNING';" )
		self.emit( 'REINDEX DATABASE "' + self.params[2] + '";' )
		self.emit( "SET client_min_messages TO 'NOTICE';" )

	def getUsername( self, username ):
		if username in self.usernames:
			newname = '3_to_4_' + username 
			print "-- NOTE: user '%s' was renamed to '%s', the password must be changed" % ( username, newname )
			return newname
		else:
			return username
			
class Pianos4InterfaceDummy( Pianos4Interface ):		
	def emit( self, sql ):
		pass
		
class Tool:
	
	def __init__( self, source, destination ):
		self.source = source
		self.destination = destination
		self.destination.setSource( self.source )
	
	def	run( self ):
		self.destination.init()
		
		self.destination.copyUsers()
		self.destination.copyGroups()
		self.destination.copyImagesDef()
		self.destination.copyFolders()
		self.destination.copyExercises()
		self.destination.copyAssignments()
		self.destination.copyImagesData()
		
		self.destination.close()
		
	def copyUsers( self ):
		users = self.source.getUsers()
		
def main():
	
	parser = optparse.OptionParser( usage = 'SOURCE [ DESTINATION ]' )
	parser.set_description( 'Converts SOURCE, which is a PiAnoS 3.x database, to the new 4.x format. If DESTINATION is not provided, write the SQL script to stdout. If specified, DESTINATION must be an empty PiAnoS 4.x database (tables and the root user will not be created - you need to use the microDB schema creation tool first). Database arguments are connection URLs, and have the form pgsql://user:password@host:port/database. The pgsql:// part is added if not present. Only the database name is mandatory, user defaults to \'pianos\', password to \'pianos\', host to \'localhost\' and port to \'5432\'. The slash before the database name is however mandatory - therefore \'pgsql://pianos:pianos@localhost:5432/mydb\' can be written as \'/mydb\'.' )
	
	options, args = parser.parse_args()

	source = None
	destination = None
	
	i = 0
	for url in args:
		params = parseConnectionString( url )
		if i == 0:
			source = Pianos3Database( params )
		elif i == 1:
			destination = Pianos4Database( params )
		i += 1

	if not source:
		print 'ERROR: you must specify the source database connection string. Try --help'
		return 1

	# if not destination:
	# 	print 'ERROR: you must specify the destination database connection string. Try --help'
	# 	return 1

	print "-- Converting"
	print "--   " + source.params[5]
	
	if destination:
		print "--  to"
		print "--   " + destination.params[5]
	
	try:
		print "--"
		print "-- Connecting to the source database ..."
		source.connect()
		print "-- OK"
		
		if destination:
			print "-- Connecting to the destination database ..."
			destination.connect()
			print "-- OK"
		
		else:
			destination = Pianos4Interface()
			
		print ""
	
		tool = Tool( source, destination )
		tool.run()
		
		print ""
		print "-- All done !"
			
	except ( pg.InternalError, pg.ProgrammingError ), err:
		print err
	
	except Error, err:
		print "ERROR:", err

	return 0
	
if __name__ == '__main__':
	sys.exit( main() )
