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

import sys
import optparse
import os
import pg
import types
import urllib
import urlparse
import math
import md5
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] )
	
	def query( self, sql ):
		res = self.db.query( sql )
		try:
			return res.dictresult()
		except AttributeError:
			return res
		
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 type( a ) is types.StringType:
				t = 'string'
			elif type( a ) is types.BooleanType:
				t = 'bool'
			else:
				t = 'verbatim'
			
		if 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 )

		
def main():
	
	parser = optparse.OptionParser( usage = 'DATABASE PWD_FILE [PREFIX=user] [NUM_USERS_PER_GROUP=100] [GROUP_PREFIX=group] [START_AT=1]' )
	parser.set_description( '''Creates users named \'userXXX\' (or starting with the specified prefix) with random passwords, putting them in groups of 100 (or the numner specified). The second argument is a file containing passwords, 1 by line. As many users as entries in the files are created.
	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()

	db = None
	inputFile = None
	userPrefix = 'user'
	numUsersPerGroup = 100
	groupPrefix = 'group'
	startAt = 1
	
	i = 0
	for arg in args:
		if i == 0:
			db = Database( parseConnectionString( arg ) )
		elif i == 1:
			inputFile = arg
		elif i == 2:
			userPrefix = arg
		elif i == 3:
			numUsersPerGroup = int( arg )
		elif i == 4:
			groupPrefix = arg
		elif i == 5:
			startAt = int( arg )
		i += 1

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

	if not inputFile:
		print 'ERROR: you must specify the input password file. Try --help'
		return 1

	
	try:
		print "Connecting to the database", db.params[5]
		db.connect()
		print "OK"
		print ""
		
		passwords = []
		print "Reading the password file '%s'" % inputFile
		f = file( inputFile, 'r' )
		for line in f:
			passwords.append( line.strip() )
		print "  OK,", len( passwords ), "entries found"
		print ""
		
		numUsers = len( passwords )
		
		print "Creating", numUsers, "users named", userPrefix + 'XXX', ",", numUsersPerGroup, "users per group named", groupPrefix + 'YYY'
		print ""
		
		db.query( "BEGIN;" )
		
		db.query( "SELECT set_config('unil.user',(SELECT person.id FROM person,membership,pgroup WHERE person.id=membership.person AND membership.pgroup = pgroup.id AND pgroup.tutors IS true ORDER BY id ASC LIMIT 1)::text,false);" )
		
		numGroups = int( math.ceil( float( numUsers ) / float( numUsersPerGroup ) ) )
		print "Creating", numGroups, "groups"

		groups = {}
		
		for i in range( numGroups ):
			gname = groupPrefix + str( i + 1 )
			print "Creating group", gname
			gid = db.query( "INSERT INTO pgroup(name) VALUES (" + sqlarg( [gname] ) + ") RETURNING id;" )[0]['id']
			groups[i] = {'id':gid, 'name':gname}
			print "  OK, id=" + str( gid )
		
		usersCreated = []
		for i in range( numUsers ):
			username = '%s%03d' % ( userPrefix, i + startAt )
			gindex = i / numUsersPerGroup
			password = passwords[i]
			print "Creating user", username, "in group", groups[gindex]['name'], ",password=", password
			m = md5.new()
			m.update( username )
			m.update( ':' )
			m.update( password )
			pwd = m.hexdigest()
			uid = db.query( "INSERT INTO person(username,name,password,active,language) VALUES (" + sqlarg( [username, username, pwd, sqlverbatim( 'true' ), sqlverbatim( '(SELECT id FROM language ORDER BY id ASC LIMIT 1)' )] ) + ") RETURNING id;" )[0]['id']
			db.query( "INSERT INTO membership(person,pgroup) VALUES (" + sqlarg( [uid, groups[gindex]['id']] ) + ");" )
			usersCreated.append( [username, password] )
			
		db.query( "COMMIT;" )
		
		print ""
		print "Now dumping users and passwords"
		for usr in usersCreated:
			print usr[0], usr[1]
			
		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() )
