from controller import Supervisor
import json
import os
import time
import random
import math

from watchdog.observers import Observer
from watchdog.events import PatternMatchingEventHandler
from referee.referee import RCJSoccerReferee
from referee.consts import (BALL_DEPTH,
                            FIELD_X_UPPER_LIMIT, FIELD_X_LOWER_LIMIT,
                            FIELD_Z_UPPER_LIMIT, FIELD_Z_LOWER_LIMIT)

STATE_FILE = "state.json"
CONTROLLERS_DIR = ".."
SUPERVISOR_NAME = os.path.split(os.getcwd())[1]

class Richo(RCJSoccerReferee):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.reset_controllers_flag = False
        self.reset_controllers_last_time = time.time()

        self.restoreState()
        self.start_watchdog()

    def saveState(self):
        try:
            data = {"Y": self.get_current_controller("Y"),
                    "B": self.get_current_controller("B"),
                    "saved_snapshot": self.saved_snapshot,
                    "check_timer_flag": self.check_timer_flag,
                    "check_progress_flag": self.check_progress_flag,
                    "check_goal_flag": self.check_goal_flag,
                    "check_robots_in_penalty_area_flag": self.check_robots_in_penalty_area_flag}
            with open(STATE_FILE, "w") as file:
                json.dump(data, file)
        except:
            print(f"El archivo {STATE_FILE} no se pudo escribir")

    def restoreState(self):
        try:
            with open(STATE_FILE) as file:
                data = json.load(file)

            self.saved_snapshot = data.get("saved_snapshot", None)
            self.check_timer_flag = data.get("check_timer_flag", True)
            self.check_progress_flag = data.get("check_progress_flag", True)
            self.check_robots_in_penalty_area_flag = data.get("check_robots_in_penalty_area_flag", True)
            self.check_goal_flag = data.get("check_goal_flag", True)
            if "Y" in data:
                self.set_controller("Y", data["Y"])
            if "B" in data:
                self.set_controller("B", data["B"])
        except:
            print(f"El archivo {STATE_FILE} no se pudo leer")
            self.saved_snapshot = None
            self.check_timer_flag = True
            self.check_progress_flag = True
            self.check_robots_in_penalty_area_flag = True
            self.check_goal_flag = True

    def start_watchdog(self):
        def on_any_event(_):
            self.reset_controllers_flag = True

        event_handler = PatternMatchingEventHandler(patterns=["*.py"],
                                                    ignore_patterns=[],
                                                    ignore_directories=True)
        event_handler.on_any_event = on_any_event
        path = CONTROLLERS_DIR

        self.observer = Observer()
        self.observer.schedule(event_handler, path, recursive=True)
        self.observer.start()

    def stop_watchdog(self):
        self.observer.stop()

    def send(self, __key, **args):
        messageString = json.dumps({"msg": __key, "args": args})
        self.wwiSendText(messageString)

    def alert(self, msg):
        self.send("alert", message=msg)

    def log(self, msg):
        self.send("log", message=msg)

    def check_progress(self):
        if self.check_progress_flag:
            super().check_progress()

    def check_goal(self):
        if self.check_goal_flag:
            super().check_goal()

    def check_robots_in_penalty_area(self):
        if self.check_robots_in_penalty_area_flag:
            super().check_robots_in_penalty_area()

    def get_current_controller(self, team_name):
        return next(map(lambda n: self.robot_nodes[n].getField("controller").getSFString(),
                    filter(lambda n: n[0] == team_name, self.robot_nodes)))

    def tick(self):
        self.checkWatchdog()
        self.checkIncomingMessages()
        self.sendCurrentState()

        return_value = super().tick()
        if self.check_timer_flag:
            return return_value
        else:
            return True

    def sendCurrentState(self):
        selected = self.getSelected()
        self.send("update",
                time=self.getTime(),
                selected=selected.getDef() if selected is not None else None,
                ball_translation=self.ball_translation,
                robot_translation=self.robot_translation,
                robot_rotation=self.robot_rotation)

    def checkWatchdog(self):
        if self.reset_controllers_flag:
            if time.time() - self.reset_controllers_last_time > 1:
                self.update_controllers_list()
                self.reset_controllers()
                self.reset_controllers_flag = False
                self.reset_controllers_last_time = time.time()

    def checkIncomingMessages(self):
        # Get the message in from the robot window(if there is one)
        message_text = self.wwiReceiveText()

        # If there is a message
        if message_text != "":
            print(message_text)
            message = json.loads(message_text)
            key = message["msg"]
            args = message["args"]
            response_id = message["response_id"]

            if key == "setup":
                self.update_controllers_list()
                self.update_flags()
            elif key == "reset":
                self.reset_controllers()
            elif key == "set_controller":
                self.set_controller(args["team"], args["controller"])
            elif key == "set_check_timer":
                self.check_timer_flag = args["enabled"]
            elif key == "set_check_progress":
                self.check_progress_flag = args["enabled"]
            elif key == "set_check_goal":
                self.check_goal_flag = args["enabled"]
            elif key == "set_check_robots_in_penalty_area":
                self.check_robots_in_penalty_area_flag = args["enabled"]
            elif key == "save_state":
                self.save_snapshot()
            elif key == "restore_state":
                self.restore_snapshot()
            elif key == "randomize_ball":
                x = random.uniform(FIELD_X_LOWER_LIMIT + 0.1, FIELD_X_UPPER_LIMIT - 0.1)
                z = random.uniform(FIELD_Z_LOWER_LIMIT + 0.1, FIELD_Z_UPPER_LIMIT - 0.1)
                self.set_ball_position([x, BALL_DEPTH, z])
            elif key == "move_object":
                self.move_object(args["object"], args["property"], args["value"])
            elif key == "move_out":
                self.move_robots_out_of_field()

            # Save the state after ANY message received
            self.saveState()

    def save_snapshot(self):
        data = {}
        data["BALL"] = {
            "translation": self.ball.getField("translation").getSFVec3f(),
            "rotation": self.ball.getField("rotation").getSFRotation(),
            "velocity": self.ball.getVelocity()
        }

        for robot_name in self.robot_nodes:
            robot = self.robot_nodes[robot_name]
            data[robot_name] = {
                "translation": robot.getField("translation").getSFVec3f(),
                "rotation": robot.getField("rotation").getSFRotation(),
            }
        self.saved_snapshot = data

    def restore_snapshot(self):
        if self.saved_snapshot is None: return
        for obj_def in self.saved_snapshot:
            obj = self.getFromDef(obj_def)
            data = self.saved_snapshot[obj_def]
            obj.getField("translation").setSFVec3f(data["translation"])
            obj.getField("rotation").setSFRotation(data["rotation"])
            if "velocity" in data:
                obj.setVelocity(data["velocity"])
            else:
                obj.resetPhysics()


    def move_object(self, object_def, property_name, property_value):
        object = self.getFromDef(object_def)
        if object is None: return

        object.setVelocity([0, 0, 0, 0, 0, 0])
        object.resetPhysics()

        if property_name == "a":
            if object_def == "BALL": return
            field = object.getField("rotation")
            field_value = [0, 1, 0, math.radians(property_value)]
            field.setSFRotation(field_value)
            self.robot_rotation[object_def] = field_value
        else:
            field = object.getField("translation")
            field_value = field.getSFVec3f()
            if property_name == "x":
                field_value[0] = property_value
            elif property_name == "y":
                field_value[2] = property_value
            field.setSFVec3f(field_value)
            if object_def == "BALL":
                self.ball_translation = field_value
                self.ball_stop = 2
            else:
                self.robot_translation[object_def] = field_value

    def move_robots_out_of_field(self):
        for robot_name in self.robot_nodes:
            yellow = robot_name[0] == "Y"
            index = int(robot_name[1]) - 1

            robot = self.getFromDef(robot_name)
            robot.setVelocity([0, 0, 0, 0, 0, 0])
            robot.resetPhysics()

            # Rotation
            rotation = [0, 1, 0, math.radians(90 * (1 if yellow else -1))]
            robot.getField("rotation").setSFRotation(rotation)
            self.robot_rotation[robot_name] = rotation

            # Translation
            translation = robot.getField("translation").getSFVec3f()
            translation[0] = -0.814 * (1 if yellow else -1)
            translation[2] = (0.283 + (0.1 * index)) * (1 if yellow else -1)
            robot.getField("translation").setSFVec3f(translation)
            self.robot_translation[robot_name] = translation


    def reset_controllers(self):
        for r in self.robot_nodes:
            self.robot_nodes[r].restartController()
        self.log("RESET")

    def set_controller(self, team_name, controller):
        for robot_name in self.robot_nodes:
            if robot_name[0] == team_name:
                robot = self.robot_nodes[robot_name]
                robot.getField("controller").setSFString(controller)

    def update_controllers_list(self):
        controllers=os.listdir(CONTROLLERS_DIR)
        controllers.remove(SUPERVISOR_NAME)
        self.send("update_controllers_list",
                controllers=controllers,
                Y=self.get_current_controller("Y"),
                B=self.get_current_controller("B"))

    def update_flags(self):
        self.send("update_flags",
                check_timer=self.check_timer_flag,
                check_progress=self.check_progress_flag,
                check_robots_in_penalty_area=self.check_robots_in_penalty_area_flag,
                check_goal=self.check_goal_flag)
