import cv2
import time
import socket
import threading


class Response(object):
    def __init__(self):
        pass

    def recv(self, data):
        pass

    def pop(self):
        pass

    def empty(self):
        pass


class Command(Response):
    def __init__(self):
        super(Command, self).__init__()
        self.response = None
        self.lock = threading.RLock()

    def recv(self, data):
        with self.lock:
            self.response = data.decode('utf-8')

    def pop(self):
        with self.lock:
            response, self.response = self.response, None
        return response

    def empty(self):
        with self.lock:
            return self.response is None


class State(Response):
    def __init__(self):
        super(State, self).__init__()
        self.response = {}
        self.lock = threading.RLock()

    def recv(self, data):
        with self.lock:
            self.response = {item.split(':')[0]:float(item.split(':')[1]) for item in data.decode('utf-8').split(';') if ':' in item}

    def pop(self):
        return self.response

    def empty(self):
        return False


class Client(object):
    def __init__(self, local_port, buffer_size, daemon, response):
        self.response = response
        self.buffer_size = buffer_size

        self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.socket.bind(('', local_port))

        self.receive_thread = threading.Thread(target=self._receive_thread)
        self.receive_thread.daemon = daemon
        self.receive_thread.start()

    def __del__(self):
        """Closes the local socket."""
        self.socket.close()

    def _receive_thread(self):
        """Listens for responses from the Tello.

        Runs as a thread, sets self.response to whatever the Tello last returned.
        """
        while True:
            try:
                self.response.recv(self.socket.recv(self.buffer_size))
            except Exception as e:
                print(e)
                break

    def empty(self):
        return self.response.empty()

    def pop(self):
        return self.response.pop()


class Video(object):
    def __init__(self, daemon=True):
        self.video = cv2.VideoCapture('udp://@0.0.0.0:11111')
        if not self.video.isOpened():
            raise RuntimeError('Failed to connect to Tello')
        self.frame = None
        self.lock = threading.RLock()
        self.thread = threading.Thread(target=self._update_thread)
        self.thread.daemon = daemon
        self.thread.start()

    def __del__(self):
        self.video.release()

    def _update_thread(self):
        while True:
            ok, frame = self.video.read()
            if ok:
                with self.lock:
                    self.frame = frame

    def empty(self):
        with self.lock:
            return self.frame is None

    def pop(self):
        with self.lock:
            frame, self.frame = self.frame, None
        return frame


class Tello(object):
    def __init__(self, local_port=9999, command_timeout=0.35, state=True, video=True):
        """Connects to Tello in command mode.

        Args:
            local_port (int): port of local machine for receiving command response.
            command_timeout (float): seconds to wait for a response of command.
            state (bool): receive state from Tello?
            video (bool): receive video from Tello?

        Raises:
            RuntimeError: If the Tello rejects the attempt to enter command mode or open the video stream.
        """
        self.command_timeout = command_timeout

        self.response_client = Client(local_port, 1024, True, Command())
        self.state_client = Client(8890, 1024, True, State()) if state else None

        self.tello_address = ('192.168.10.1', 8889)
        self.enter_command_mode()
        self.video_client = None
        if video:
            self.open_video_stream()
            self.video_client = Video(True)

    def send_command(self, command, with_return=True):
        """Sends a command to the Tello and waits for a response.

        If self.command_timeout is exceeded before a response is received,
        a RuntimeError exception is raised.

        Args:
            command (str): Command to send.

        Returns:
            str: Response from Tello.

        Raises:
            RuntimeError: If no response is received within self.timeout seconds.
        """
        self.response_client.pop()
        self.response_client.socket.sendto(command.encode('utf-8'), self.tello_address)

        if not with_return:
            return

        st = time.time()
        while self.response_client.empty():
            if time.time() - st >= self.command_timeout:
                raise RuntimeError('No response to command')
        return self.response_client.pop()

    def state(self):
        return self.state_client.pop() if self.state_client else None

    def read_frame(self):
        if self.video_client is None:
            raise RuntimeError('Video is not available')
        while self.video_client.empty():
            pass
        return self.video_client.pop()

    def enter_command_mode(self):
        if self.send_command('command') != 'ok':
            raise RuntimeError('Tello rejected the attempt to enter command mode')

    def take_off(self):
        """
        return: 'ok' or 'error'
        """
        return self.send_command('takeoff')

    def land(self):
        """
        return: 'ok' or 'error'
        """
        return self.send_command('land')

    def open_video_stream(self):
        if self.send_command('streamon') != 'ok':
            raise RuntimeError('Tello rejected to open the video stream')

    def close_video_stream(self):
        """
        return: 'ok' or 'error'
        """
        return self.send_command('streamoff')

    def emergency_shutdown(self):
        """
        return: 'ok' or 'error'
        """
        return self.send_command('emergency')

    def move_up(self, x, with_return=False):
        """
        param x: int, [20, 500]
        param with_return: bool
        return: 'ok' or 'error'
        """
        return self.send_command('up {}'.format(x), with_return)

    def move_down(self, x, with_return=False):
        """
        param x: int, [20, 500]
        param with_return: bool
        return: 'ok' or 'error'
        """
        return self.send_command('down {}'.format(x), with_return)

    def move_left(self, x, with_return=False):
        """
        param x: int, [20, 500]
        param with_return: bool
        return: 'ok' or 'error'
        """
        return self.send_command('left {}'.format(x), with_return)

    def move_right(self, x, with_return=False):
        """
        param x: int, [20, 500]
        param with_return: bool
        return: 'ok' or 'error'
        """
        return self.send_command('right {}'.format(x), with_return)

    def move_forward(self, x, with_return=False):
        """
        param x: int, [20, 500]
        param with_return: bool
        return: 'ok' or 'error'
        """
        return self.send_command('forward {}'.format(x), with_return)

    def move_backward(self, x, with_return=False):
        """
        param x: int, [20, 500]
        param with_return: bool
        return: 'ok' or 'error'
        """
        return self.send_command('back {}'.format(x), with_return)

    def rotate_clockwise(self, x, with_return=False):
        """
        param x: int, [1, 3600]
        param with_return: bool
        return: 'ok' or 'error'
        """
        return self.send_command('cw {}'.format(x), with_return)

    def rotate_counter_clockwise(self, x, with_return=False):
        """
        param x: int, [1, 3600]
        param with_return: bool
        return: 'ok' or 'error'
        """
        return self.send_command('ccw {}'.format(x), with_return)

    def flip_left(self, with_return=False):
        """
        param with_return: bool
        return: 'ok' or 'error'
        """
        return self.send_command('flip l', with_return)

    def flip_right(self, with_return=False):
        """
        param with_return: bool
        return: 'ok' or 'error'
        """
        return self.send_command('flip r', with_return)

    def flip_forward(self, with_return=False):
        """
        param with_return: bool
        return: 'ok' or 'error'
        """
        return self.send_command('flip f', with_return)

    def flip_backward(self, with_return=False):
        """
        param with_return: bool
        return: 'ok' or 'error'
        """
        return self.send_command('flip b', with_return)

    def goto(self, x, y, z, speed, with_return=False):
        """
        param x: int, [20, 500]
        param y: int, [20, 500]
        param z: int, [20, 500]
        param speed: int, [10-100]
        param with_return: bool
        return: 'ok' or 'error'
        """
        return self.send_command('go {} {} {} {}'.format(x, y, z, speed), with_return)

    def goto_curve(self, x1, y1, z1, x2, y2, z2, speed, with_return=False):
        """fly a curve defined by (0, 0, 0), (x1, y1, z1), (x2, y2, z2) with speed
        param x1, x2: int, [-500, 500]
        param y1, y2: int, [-500, 500]
        param z1, z2: int, [-500, 500]
        param speed: int, [10-60]
        param with_return: bool
        return: 'ok' or 'error'
        """
        return self.send_command('curve {} {} {} {} {} {} {}'.format(x1, y1, z1, x2, y2, z2, speed), with_return)

    def set_speed(self, speed, with_return=False):
        """
        param speed: int, [10-100]
        param with_return: bool
        return: 'ok' or 'error'
        """
        return self.send_command('speed {}'.format(speed), with_return)

    def set_remote_controller_command(self, left_right_velocity, forward_backward_velocity, up_down_velocity, rotate_velocity, with_return=False):
        """
        param left_right_velocity: int, [-100, 100]
        param forward_backward_velocity: int, [-100, 100]
        param up_down_velocity: int, [-100, 100]
        param rotate_velocity: int, [-100, 100]
        param with_return: bool
        return: 'ok' or 'error'
        """
        return self.send_command('rc {} {} {} {}'.format(left_right_velocity, forward_backward_velocity, up_down_velocity, rotate_velocity), with_return)

    def get(self, command, split=False):
        """
        param command
        param split: bool, multiple values?
        return: int or list(int)
        """
        result = self.send_command(command)
        if split:
            return [int(x) for x in result.split(' ')]
        else:
            return int(result)

    def get_speed(self):
        """
        return: int, [10, 100]
        """
        return self.get('speed?')

    def get_battery(self):
        """
        return: int, [0, 100]
        """
        return self.get('battery?')

    def get_flight_time(self):
        """
        return: int
        """
        return self.get('time?')

    def get_relative_height(self):
        """
        return: int, [10, 3000]
        """
        return self.get('height?')

    def get_temperature(self):
        """
        return: int, [0, 90]
        """
        return self.get('temp?')

    def get_imu_pose(self):
        """[pitch, roll, yaw]
        return: list(int), [[-89, 89], [-179, 179], [-179, 179]]
        """
        return self.get('attitude?', split=True)

    def get_absolute_height(self):
        """
        return: int
        """
        return self.get('baro?')

    def get_imu_acceleration(self):
        """
        return: list(int)
        """
        return self.get('acceleration?', split=True)

    def get_tof_height(self):
        """
        return: int, [10, 400]; 6553: out of bounds
        """
        return self.get('tof?')

[((25, 20, 25, 37), 'threading.RLock', 'threading.RLock', ({}, {}), '()', False, 'import threading\n'), ((45, 20, 45, 37), 'threading.RLock', 'threading.RLock', ({}, {}), '()', False, 'import threading\n'), ((63, 22, 63, 70), 'socket.socket', 'socket.socket', ({(63, 36, 63, 50): 'socket.AF_INET', (63, 52, 63, 69): 'socket.SOCK_DGRAM'}, {}), '(socket.AF_INET, socket.SOCK_DGRAM)', False, 'import socket\n'), ((66, 30, 66, 75), 'threading.Thread', 'threading.Thread', (), '', False, 'import threading\n'), ((95, 21, 95, 61), 'cv2.VideoCapture', 'cv2.VideoCapture', ({(95, 38, 95, 60): '"""udp://@0.0.0.0:11111"""'}, {}), "('udp://@0.0.0.0:11111')", False, 'import cv2\n'), ((99, 20, 99, 37), 'threading.RLock', 'threading.RLock', ({}, {}), '()', False, 'import threading\n'), ((100, 22, 100, 66), 'threading.Thread', 'threading.Thread', (), '', False, 'import threading\n'), ((170, 13, 170, 24), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((172, 15, 172, 26), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n')]