#DronePool module which handles interaction with SITLs

from dronekit import Vehicle, VehicleMode, connect
from dronekit_sitl import SITL
from threading import Lock
import node, time
import mavparser
import threadrunner

drone_pool = {}
instance_count = 0
env_test = False

q = None
mq = None
lock = Lock()

class Sim(SITL, object):
	def __init__(self, instance=1, home=None):
		super(Sim, self).download("copter", "3.3", verbose=not env_test)
		self.instance = instance

		if home:
			self.home = home
		else:
			self.home = {"lat":6.9271, "lon":79.8612, "alt": 1}

		self.p = None
		return

	def connection_string(self):
		return super(Sim, self).connection_string()[:-4] + str(5760 + self.instance * 10)

	def launch(self):
		home_str = str(self.home['lat']) + ',' + str(self.home['lon']) + ',0,353'
		super(Sim, self).launch(["--instance", str(self.instance), "--home", home_str], await_ready=True, verbose=not env_test)

	def get_sitl_status(self):
		return { 'id': self.instance, 'home': self.home }

def initialize():
	global q, mq, instance_count
	q = threadrunner.q
	mq = threadrunner.mq

	drones = node.get_drones()['drones']

	if not drones:
		return

	for drone_id in drones:
		if drone_id not in list(drone_pool.keys()):
			drone = node.get_drone_by_id(drone_id)
			location = drone['location']
			q.put((create_new_drone, { "db_key" : drone_id, "home" : location }))

			if 'status' in list(drone.keys()) and drone['status'] == 'FLYING':
				q.put((resume_flight, { "drone_id" : drone_id }))

def resume_flight(kwargs):
	drone_id = kwargs.get("drone_id", None)
	drone = node.get_drone_by_id(drone_id)

	waypoints = []

	for wp in sorted(drone['waypoints']):
		waypoints.append(drone['waypoints'][wp])

	next_waypoint = waypoints.index(drone['waypoint'])
	print (next_waypoint)

	q.put((takeoff_drone, { "drone_id" : drone_id, "waypoints" : waypoints[next_waypoint:] }))

def create_new_drone(kwargs):
	global instance_count
	instance_count += 1
	home = kwargs.get("home", None)
	db_key = kwargs.get("db_key", None)

	retries = 3

	drone = Sim(instance_count, home)
	drone.launch()

	while retries > 0:
		try:
			drone_conn = connect(drone.connection_string(), wait_ready=True)
			break
		except:
			print ("Retrying...")
			retries -= 1


	drone_pool[db_key] = drone_conn

	res = { "status" : "OK", "id" : db_key }
	return res

def remove_drone(kwargs):
	drone_id = kwargs.get("drone_id", None)
	if drone_id not in drone_pool:
		return { "status" : "ERROR", "msg" : "Drone instance not found" }

	drone = drone_pool[drone_id]

	if drone.mode == VehicleMode('AUTO'):
		return { "status" : "ERROR", "msg" : "Drone in operation" }

	del drone_pool[drone_id]

	return { "status" : "OK", "id" : drone_id }


def run_mission(drone, target_height, waypoints):
	while True:
		print(("Reaching target alt : " + str(drone.location.global_relative_frame.alt)))
		if drone.location.global_relative_frame.alt >= target_height * 0.9:
			break

	print ('target alt reached')

	mavparser.create_mission(drone, waypoints)
	print ('mission acquired')

	drone.mode = VehicleMode('AUTO')
	print ('initiating sequence')

	print ('in mission')

def attach_listener(kwargs):
	attr = kwargs.get('attr', None)
	fn = kwargs.get('fn', None)
	attach_fn = kwargs.get('attach_fn', None)

	if not fn == None and not attr == None and not attach_fn == None:
		attach_fn(attr, fn)

def takeoff_drone(kwargs):
	global q

	drone_id = kwargs.get("drone_id", None)
	target_height = kwargs.get("target_height", 10)
	waypoints = kwargs.get("waypoints", None)

	try:
		drone = drone_pool[drone_id]
	except:
		raise

	drone.initialize()

	drone.mode = VehicleMode('GUIDED')
	drone.armed = True

	while not drone.armed:
		time.sleep(1)

	drone.simple_takeoff(target_height)

	print (waypoints)

	if waypoints:
		run_mission(drone, target_height, waypoints)

	def detach_event_listeners(drone, value, status):
		drone.remove_attribute_listener('location', update_location)
		drone.remove_attribute_listener('airspeed', update_airspeed)
		drone.remove_attribute_listener('attitude', udpate_attitude)
		drone.remove_attribute_listener('heading', update_heading)
		node.update_drone(drone_id, { "location" : {"lat": value.global_relative_frame.lat, "lon": value.global_relative_frame.lon, "alt": value.global_relative_frame.alt}, "status": status})
		return

	def update_location(self, attr_name, value):

		node.update_drone(drone_id, { "location" : {"lat": value.global_relative_frame.lat, "lon": value.global_relative_frame.lon, "alt": value.global_relative_frame.alt}, "status": "FLYING"})

		command_len = len(drone.commands)
		wp_len = len(waypoints)

		if command_len >= wp_len :
			diff = command_len - wp_len
			next_wp = max(drone.commands.__next__ - diff, 0) % len(waypoints)
			waypoint = waypoints[next_wp]
			# print "df: " + `diff`
			# print next_wp
			node.update_drone(drone_id, { "waypoint" : waypoint })

		if drone.mode == VehicleMode('LAND') and drone.location.global_relative_frame.alt <= 0.1:
			detach_event_listeners(drone, value, "HALTED")
			return

		if drone.commands.__next__ == len(drone.commands):
			detach_event_listeners(drone, value, "FINISHED")
			return


	def update_airspeed(self, attr_name, value):
		node.update_drone(drone_id, {"airspeed": value})

	def udpate_attitude(self, attr_name, value):
		node.update_drone(drone_id, { "pitch": value.pitch, 'roll': value.roll, 'yaw': value.yaw })

	def update_heading(self, attr_name, value):
		node.update_drone(drone_id, { "heading": value })

	mq.put((attach_listener, { "attach_fn" : drone.add_attribute_listener, "attr" : 'location', "fn" : update_location }))
	mq.put((attach_listener, { "attach_fn" : drone.add_attribute_listener, "attr" : 'airspeed', "fn" : update_airspeed }))
	mq.put((attach_listener, { "attach_fn" : drone.add_attribute_listener, "attr" : 'attitude', "fn" : udpate_attitude }))
	mq.put((attach_listener, { "attach_fn" : drone.add_attribute_listener, "attr" : 'heading', "fn" : update_heading }))

	print ('took off')

	return True

def land_drone(kwargs):
	drone_id = kwargs.get("drone_id", None)
	try:
		drone = drone_pool[drone_id]
	except:
		raise

	if not drone.armed:
		return False

	cmds = drone.commands
	cmds.wait_ready()
	cmds.clear()

	drone.mode = VehicleMode('LAND')
	print((drone.mode))
	return True

[((16, 7, 16, 13), 'threading.Lock', 'Lock', ({}, {}), '()', False, 'from threading import Lock\n'), ((62, 9, 62, 39), 'node.get_drone_by_id', 'node.get_drone_by_id', ({(62, 30, 62, 38): 'drone_id'}, {}), '(drone_id)', False, 'import node, time\n'), ((122, 1, 122, 43), 'mavparser.create_mission', 'mavparser.create_mission', ({(122, 26, 122, 31): 'drone', (122, 33, 122, 42): 'waypoints'}, {}), '(drone, waypoints)', False, 'import mavparser\n'), ((125, 14, 125, 33), 'dronekit.VehicleMode', 'VehicleMode', ({(125, 26, 125, 32): '"""AUTO"""'}, {}), "('AUTO')", False, 'from dronekit import Vehicle, VehicleMode, connect\n'), ((152, 14, 152, 35), 'dronekit.VehicleMode', 'VehicleMode', ({(152, 26, 152, 34): '"""GUIDED"""'}, {}), "('GUIDED')", False, 'from dronekit import Vehicle, VehicleMode, connect\n'), ((229, 14, 229, 33), 'dronekit.VehicleMode', 'VehicleMode', ({(229, 26, 229, 32): '"""LAND"""'}, {}), "('LAND')", False, 'from dronekit import Vehicle, VehicleMode, connect\n'), ((46, 10, 46, 27), 'node.get_drones', 'node.get_drones', ({}, {}), '()', False, 'import node, time\n'), ((106, 18, 106, 37), 'dronekit.VehicleMode', 'VehicleMode', ({(106, 30, 106, 36): '"""AUTO"""'}, {}), "('AUTO')", False, 'from dronekit import Vehicle, VehicleMode, connect\n'), ((156, 2, 156, 15), 'time.sleep', 'time.sleep', ({(156, 13, 156, 14): '(1)'}, {}), '(1)', False, 'import node, time\n'), ((170, 2, 170, 185), 'node.update_drone', 'node.update_drone', ({(170, 20, 170, 28): 'drone_id', (170, 30, 170, 184): "{'location': {'lat': value.global_relative_frame.lat, 'lon': value.\n    global_relative_frame.lon, 'alt': value.global_relative_frame.alt},\n    'status': status}"}, {}), "(drone_id, {'location': {'lat': value.\n    global_relative_frame.lat, 'lon': value.global_relative_frame.lon,\n    'alt': value.global_relative_frame.alt}, 'status': status})", False, 'import node, time\n'), ((175, 2, 175, 187), 'node.update_drone', 'node.update_drone', ({(175, 20, 175, 28): 'drone_id', (175, 30, 175, 186): "{'location': {'lat': value.global_relative_frame.lat, 'lon': value.\n    global_relative_frame.lon, 'alt': value.global_relative_frame.alt},\n    'status': 'FLYING'}"}, {}), "(drone_id, {'location': {'lat': value.\n    global_relative_frame.lat, 'lon': value.global_relative_frame.lon,\n    'alt': value.global_relative_frame.alt}, 'status': 'FLYING'})", False, 'import node, time\n'), ((198, 2, 198, 50), 'node.update_drone', 'node.update_drone', ({(198, 20, 198, 28): 'drone_id', (198, 30, 198, 49): "{'airspeed': value}"}, {}), "(drone_id, {'airspeed': value})", False, 'import node, time\n'), ((201, 2, 201, 93), 'node.update_drone', 'node.update_drone', ({(201, 20, 201, 28): 'drone_id', (201, 30, 201, 92): "{'pitch': value.pitch, 'roll': value.roll, 'yaw': value.yaw}"}, {}), "(drone_id, {'pitch': value.pitch, 'roll': value.roll,\n    'yaw': value.yaw})", False, 'import node, time\n'), ((204, 2, 204, 51), 'node.update_drone', 'node.update_drone', ({(204, 20, 204, 28): 'drone_id', (204, 30, 204, 50): "{'heading': value}"}, {}), "(drone_id, {'heading': value})", False, 'import node, time\n'), ((53, 11, 53, 41), 'node.get_drone_by_id', 'node.get_drone_by_id', ({(53, 32, 53, 40): 'drone_id'}, {}), '(drone_id)', False, 'import node, time\n'), ((186, 3, 186, 57), 'node.update_drone', 'node.update_drone', ({(186, 21, 186, 29): 'drone_id', (186, 31, 186, 56): "{'waypoint': waypoint}"}, {}), "(drone_id, {'waypoint': waypoint})", False, 'import node, time\n'), ((188, 19, 188, 38), 'dronekit.VehicleMode', 'VehicleMode', ({(188, 31, 188, 37): '"""LAND"""'}, {}), "('LAND')", False, 'from dronekit import Vehicle, VehicleMode, connect\n')]