import abc
import enum
import os
import time

from PIL import Image, ImageDraw
import tkinter as tk
import datetime
from pynput import mouse, keyboard
from typing import Tuple, Dict, Union

import ctypes

PROCESS_PRE_MONITOR_DPI_AWARE = 2
ctypes.windll.shcore.SetProcessDpiAwareness(PROCESS_PRE_MONITOR_DPI_AWARE)

Position = Tuple[int, int]
Size = Tuple[int, int]
Color = Tuple[int, int, int, int]


class Colors:
    Left = (0, 255, 0, 100)
    Right = (255, 0, 0, 100)
    Middle = (255, 255, 0, 100)


class Button(tk.Button):
    def __init__(self, *args, **kwarges):
        super(Button, self).__init__(*args, **kwarges)
        self.pack(pady=10)
        self["state"] = "normal"

    def switch(self):
        self["state"] = "disable" if self["state"] == "normal" else "normal"


class Checkbutton(tk.Checkbutton):
    def __init__(self, *args, **kwargs):
        super(Checkbutton, self).__init__(*args, **kwargs)
        self.pack(pady=2)


class EventType(enum.IntEnum):
    MOUSE_MOVE = 0
    MOUSE_LEFT_CLICK = 1
    MOUSE_RIGHT_CLICK = 2
    MOUSE_MIDDLE_CLICK = 3
    KEY_BOARD = 5


class Handler(abc.ABC):
    @abc.abstractmethod
    def start(self):
        raise NotImplementedError

    @abc.abstractmethod
    def on_event(self, type_: EventType, content: tuple):
        raise NotImplementedError

    @abc.abstractmethod
    def stop(self):
        raise NotImplementedError


class RecordHandler(Handler):
    def __init__(self, record_fn=""):
        self._record_fn = record_fn
        self._records = []

    def start(self):
        pass

    def on_event(self, type_: EventType, content: tuple):
        tp = time.time_ns()
        data = f"{tp}-{type_}"
        for ele in content:
            data = f"{data}-{ele}"
        data = f"{data}\n"
        self._records.append(data)

    def stop(self):
        if self._record_fn:
            fn = self._record_fn
        else:
            now = datetime.datetime.now()
            fn = f"record-{now.year}-{now.month}-{now.day}-{now.hour}-{now.minute}-{now.second}.rec"
        with open(fn, 'w+') as file:
            for rec in self._records:
                file.write(rec)
            file.close()
        self._records.clear()


class DrawHandler(Handler):
    def __init__(self, size: Size):
        """Image Cache"""
        self._size = size
        self._cache = None
        self._position = None

    def start(self):
        self._cache = Image.new("RGBA", self._size, (0, 0, 0, 255))

    def on_event(self, type_: EventType, content: tuple):
        if type_ == EventType.MOUSE_MOVE:
            position = (content[0], content[1])
            if self._position:
                self._line(start=self._position, end=position)
            self._position = position
            return
        if type_ == EventType.MOUSE_LEFT_CLICK:
            self._ellipse(content[0], content[1], color=Colors.Left)
            return
        if type_ == EventType.MOUSE_RIGHT_CLICK:
            self._ellipse(content[0], content[1], color=Colors.Right)
            return
        if type_ == EventType.MOUSE_MIDDLE_CLICK:
            self._ellipse(content[0], content[1], color=Colors.Middle)
            return

    def stop(self):
        dirname = 'out'
        dir_path = os.path.join(os.path.dirname(__file__), dirname)
        os.makedirs(dir_path, exist_ok=True)
        now = datetime.datetime.now()
        file_path = os.path.join(
            dir_path,
            f"mouse_track-{now.year}-{now.month}-{now.day}-{now.hour}-{now.minute}-{now.second}.png",
        )
        self._cache.save(file_path)
        print(f"轨迹图像已保存: {file_path}")

    def _line(self, start: Position, end: Position):
        """
        Draw a line
        Parameters:
        - start: tuple of the line's start
        - end: tuple of the line's end
        """
        self._draw_transp_line(xy=[start, end], fill=(255, 255, 255, 50), width=2)

    def _ellipse(self, x, y, color: Color, radius=10):
        """
        Draw a point at `(x, y)`
        :param x:
        :param y:
        :param color: 注意：有四个值，最后一个值的不透明度最大值是255，不是float的0~1
        :param radius:
        :return: None
        """
        self._draw_transparent_ellipse(
            [(x - radius, y - radius), (x + radius, y + radius)],
            fill=color,
        )

    def _draw_transp_line(self, xy, **kwargs):
        """
        Draws a line inside the given bounding box onto given image.
        Supports transparent colors
        """
        transp = Image.new("RGBA", self._size, (0, 0, 0, 0))  # Temp drawing image.
        draw = ImageDraw.Draw(transp, "RGBA")
        draw.line(xy, **kwargs)
        # Alpha composite two images together and replace first with result.
        self._cache.paste(Image.alpha_composite(self._cache, transp))

    def _draw_transparent_ellipse(self, xy, **kwargs):
        """
        Draws an ellipse inside the given bounding box onto given image.
        Supports transparent colors
        https://stackoverflow.com/a/54426778
        """
        transp = Image.new("RGBA", self._size, (0, 0, 0, 0))  # Temp drawing image.
        draw = ImageDraw.Draw(transp, "RGBA")
        draw.ellipse(xy, **kwargs)
        # Alpha composite two images together and replace first with result.
        self._cache.paste(Image.alpha_composite(self._cache, transp))


class Agent:
    def __init__(self):
        self._mouse_listener = mouse.Listener(on_move=self._on_mouse_move, on_click=self._on_mouse_click)
        self._keyboard_listener = keyboard.Listener(on_press=self._on_keyboard_press,
                                                    on_release=self._on_keyboard_release)
        self._handlers = []
        self._key_set = set()

    def add_listener(self, handler: Handler):
        self._handlers.append(handler)

    def start(self):
        self._mouse_listener.start()
        self._keyboard_listener.start()
        for handler in self._handlers:
            handler.start()

    def stop(self):
        self._mouse_listener.stop()
        self._keyboard_listener.stop()
        for handler in self._handlers:
            handler.stop()

    def _notify(self, type_: EventType, content: tuple):
        for handler in self._handlers:
            handler.on_event(type_, content)

    def _on_keyboard_press(self, key: Union[keyboard.Key, keyboard.KeyCode, None]):
        if key is None:
            return
        if isinstance(key, keyboard.Key):
            return self._key_set.add(key)
        if isinstance(key, keyboard.KeyCode):
            return
        raise TypeError("Invalid input key: ", key)

    def _on_keyboard_release(self, key: Union[keyboard.Key, keyboard.KeyCode, None]):
        if key is None:
            return
        if isinstance(key, keyboard.Key):
            if key in self._key_set:
                self._key_set.remove(key)
            return
        if isinstance(key, keyboard.KeyCode):
            content = tuple(self._key_set)
            content = content.__add__((key,))
            return self._notify(EventType.KEY_BOARD, content)
        raise TypeError("Invalid input key: ", key)

    def _on_mouse_click(self, x, y, button, pressed):
        if not pressed:
            return
        if not isinstance(button, mouse.Button):
            raise TypeError("Invalid input button: ", button)
        if button == mouse.Button.left:
            return self._notify(EventType.MOUSE_LEFT_CLICK, (x, y))
        if button == mouse.Button.right:
            return self._notify(EventType.MOUSE_RIGHT_CLICK, (x, y))
        if button == mouse.Button.middle:
            return self._notify(EventType.MOUSE_MIDDLE_CLICK, (x, y))

    def _on_mouse_move(self, x: int, y: int):
        self._notify(EventType.MOUSE_MOVE, (x, y))


class App(tk.Tk):
    def __init__(self):
        super(App, self).__init__()

        self.title("Mouse Tracker")
        self.geometry("400x100")

        self.start_button = Button(self, text="开始记录", command=self.start_tracking)
        self.stop_button = Button(self, text="停止记录", command=self.stop_tracking)
        self.stop_button.switch()
        self.agent = Agent()
        self.recoder = RecordHandler()
        self.drawer = DrawHandler(size=(self.winfo_screenwidth(), self.winfo_screenheight()))
        self.agent.add_listener(self.recoder)
        self.agent.add_listener(self.drawer)

    def start_tracking(self):
        """点击开始记录"""
        print("----------------- start_tracking")
        self.start_button.switch()
        self.stop_button.switch()
        self.agent.start()

    def stop_tracking(self):
        """点击结束记录"""
        print("----------------- stop_tracking")
        self.stop_button.switch()
        self.start_button.switch()
        self.agent.stop()


def main():
    app = App()
    app.mainloop()


if __name__ == "__main__":
    main()