#!/usr/bin/env python

import commands, sys, os, traceback, time
import subprocess, time, tarfile, json, glob
import boto3, shutil

from common import write_pidfile_or_die, remove_pidfile, db_type, assume_role, s3_assume_role_client, sqs_assume_role_client

# Psuedo code
# Should check if there is already a datacollection running
# If no datacollection start datacolleciton, using the datacollection path
# Then should gzip csv files from datacollection
# Then upload to s3, send sqs to note that datacollection files are ready to be analyzed


def datacollection(db_name, times, tt_collect):
	datacollection_path = "/var/lib/postgresql/healthcheck/datacollection/datacollection.py"
	start_time = time.strftime("%Y_%m_%d/%H")
	success = os.system('python %s -d \"%s\" -c /etc/postgresql/9.3/main/postgresql.conf -%s %s --clear' %(datacollection_path, db_name, times, tt_collect))
	print success
	end_time = time.strftime("%Y_%m_%d/%H")
	times = {'start_time': start_time, 'end_time': end_time}
	if success == 0:
		print "returning times"
		return times
	else:
		print "datacollection error"
		clean_up_datacollection()
		exit (1)

def mv_data_collection_to_logs(s3_bucket, s3_key, count):
	print "Moving datacollection files to %s." % s3_key
	csv_Path="OUTPUT/project/logs/"
	datacollection_files=glob.glob('/var/log/postgresql/datacollection*')
	s3credentials=assume_role("assumePushcsvtoS3")
	s3_client=s3_assume_role_client(s3credentials)
	for dc_filepath in datacollection_files:
		s3_prefix=dc_filepath.split('/', 1)[1]
		print ("gzip %s" %dc_filepath)
		if dc_filepath.endswith('.csv'):
			print("uploading csv file to s3")
			print("csv file path : %s" %dc_filepath)
			os.system("nice gzip %s" % dc_filepath)
			print "s3 long path: " + s3_key+'/'+csv_Path+dc_filepath.rsplit('/', 1)[1]+'.gz'
			s3response=s3_client.upload_file(dc_filepath+'.gz', s3_bucket, s3_key+'/'+csv_Path+dc_filepath.rsplit('/', 1)[1]+'.gz')
			print ("status: %s" %s3response)
			if s3response:
				if count > 0:
					print "Retry, uploading csv file to s3"
					time.sleep(5)
					mv_data_collection_to_logs(s3_bucket, s3_key, count-1)


def uploadDirectory(s3_client, s3_prefix_path, path, bucketname):
	for dirs, subdir, files in os.walk(path):
		# print path to all filenames.
		for file in files:
			filename=os.path.join(dirs, file)
			try:
				s3_client.upload_file(filename, bucketname, s3_prefix_path+'/'+filename)	
			except Exception as e:
				return 0

#needs s3 paths and file nameing convention. Fix building tar ball then compress
#this uploads everything in the OUTPUT directory of the datacollection, csv files not included
def upload_to_S3(schema, fqdn, times, database, db_type, count):
	print "Uploading output directory to s3"
	date_stamp = times['start_time']
	s3_bucket="d8a/"
	s3_key="clio/"+database+"/"+schema+"/"+db_type+"/"+fqdn+"/"+date_stamp
	s3_path=s3_bucket+s3_key
	s3_filename="OUTPUT/"
	s3pushcredentials=assume_role("assumePushoutputtoS3")
	s3_client=s3_assume_role_client(s3pushcredentials)
	response=uploadDirectory(s3_client, s3_key, s3_filename, s3_bucket[:-1])
	if response != 0:
		print "s3_key: " +s3_key
		mv_data_collection_to_logs(s3_bucket[:-1], s3_key, 5)
		send_sqs(fqdn, s3_path, s3_filename, schema, times, database, db_type)
	else:
		print "Retry, uploading output folder to s3"
		if count > 0:
			time.sleep(5)
			upload_to_S3(schema, fqdn, times, database, db_type, count-1)
		clean_up_datacollection()
		exit (1)

def send_sqs(fqdn, s3_path, s3_filename, schema, times, database, db_type):
	queue_url='https://sqs.us-west-2.amazonaws.com/673385534282/Clio_Pull_Request'
	sqscredentials=assume_role("assumeSQS")
	sqs_client=sqs_assume_role_client(sqscredentials)
	response = sqs_client.send_message(
	    QueueUrl=queue_url,
	    MessageBody='ready to pull',
	    MessageAttributes={
	        'date_stamp': {
	            'StringValue': times['start_time'],
	            'DataType': 'String'
	        },
	        'end_time': {
	            'StringValue': times['end_time'],
	            'DataType': 'String'
	        },
	        'fqdn': {
	            'StringValue': fqdn,
	            'DataType': 'String'
	        },
	        's3_path': {
	            'StringValue': s3_path,
	            'DataType': 'String'
	        },
	        's3filename': {
	            'StringValue': s3_filename,
	            'DataType': 'String'
	        },
	        'schema': {
	            'StringValue': schema,
	            'DataType': 'String'
	        },
	        'database': {
	            'StringValue': database,
	            'DataType': 'String'
	        },
	        'db_type': {
	            'StringValue': db_type,
	            'DataType': 'String'
	        },
	    }
	)
	if response:
		print "sqs message successful"
		clean_up_datacollection()
	else:
		print "Could not send sqs message"
		clean_up_datacollection()
		exit(1)

def clean_up_datacollection():
	os.system('rm -rf /var/log/postgresql/datacollection*')
	os.system('rm -rf OUTPUT/')
	remove_pidfile("/var/lib/postgresql/healthcheck/datacollection.pid")


if __name__ == '__main__':

	try:
		prog = "datacollection.sh"
		
		db_type = db_type()

		if len(sys.argv) > 4:
			write_pidfile_or_die("/var/lib/postgresql/healthcheck/datacollection.pid")
			times = datacollection(sys.argv[1], sys.argv[2],sys.argv[3])
			upload_to_S3(sys.argv[1], sys.argv[4], times, sys.argv[5], db_type, 5)  #argv[1] = schema, argv[4] = fqdn, argv[5] = database e.g. sitedb,usherdb, argv[6] = db_type
		else:
			print("Error: Not enough parameters")
			exit(1)

	except Exception as e:
		clean_up_datacollection()
		traceback.print_exc()
		exit(1)

