'''AWS EC2 Wrappers'''

# General imports
import base64
import json
import os
import time
import re

# Libs imports
from botocore.client import Config
from botocore.exceptions import ClientError

# Module imports
from basepair.modules.aws.handler.exception import ExceptionHandler
from basepair.modules.aws.service import Service


class EC2(Service):
	'''Wrapper for EC2 services'''
	def __init__(self, cfg):
		super().__init__(cfg, 'EC2')
		self.client = self.session.client(**{
			'config': Config(retries={'max_attempts': 10, 'mode': 'standard'}),
			'service_name': 'ec2',
		})
		self.default = {
			'ami': cfg.get('ami'),
			'availability_zone': cfg.get('availability_zone'),
			'instance_type': cfg.get('instance_type', 'm1.large'),
			'key_name': cfg.get('key_name'),
			'region': cfg.get('region'),
			'role_name': cfg.get('role'),
			'security_groups': cfg.get('security_groups') or ['worker'],
		}
		self.resource = self.session.resource('ec2')

	def attach_role(self, instance_id, role_name):
		'''Attach a role to instance'''
		try:
			response = self.client.associate_iam_instance_profile(
				IamInstanceProfile={'Name': role_name},
				InstanceId=instance_id,
			)
		except ClientError as error:
			response = self.get_log_msg({
				'exception': error,
				'msg': f'Not able to attach role {role_name} to instance {instance_id}.',
			})
			if ExceptionHandler.is_throttled_error(exception=error):
				raise error
		return response

	def describe_images(self, owners=None, filters=None):
		'''Look for images'''
		if owners is None:
			owners = []
		if filters is None:
			filters = []
		try:
			response = self.client.describe_images(Owners=owners, Filters=filters)
		except ClientError as error:
			response = self.get_log_msg({
				'exception': error,
				'msg': f'Not able to describe images due to error :\n{error}',
			})
			if ExceptionHandler.is_throttled_error(exception=error):
				raise error
		return response

	def get_instances(self, filters=None, ids=[], instance_id=None): # pylint: disable=dangerous-default-value
		'''Look for instances'''
		if instance_id:
			return self.resource.Instance(instance_id)  # pylint: disable=no-member
		return self.resource.instances.filter(InstanceIds=ids, Filters=filters)  # pylint: disable=no-member

	def get_security_group_by_name(self, group_name):
		'''Get security group by its name'''
		try:
			response = self.client.describe_security_groups(
				GroupNames=[group_name]
			)
			if response and response.get('SecurityGroups'):
				response = response.get('SecurityGroups')[0]
		except ClientError as error:
			response = self.get_log_msg({
				'exception': error,
				'msg': f'Not able to get group {group_name}.',
			})
			if ExceptionHandler.is_throttled_error(exception=error):
				raise error
		return response

	def get_spot_price_history(self, options, use_cache=True):
		'''Look for the spot price history'''
		region = self.default.get('region')
		key = f"{options.get('availability_zone') or region}_{options.get('instance_type')}"
		cache_file = f'/tmp/{key}'
		ten_minutes_ago = time.time() - 600  # 600 seconds = 10 minutes
		if use_cache and os.path.isfile(cache_file):
			# reading the price from cache
			cache = {}
			with open(cache_file, 'r', encoding='utf-8') as file:
				cache = json.load(file)
			# expire the cache if needed
			filte_stats = os.stat(cache_file)
			if filte_stats.st_mtime < ten_minutes_ago:
				os.unlink(cache_file)
			# returning cached price
			return cache.get('price', [])

		try:
			response = self.client.describe_spot_price_history(
				AvailabilityZone=options.get('availability_zone', ''),
				EndTime=options.get('end_time'),
				InstanceTypes=[options.get('instance_type')],
				ProductDescriptions=options.get('product_description'),
				StartTime=options.get('start_time'),
			)
		except ClientError as error:
			response = self.get_log_msg({
				'exception': error,
				'msg': 'Not able to describe spot price history.',
			})
			if ExceptionHandler.is_throttled_error(exception=error):
				raise error

		price = response.get('SpotPriceHistory', [])
		if price and use_cache:
			# removing Timestamp so we can json it
			for item in price:
				item.pop('Timestamp')
			with open(cache_file, 'w', encoding='utf-8') as file:
				json_content = json.dumps({'price': price})
				file.write(json_content)
		return price

	def launch_instance(self, settings):  # pylint: disable=too-many-locals, too-many-branches
		'''Launch instance'''
		iam_role = settings.get('role_name') or self.default.get('role_name')
		security_groups = settings.get('security_groups') or self.default.get('security_groups')
		specification = {
			'BlockDeviceMappings': settings.get('block_device_map') or [],
			'ImageId': settings.get('ami', self.default.get('ami')),
			'InstanceType': settings.get('instance_type') or self.default.get('instance_type'),
			'SecurityGroupIds': self._get_security_groups_ids(security_groups),
			'UserData': base64.b64encode((settings.get('user_data') or '').encode()).decode('ascii')
		}
		if settings.get('ebs_optimized') is not None:
			specification['EbsOptimized'] = settings.get('ebs_optimized')
		is_spot = settings.get("spot")
		if not is_spot:
			specification['TagSpecifications'] = settings.get("tags_specifications") or []

		key_name = settings.get('key_name') or self.default.get('key_name')
		if key_name:
			specification['KeyName'] = key_name

		if iam_role:
			specification['IamInstanceProfile'] = {'Name': iam_role}

		availability_zone = settings.get('availability_zone') or self.default.get('availability_zone')
		if settings.get('subnet_id'):
			specification['SubnetId'] = settings.get('subnet_id')
		elif availability_zone:
			specification['Placement'] = {'AvailabilityZone': availability_zone}
			subnet_ids = settings.get('subnet_ids')
			if subnet_ids:
				for subnet_id in subnet_ids:
					subnet = self.resource.Subnet(subnet_id)
					if subnet.availability_zone == availability_zone:
						specification['SubnetId'] = subnet_id
						break
				else:
					return self.get_log_msg({
						'exception': None,
						'msg': 'No subnet found in the availability zone.',
					})

		price = settings.get('spot_price') or ''
		try:
			response = self.launch_spot(specification, price) if is_spot else self.launch_ondemand(specification)
			if not isinstance(response, dict):
				while (getattr(response, 'state', {}) or {}).get('Code') == 0:
					time.sleep(5)
					response.load()
				# update tags
				tags = []
				for tag, value in settings.get('tags', {}).items():
					tags.append({'Key': tag, 'Value': str(value)})
				response.create_tags(Tags=tags)

		except ClientError as error:
			response = self.get_log_msg({
				'exception': error,
				'msg': 'Not able to launch instance.',
			})
			if ExceptionHandler.is_throttled_error(exception=error) or ExceptionHandler.is_insufficient_capacity_error(exception=error):
				raise error
		return response

	def launch_ondemand(self, specification):
		'''Start an ondemand intance'''
		specification['MaxCount'] = 1
		specification['MinCount'] = 1
		response = self.resource.create_instances(# pylint: disable=no-member
			**specification,
		)
		return next(iter(response), None)

	def launch_spot(self, specification, price=''):
		'''Request a spot intance and wait till it is available'''
		response = self.client.request_spot_instances(
			LaunchSpecification=specification,
			SpotPrice=price,
		)
		try:
			request = response.get('SpotInstanceRequests', [])[0]
			response = self._wait_for_spot_instance_request(request.get('SpotInstanceRequestId'))
		except KeyError as error:
			response = self.get_log_msg({
				'exception': error,
				'msg': 'Not able to request spot instance.',
			})
		return response

	def look_for_instance(self, instance_id, retry=10):
		'''Look for the instance with retry'''
		instance = None
		# we break the loop after n retries
		while retry > 0:
			time.sleep(15)
			instance = self.get_instances(instance_id=instance_id)
			if instance and instance.instance_id:
				break
			retry -= 1

		if retry == 0:
			instance = self.get_log_msg({
				'exception': 'Looking for instance stop after a few retries.',
				'msg': f'Not able to find the instance {instance_id}.',
			})
		return instance

	def tag(self, instance, tags):
		'''Helper to add tags to instances'''
		response = {
			'detail': None,
			'error': True,
			'msg': 'Not valid instance to tag.',
		}
		if hasattr(instance, 'create_tags'):
			try:
				tags_to_push = []
				for tag, value in tags.items():
					tags_to_push.append({'Key': tag, 'Value': str(value)})
				return instance.create_tags(Tags=tags_to_push)
			except ClientError as error:
				response = self.get_log_msg({
					'exception': error,
					'msg': 'Not able to tag instance.',
				})
				if ExceptionHandler.is_throttled_error(exception=error):
					raise error
		return response

	def _get_security_groups_ids(self, security_groups):
		'''If security group is of the format sg-words, returns it as the
		security group id, otherwise, translates groups name to groups ids.
		The group name translation works only if the security group is in
		the default VPC'''
		security_group_ids = []
		security_group_id_pattern = re.compile(r'^sg-\w+$')
		for group in security_groups:
			if security_group_id_pattern.match(group):
				# group is a security group id
				security_group_ids.append(group)
			else:
				# group is probably a security group name.
				# this works only if it is in the default vpc.
				group_obj = self.get_security_group_by_name(group)
				if group_obj.get('GroupId'):
					security_group_ids.append(group_obj.get('GroupId'))
		return security_group_ids

	def _wait_for_spot_instance_request(self, request_id, retry=4):
		'''Wait till we get the spot instance or an error'''
		time.sleep(10)
		request_instance = None
		try:
			requests = self.client.describe_spot_instance_requests(
				SpotInstanceRequestIds=[request_id]
			)
			request_instance = (requests or {}).get('SpotInstanceRequests', [])[0]
		except (ClientError, KeyError) as error:
			self.get_log_msg({
				'exception': error,
				'msg': f'Not able to describe the spot instance request {request_id}.',
			})
			if ExceptionHandler.is_throttled_error(exception=error):
				raise error

		if request_instance:
			status = request_instance.get('Status', {})
			if status.get('Code') == 'fulfilled' and request_instance.get('InstanceId'):
				return self.look_for_instance(instance_id=request_instance.get('InstanceId'))

			error_status = [
				'canceled-before-fulfillment',
				'capacity-not-available',
				'capacity-oversubscribed',
				'failed',
				'instance-terminated-no-capacity',
				'price-too-low',
				'bad-parameters',
				'system-error'
			]
			if retry == 0 or status.get('Code') in error_status:
				response = {
					'detail': status.get('Code') or 'Describe spot instance request expired.',
					'error': True,
					'msg': 'Not able to create spot instance.',
				}
				try:
					self.client.cancel_spot_instance_requests(SpotInstanceRequestIds=[request_id])
				except ClientError as error:
					response = self.get_log_msg({
						'exception': error,
						'msg': 'Not able to cancel the spot instance request.',
					})
					if ExceptionHandler.is_throttled_error(exception=error):
						raise error
				return response
		time.sleep(10)
		return self._wait_for_spot_instance_request(request_id, retry=retry - 1)
