#!/usr/bin/env python

import subprocess, sys, os, argparse, boto3
from objects.exceptions import CommandError

def get_aws_session(profile_name=None):
	return boto3.Session(profile_name=profile_name, region_name='eu-west-1')

def get_instances_list(args):

	conn_eb = get_aws_session(args.profile).client('elasticbeanstalk')

	instances_list = [ instance['Id'] for instance in conn_eb.describe_environment_resources(EnvironmentName=args.environment_name)['EnvironmentResources']['Instances'] ]

	return instances_list

def get_input(output, default):

	result = str(raw_input(output + ': ')).strip() or default
	return result


def prompt_for_instance_in_list(instances_list, default=1):
	for x in range(0, len(instances_list)):
		print str(x + 1) + ')', instances_list[x]

	while True:
		try:
			choice = int(get_input('(default is ' + str(default)+')', default))
			if not (0 < choice <= len(instances_list)):
				raise ValueError  # Also thrown by non int numbers
			else:
				break
		except ValueError:
			print 'Sorry, that is not a valid choice. Please choose a number between 1 and ' + str(len(instances_list)) + '.'

	return choice - 1

def ssh_into_instance(args, instance_id, custom_ssh=None, command=None):

	conn_ec2 = get_aws_session(args.profile).client('ec2')

	instance = conn_ec2.describe_instances(InstanceIds=[instance_id])['Reservations'][0]['Instances'][0]

	try:
		keypair_file = '~/.ssh/'+instance['KeyName']
	except KeyError:
		raise
	try:
		ip = instance['PrivateIpAddress']
	except KeyError:
		# Now allows access to private subnet
		if 'PrivateIpAddress' in instance and 'PrivateDnsName' in instance:
			ip = instance['PrivateDnsName']
		else:
			raise

	user = 'ec2-user'

	# do ssh
	try:
		ssh_command = ['ssh', '-i', keypair_file, user + '@' + ip]

		print 'INFO: Running ' + ' '.join(ssh_command)
		returncode = subprocess.call(ssh_command)
		if returncode != 0:
			raise CommandError('An error occurred while running: ' + ssh_command[0] + '.')
	except OSError:
		raise

def main(args):

	# Get the instances list
	instances_list = get_instances_list(args)

	# Prompt the user to chose an instance
	instance_chosen = prompt_for_instance_in_list(instances_list)

	# Open the ssh connection using the internal IP
	ssh_into_instance(args, instances_list[instance_chosen])


### Init ###

if 'ssh' in [arguments for arguments in sys.argv]:

	parser = argparse.ArgumentParser()

	parser.add_argument('ssh')
	parser.add_argument('--profile', help='--profile <application-[production|staging]>')
	parser.add_argument('environment_name', help='<application-[production|staging]-[web|activejobs|cronjobs]>')

	args = parser.parse_args()

	main(args)

else:
	print 'Forwarding command to the official awsebcli..'

	aws_eb_cli = ['eb']

	for eb_arguments in sys.argv[1:]:
		aws_eb_cli.append(eb_arguments)

	subprocess.call(aws_eb_cli)




