import sys
import socket
import threading
import time

from mpc_controller import MPCController
from robot import Robot
from aruco_estimator import ArucoEstimator

MSGLEN = 64

def myreceive(sock):
    chunks = []
    bytes_recd = 0
    while bytes_recd < MSGLEN:
        chunk = sock.recv(1)
        if chunk == b'':
            raise RuntimeError("socket connection broken")
        chunks.append(chunk)
        bytes_recd = bytes_recd + len(chunk)

        if chunks[-1] == b'\n':
            break

    return b''.join(chunks)


class RoboRallyRobot(Robot):
    # dictionary mapping the current orientation and a turn command to the resulting orientation
    resulting_orientation = {
        '^': {'turn left': '<', 'turn right': '>', 'turn around': 'v'},
        '>': {'turn left': '^', 'turn right': 'v', 'turn around': '<'},
        'v': {'turn left': '>', 'turn right': '<', 'turn around': '^'},
        '<': {'turn left': 'v', 'turn right': '^', 'turn around': '>'},
    }

    # dictionary mapping an orientation to its opposite
    opposites = {'^': 'v', '>': '<', 'v': '^', '<': '>'}

    def __init__(self, id, ip, x, y, orientation):
        super().__init__(id, ip)

        self.grid_x = x
        self.grid_y = y
        self.grid_orientation = orientation

    def get_neighbor_coordinates(self, direction):
        # get the coordinates of the neighboring tile in the given direction
        if direction == '^':
            return self.grid_x, self.grid_y - 1
        elif direction == '>':
            return self.grid_x + 1, self.grid_y
        elif direction == 'v':
            return self.grid_x, self.grid_y + 1
        elif direction == '<':
            return self.grid_x - 1, self.grid_y
        else:
            print("error: unknown direction")
            sys.exit(1)

    def move(self, move_type):
        if move_type == 'forward':
            target_tile = self.get_neighbor_coordinates(self.grid_orientation)
            self.grid_x = target_tile[0]
            self.grid_y = target_tile[1]
        elif move_type == 'backward':
            opposite_orientation = RoboRallyRobot.opposites[self.grid_orientation]
            target_tile = self.get_neighbor_coordinates(opposite_orientation)
            self.grid_x = target_tile[0]
            self.grid_y = target_tile[1]
        elif 'turn' in move_type:
            self.grid_orientation = RoboRallyRobot.resulting_orientation[self.grid_orientation][move_type]
        elif 'nop' in move_type:
            pass    # nop command -> robot grid position does not change (used e.g. for driving the robot to initial
                    #                                                      position)
        else:
            print(f"error: invalid move: {move_type}")
            sys.exit(1)

    def __str__(self):
        return self.__repr__()

    def __repr__(self):
        return f"grid x: {self.grid_x}, grid y: {self.grid_y}, grid orientation: {self.grid_orientation}"


class RemoteController:
    valid_cmds = ['forward', 'backward', 'turn left', 'turn right', 'turn around', 'nop', 'get position',
                  'set position', 'initialize_robot', 'initialize_grid']

    def __init__(self):
        self.robots = []
        #self.robots = [RoboRallyRobot(12, '192.168.1.12', x=1, y=1, orientation='>')]

        self.robot_ids = {}
        for r in self.robots:
            self.robot_ids[r.id] = r

        # socket for movement commands
        self.comm_socket = socket.socket()
        self.comm_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

        for robot in self.robots:
            robot.connect()

        self.comm_socket.bind(('', 1337))
        self.comm_socket.listen(5)

        self.t = time.time()

        # start thread for marker position detection
        self.estimator = ArucoEstimator(self.robot_ids.keys())
        self.estimator_thread = threading.Thread(target=self.estimator.run_tracking)
        self.estimator_thread.start()

        self.controller = MPCController(self.estimator)

        print("waiting for corner and robot detection..")
        while not self.estimator.all_robots_detected() or not self.estimator.all_corners_detected():
            pass
        print("everything detected!")
        # drive robots to initial position
        for robot_id in self.robot_ids:
            self.grid_control(robot_id, 'nop')

    def run(self):
        print("waiting until all markers are detected...")
        while not self.estimator.all_corners_detected():
            pass
        print("starting control")

        running = True
        while running:
            (clientsocket, address) = self.comm_socket.accept()
            clientsocket.settimeout(None)

            connected = True
            while connected:
                try:
                    data = myreceive(clientsocket)
                    print("data received: ", data)

                    inputs = data.split(b',')
                    cmd = inputs[0]
                    cmd = cmd.strip().decode()

                    if len(inputs) > 1:
                        if cmd in RemoteController.valid_cmds:
                            if cmd == 'initialize_grid':
                                try:
                                    grid_columns = int(inputs[1])
                                    grid_rows = int(inputs[2])

                                    self.estimator.grid_columns = grid_columns
                                    self.estimator.grid_rows = grid_rows
                                    clientsocket.sendall(b'OK\n')
                                except ValueError:
                                    print("could not initialize grid!")
                                    clientsocket.sendall(b'"could not initialize grid!\n'
                                                         b'expected: initialize_robot, <grid columns>, <grid rows>')
                                except IndexError:
                                    print("could not initialize grid!")
                                    clientsocket.sendall(b'"could not initialize grid!\n'
                                                         b'expected: initialize_robot, <grid columns>, <grid rows>')
                            else:   # robot command
                                try:
                                    robot_id = int(inputs[1])
                                except ValueError:
                                    robot_id = None
                                    print("could not read robot id!")
                                    clientsocket.sendall(b'Could not read robot id!\n')

                                if cmd == 'initialize_robot':
                                    # add a new robot to the game
                                    try:
                                        id = int(inputs[1])
                                        ip = inputs[2].decode().strip()
                                        x = int(inputs[3])
                                        y = int(inputs[4])
                                        orientation = inputs[5].decode().strip()

                                        print(f"initializing new robot with id {id} and ip {ip} at pos ({x},{y}) with "
                                              f"orientation '{orientation}'")
                                        if id not in self.robot_ids:
                                            new_robot = RoboRallyRobot(id=id, ip=ip, x=x, y=y, orientation=orientation)
                                            new_robot.connect()

                                            # store the new robot in the list of robots
                                            self.robots.append(new_robot)
                                            self.robot_ids[new_robot.id] = new_robot  # this also means the estimator
                                            # will track the new robot because
                                            # it got a reference to the list of
                                            # robot ids to keep an eye out for

                                        else: # robot is already created -> just re-initialize its position
                                            new_robot = self.robot_ids[id]
                                            new_robot.grid_x = x
                                            new_robot.grid_y = y
                                            new_robot.orientation = orientation

                                        if new_robot.connected:
                                            print("created new robot and successfully connected to it!")

                                            while not self.estimator.all_robots_detected(): # wait until the robot gets detected
                                                pass

                                            # drive the new robot to its starting position
                                            self.grid_control(new_robot.id, 'nop')

                                            clientsocket.sendall(b'OK\n')
                                        else:
                                            clientsocket.sendall(f"error: could not connect to new robot {new_robot}".encode())

                                    except IndexError:
                                        print("could not initialize a new robot")
                                        clientsocket.sendall('could not initialize a new robot: invalid command format\n'
                                            'expected: initialize_robot, <id>, <ip>, <x>, <y>, <orientation>\n'.encode())
                                    except ValueError:
                                        print("could not initialize a new robot")
                                        clientsocket.sendall('could not initialize a new robot: invalid command format\n'
                                            'expected: initialize_robot, <id>, <ip>, <x>, <y>, <orientation>\n'.encode())

                                elif robot_id in self.robot_ids:
                                    if cmd == b'get position':
                                        clientsocket.sendall(bytes('{}\n'.format(self.robot_ids[robot_id].grid_pos)))
                                    elif cmd == b'set position':
                                        try:
                                            pos_data = ",".join(inputs[2:])
                                            new_grid_pos = tuple(map(lambda x: int(x[1]) if x[0] < 2 else float(x[1]),
                                                                     enumerate(pos_data.strip().strip('()').split(','))))
                                            self.robot_ids[robot_id].grid_pos = new_grid_pos
                                            clientsocket.sendall(b'OK\n')
                                        except IndexError as e:
                                            print("could not set grid position!")
                                            clientsocket.sendall(bytes(
                                                'could not set grid position! (invalid format)\n'.format(
                                                    self.robot_ids[robot_id].grid_pos)))
                                        except ValueError as e:
                                            print("could not set grid position!")
                                            clientsocket.sendall(bytes(
                                                'could not set grid position! (invalid format)\n'.format(
                                                    self.robot_ids[robot_id].grid_pos)))
                                    else:
                                        self.grid_control(robot_id, cmd)
                                        clientsocket.sendall(b'OK\n')

                                else:
                                    print("invalid robot id!")
                                    clientsocket.sendall(b'Invalid robot id!\n')
                        else:
                            clientsocket.sendall(b'Invalid command!\n')
                    else:  # len(inputs) <= 1
                        if b'quit' in inputs[0]:
                            clientsocket.close()
                            self.comm_socket.close()
                            connected = False
                            running = False
                        else:
                            print("could not process command!")
                            clientsocket.sendall(b'Could not process request!\n')

                except RuntimeError:
                    print("disconnected")
                    connected = False
                    clientsocket.close()

    def grid_control(self, robot_id, cmd):
        robot = self.robot_ids[robot_id]  # get robot to be controlled

        print("robot grid pos before move: ", robot)
        robot.move(cmd)
        print("robot grid pos after move: ", robot)

        target = self.estimator.get_pos_from_grid_point(robot.grid_x, robot.grid_y, robot.grid_orientation)

        self.controller.move_to_pos(target, robot)


def main(args):
    rc = RemoteController()
    rc.run()


if __name__ == '__main__':
    main(sys.argv)
