import sys, pathlib; sys.path.append(str(pathlib.Path(__file__).parents[1]))

import os
import json
from desktop_env.envs.desktop_env import Action, MouseClick

class DuckTrackEventActionConverter:
    def __init__(self, human_readable: str, compress_move: bool = True):
        self.human_readable = human_readable
        self.compress_move = compress_move

    def enum_to_str(self, enum):
        """Converts an enum to its string representation if HUMAN_READABLE is True, otherwise returns its value."""
        return enum.name if self.human_readable else enum.value

    def compress_mouse_move(self, data: list[dict], index: int):
        """Compresses consecutive mouse move events into first and last move events."""
        first_move, last_move = data[index], data[index]
        while index < len(data) and data[index]["action"] == "move":
            last_move = data[index]
            index += 1
        return first_move, last_move, index

    def move_event_to_action(self, event: dict):
        return {"action_type": self.enum_to_str(Action.MOUSE_MOVE),
                "x": event["x"],
                "y": event["y"]}

    def click_event_to_action(self, event: dict):
        action = {}
        mouse_button = event["button"]
        mouse_pressed = event["pressed"]

        if mouse_pressed == True:
            action["action_type"] = self.enum_to_str(Action.MOUSE_DOWN)
        elif mouse_pressed == False:
            action["action_type"] = self.enum_to_str(Action.MOUSE_UP)
        else:
            raise NotImplementedError(mouse_pressed)

        if mouse_button == "left":
            action["click_type"] = self.enum_to_str(MouseClick.LEFT)
        elif mouse_button == "right":
            action["click_type"] = self.enum_to_str(MouseClick.RIGHT)
        elif mouse_button == "middle":
            action["click_type"] = self.enum_to_str(MouseClick.MIDDLE)
        else:
            raise NotImplementedError(mouse_button)

        return action

    def press_event_to_action(self, event: dict):
        return {"action_type": self.enum_to_str(Action.KEY_DOWN),
                "key": [ord(c) for c in event["name"]]}

    def release_event_to_action(self, event: dict):
        return {"action_type": self.enum_to_str(Action.KEY_UP),
                "key": [ord(c) for c in event["name"]]}

    def scroll_event_to_action(self, event: dict):
        # TODO: need to confirm if df < 0 means scroll up or down
        if event["dy"] < 0:
            down = False
        else:
            down = True

        return {"action_type": self.enum_to_str(Action.CLICK),
                "click_type": self.enum_to_str(MouseClick.WHEEL_DOWN) if down else self.enum_to_str(MouseClick.WHEEL_UP)}

    def event_to_action(self, event: dict):
        """Converts an event to its corresponding action based on the event type."""
        if event["action"] == "move":
            return self.move_event_to_action(event)
        elif event["action"] == "click":
            return self.click_event_to_action(event)
        elif event["action"] == "press":
            return self.press_event_to_action(event)
        elif event["action"] == "release":
            return self.release_event_to_action(event)
        elif event["action"] == "scroll":
            return self.scroll_event_to_action(event)
        else:
            raise NotImplementedError(event["action"])

    def ducktrack_event_file_to_action(self, ducktrack_event_file: str, out_file: str, compress_move: bool = None):
        """Converts DuckTrack event data to a list of actions and saves them to a file."""
        if not os.path.exists(ducktrack_event_file):
            raise FileNotFoundError(ducktrack_event_file)
        
        # set to default
        if compress_move == None:
            compress_move = self.compress_move
        
        with open(ducktrack_event_file, 'r') as file:
            data = [json.loads(line) for line in file]
        
        result = {"action": [], "event": []}
        index = 0
        
        while index < len(data):
            event = data[index]
            if event["action"] == "move" and compress_move:
                first_move, last_move, index = self.compress_mouse_move(data, index)
                result["action"].extend([self.event_to_action(first_move), self.event_to_action(last_move)])
                result["event"].extend([first_move, last_move])
            else:
                result["action"].append(self.event_to_action(event))
                result["event"].append(event)
                index += 1
        
        with open(out_file, "w") as f:
            json.dump(result, f)

if __name__ == "__main__":
    converter = DuckTrackEventActionConverter(human_readable=False)
    converter.ducktrack_event_file_to_action(ducktrack_event_file="sample.jsonl",
                                             out_file="output.json",
                                             compress_move=True)