from typing import Optional
from xml.etree import ElementTree

from omegaconf import DictConfig

import pygame
from pygame.locals import *
from pygame_gui import UIManager
from pygame_gui.windows import UIConsoleWindow
from pygame_gui.elements import UILabel

import pymasep as pm
from pymasep.communication import Message, Client, SocketHandlerQueue, Frame
from pymasep.application.sub_app import SubApp
from pymasep.application.connected_sub_app_info import ConnectedSubAppInfo
from pymasep.common.exception import CreationException


class Interface(SubApp):
    """
    Sub application that display an environment and interact with user

    :param app: main application
    :param received_q_id: id of the queue used by the sub app to receive messages.
    :param cfg: configuration of the interface
    :param game_class_: class of the game
    :param interface_role: Role of the interface. see @ConnectedSubAppInfo.
    :param remote_engine_host: host of the remote engine. May be None if interface is in the same process of engine
    :param remote_engine_port: port of the remote engine. May be None if interface is in the same process of engine
    :param search_remote: True if the interface have to search the server on LAN. Default False
    """

    def __init__(self, app, received_q_id,
                 cfg: DictConfig,
                 game_class_=None,  # None for subclass, mandatory in the base class
                 interface_role=None,  # None for subclass, mandatory in the base class
                 remote_engine_host: Optional[str] = None,
                 remote_engine_port: Optional[int] = None,
                 search_remote: bool = False):
        super().__init__(app=app, received_q_id=received_q_id, cfg=cfg)
        # display related information
        self.font = None
        self.screen = None
        self.virtual_screen = None
        self.color = (255, 255, 0)
        self.backgroundColor = (0, 0, 100)
        self.fontSize = 60
        self.textPos = (0, 0)
        self.virtual_screen_size = (100, 100)
        self.end = False
        self.root_path = self.config['root_path']
        self.ui_manager = None
        self.ui_theme_file = None

        # environment related information
        if game_class_ is None or interface_role is None:
            raise CreationException('game_class_ and interface_role must not be None')
        self.interface_role = interface_role
        self.environment = pm.interface.EnvironmentInterface(id_env=self.id)
        self.environment.game = game_class_(cfg=cfg)
        self.controller_interface_class_ = None
        self.controller = None
        self.last_step_choose_action = None

        # interface element related information
        self.console = None

        # debug related information
        self.interface_dbg_info_index = 0
        self.debug_info_labels = []

        # client/server/engine related information
        self.engine_id = None
        self.remote_engine_host = remote_engine_host
        self.remote_engine_port = remote_engine_port
        if search_remote or remote_engine_host is not None:
            self.remote_communication = Client(ident=self.id,
                                               host=remote_engine_host,
                                               port=remote_engine_port,
                                               socket_handler_class=SocketHandlerQueue,
                                               socket_handler_args=(self.messages_received,))
            if search_remote:
                self.remote_communication.search_server()
                if self.remote_communication.host is None or self.remote_communication.port is None:
                    raise ConnectionError
            self.remote_communication.connect()

    def init_graphics(self):
        """
        Init the graphical elements
        """
        pygame.init()
        pygame.font.init()
        self.font = pygame.font.SysFont("dejavusansmono", self.fontSize)

        # Warning : pygame.SCALED slows down fps
        # flags = pygame.SCALED
        flags = 0
        # Keep 16:9 ratio (1368, 768) (1920, 1080) (2560, 1440)
        # https://www.sven.de/dpi/
        # https://developer.android.com/training/multiscreen/screendensities
        # https://benui.ca/unreal/ui-resolution/
        # https://blogs.windows.com/windowsdeveloper/2017/05/19/improving-high-dpi-experience-gdi-based-desktop-apps/
        screen_resolution = self.config.interface.get('resolution', [1368, 768])
        self.screen = pygame.display.set_mode(tuple(screen_resolution), flags=flags)
        self.virtual_screen = pygame.Surface(self.virtual_screen_size)

        self.ui_manager = UIManager(self.virtual_screen_size,
                                    theme_path=self.ui_theme_file)
        ratio_x = (self.virtual_screen.get_width() / self.screen.get_width())
        ratio_y = (self.virtual_screen.get_height() / self.screen.get_height())
        # undocumented pygame_gui feature... hope it will stay as is :)
        # see https://github.com/MyreMylar/pygame_gui/issues/210
        self.ui_manager.mouse_pos_scale_factor = [ratio_x, ratio_y]

        for k_gr in self.resources.keys():
            self.resources[k_gr] = self.resources[k_gr].convert(self.virtual_screen)

        console_size = (self.virtual_screen_size[0], self.virtual_screen_size[1] / 3)
        self.console = UIConsoleWindow(pygame.Rect((0, 0), console_size),
                                       manager=self.ui_manager,
                                       window_title='Console',
                                       object_id='#console',
                                       visible=False)
        self.console.change_layer(200)

        self.interface_dbg_info_index = self.add_debug_label()

    def add_debug_label(self):
        """
        Adding a new graphical debug labels. A text can be set later to add debug info, displayed on the top of all.

        :return the index of the new added label
        """
        pos_y = 0
        font_info = self.ui_manager.get_theme().get_font_info(["#fps_counter"])
        if self.debug_info_labels:
            pos_y = self.debug_info_labels[-1].get_relative_rect()[1] + int(int(font_info['size']) * 1.5)
        label = UILabel(pygame.Rect(0,
                                    pos_y,
                                    int(self.virtual_screen.get_width() * 3 / 4),
                                    int(int(font_info['size']) * 1.5)),
                        "",
                        self.ui_manager,
                        object_id='#fps_counter')

        label.change_layer(1000)
        label.hide()
        self.debug_info_labels.append(label)
        return len(self.debug_info_labels) - 1

    def init_game(self):
        """
        Initialize game : Connect to engine, define the game and register self as a new interface
        """
        if self.remote_communication is not None:
            self.engine_id = self.remote_communication.server_id
            self.register_connected_sub_app(sub_app_id=self.engine_id,
                                            role=ConnectedSubAppInfo.ROLE_SERVER,
                                            sent_q=None)
        else:
            self.engine_id = list(self.connected_sub_app.keys())[0]

        self.controller = self.controller_interface_class_(self.environment)

        init_message = Message(Message.MESSAGE_TYPE_DEFINE, {'game': self.environment.game.__class__}, self.id)
        self.send_message(init_message, self.engine_id)

        register_message = Message(Message.MESSAGE_TYPE_REGISTER,
                                   {'id': self.id,
                                    'queue_id': self.messages_received_q_id,
                                    'role': self.interface_role},
                                   self.id)
        self.send_message(register_message, self.engine_id)

    def receive_observation(self):
        """
        Wait for a OBSERVATION message (or END_GAME) from the engine and update the interface environment
        """
        obs_msg = self.wait_message([Message.MESSAGE_TYPE_OBSERVATION, Message.MESSAGE_TYPE_END_GAME], block=False)
        if obs_msg is not None:
            if obs_msg.msg_type == Message.MESSAGE_TYPE_OBSERVATION:
                observation_xml = obs_msg.params['observation']
                self.logger.info('Observation received ' + str(ElementTree.tostring(observation_xml)))

                additional_info = dict()
                if 'additional_information' in obs_msg.params:
                    additional_info = obs_msg.params['additional_information']
                    self.logger.info('Additional info received' + str(additional_info))
                    if 'log' in additional_info:
                        self.console.add_output_line_to_log(str(additional_info['log']))

                self.environment.update_current_next(observation_xml, additional_info)

                for ag in self.environment.agents:
                    if ag.control == self.id:
                        ag.controller = self.controller
                        self.environment.controlled_agent = ag
                self.messages_received.task_done()

                self.get_current_app_state().on_receive_observation()
            if obs_msg.msg_type == Message.MESSAGE_TYPE_END_GAME:
                self.logger.info('End Game message received from Engine')
                self.end = True
        else:
            self.logger.debug('No observation received')

    def handle_event(self):
        """
        Handle interface events
        """
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                self.end = True
            if event.type == pygame.KEYDOWN:
                if event.key == K_HASH:
                    if self.console.visible == 0:
                        self.console.show()
                    else:
                        self.console.hide()
                if event.key == K_d:
                    if event.mod & pygame.KMOD_CTRL:
                        for d in self.debug_info_labels:
                            if d.visible == 0:
                                d.show()
                            else:
                                d.hide()
                    if event.mod & pygame.KMOD_ALT:
                        self.ui_manager.set_visual_debug_mode(True)
                    else:
                        self.ui_manager.set_visual_debug_mode(False)

            self.ui_manager.process_events(event)
            self.get_current_app_state().handle_event(event)

    def choose_action(self):
        """
        Let the agent choose action if conditions are respected :
        - The interface is a ROLE_ACTOR
        - The interface has received a new observation since the agent's last action.
        """
        if self.interface_role == ConnectedSubAppInfo.ROLE_ACTOR:
            # last_step_choose_action is used to choose the (interface) step when an agent should choose an action
            # It occurs when the interface received a new observation
            if self.last_step_choose_action is None or self.last_step_choose_action < self.environment.current_step:
                self.environment.choose_action()
                if self.environment.controlled_agent is not None and self.environment.controlled_agent.action:
                    self.last_step_choose_action = self.environment.current_step

    def send_action(self):
        """
        Send the action chosen by the interface. Only if the interface is a ROLE_ACTOR
        """
        if self.interface_role == ConnectedSubAppInfo.ROLE_ACTOR:
            if self.environment.controlled_agent is not None:
                if self.environment.controlled_agent.action is not None:
                    act_msg = Message(msg_type=Message.MESSAGE_TYPE_ACTION,
                                      param_dict={'action': self.environment.controlled_agent.action},
                                      src_sub_app_id=self.id)
                    self.send_message(act_msg, self.engine_id)
                    self.logger.info(
                        'Action sent ' + str(ElementTree.tostring(self.environment.controlled_agent.action.to_xml())))
                    # send the action only once
                    self.environment.controlled_agent.action = None

    def update(self):
        """ update the current app state """
        self.get_current_app_state().update()

    def render(self, time_delta):
        """
        Render the environment current state and other information on screen
        """
        self.virtual_screen.fill('black')
        if self.environment.current_state is not None:
            self.ui_manager.update(time_delta)
            self.ui_manager.draw_ui(self.virtual_screen)
            self.get_current_app_state().render()
        self.screen.blit(pygame.transform.smoothscale(self.virtual_screen, self.screen.get_rect().size), (0, 0))
        # self.screen.blit(pygame.transform.scale(self.virtual_screen, self.screen.get_rect().size), (0, 0))
        # self.screen.blit(self.virtual_screen, (0, 0))
        pygame.display.flip()
        if time_delta != 0:
            mouse_pos = self.ui_manager.get_mouse_position()
            self.debug_info_labels[self.interface_dbg_info_index].set_text(f'FPS: {1 / time_delta:>10.2f}, '
                                                                           f'Mouse pos : {str(pygame.mouse.get_pos())} '
                                                                           f'Mouse pygame_gui pos : {str(mouse_pos)}')

    def run(self):
        """
        Main interface loop
        """
        self.init_graphics()
        self.init_game()

        self.receive_observation()

        self.render(0.0)

        clock = pygame.time.Clock()
        self.end = False
        while not self.end:
            time_delta = clock.tick(500) / 1000
            self.handle_event()
            self.update()
            self.render(time_delta)

        msg_quit = Message(msg_type=Message.MESSAGE_TYPE_QUIT_GAME, param_dict=None, src_sub_app_id=self.id)
        self.send_message(msg_quit, self.engine_id)
        if self.remote_communication is not None:
            self.remote_communication.send_frame_to(self.engine_id, Frame.bye_frame())
