from enum import Enum

import numpy as np
from PIL import Image

from simulator import EmulatorSimulator, MouseClick


class Action(Enum):
    CLICK = "click"
    MOUSE_DOWN = "mousedown"
    MOUSE_UP = "mouseup"
    MOUSE_MOVE = "mousemove"
    KEY = "key"
    TYPE = "type"


class Controller:
    def __init__(self, vm_name: str, username: str, password: str, host: str) -> None:
        self.simulator = EmulatorSimulator(vm_name=vm_name, username=username,
                                           password=password, host=host)
    
    def get_state(self) -> np.ndarray:
        image_path = self.simulator.get_screenshot()
        with Image.open(image_path) as img:
            return np.array(img)
    
    def step(self, action: Action, **action_args) -> None:
        if action == Action.CLICK:
            print(action_args)
            self.simulator.click(**action_args)
        elif action == Action.MOUSE_DOWN:
            self.simulator.mousedown(**action_args)
        elif action == Action.MOUSE_UP:
            self.simulator.mouseup(**action_args)
        elif action == Action.MOUSE_MOVE:
            self.simulator.mouse_move(**action_args)
        elif action == Action.KEY:
            self.simulator.key(**action_args)
        elif action == Action.TYPE:
            self.simulator.type(**action_args)