import json
import logging
import time
import xml.etree.ElementTree as ElementTree
from pathlib import Path

import gym
import marlo
import numpy as np
from gym.spaces import Discrete

from domain.actions import Turn, Teleport, Pitch, Yaw
from symbols.domain.option_space import OptionSpace

logger = logging.getLogger(__name__)

#C:\Users\\Anaconda3\envs\marlo\Minecraft


def create_basic_minecraft_env(mission_xml):
    client_pool = [('127.0.0.1', 10000), ('127.0.0.1', 10001), ('127.0.0.1', 10002),  ('127.0.0.1', 10003),  ('127.0.0.1', 10004), ('127.0.0.1', 10005)]
    params = {
        "client_pool": client_pool,
        "suppress_info": False,
        # "allowContinuousMovement": ["move", "turn"],
        'allowDiscreteMovement': True,
        'allowContinuousMovement': False,
        'allowAbsoluteMovement': False,
        'continuous_discrete': True,
        "recordMP4": [30, 2000000],
        "recordObservations": True,
        "recordDestination": "video.tgz",
        'recordCommands': True
    }

    if Path(mission_xml).is_file():
        join_tokens = marlo.make(mission_xml,
                                 params=params)
    else:
        params["mission_xml"] = mission_xml
        env = gym.make("MarLo-RawXMLEnv-v0")
        join_tokens = env.init(params, dry_run=True)


    # As this is a single agent scenario,
    # there will just be a single token
    assert len(join_tokens) == 1
    join_token = join_tokens[0]
    return marlo.init(join_token)


def create_minecraft_env(mission_xml, wrap=True):
    client_pool = [('127.0.0.1', 10000), ('127.0.0.1', 10001), ('127.0.0.1', 10002)]
    params = {
        "client_pool": client_pool,
        "suppress_info": False,
        "recordMP4": [30, 2000000],
        "recordObservations": True,
        "recordDestination": "video.tgz",
        'recordCommands': True
    }

    if Path(mission_xml).is_file():
        join_tokens = marlo.make(mission_xml,
                                 params=params)
    else:

        params["mission_xml"] = mission_xml
        env = gym.make("MarLo-RawXMLEnv-v0")
        join_tokens = env.init(params, dry_run=True)


    # As this is a single agent scenario,
    # there will just be a single token
    assert len(join_tokens) == 1
    join_token = join_tokens[0]

    if not wrap:
        return marlo.init(join_token)

    return TeleportWrapper(marlo.init(join_token))


class TeleportWrapper(gym.ActionWrapper):


    def _desc(self, x):
        if x == 0:
            return "WalkToItem"
        if x == 1:
            return "AttackItem"
        if x == 2:
            return "PickupItem"
        if x == 3:
            return "WalkNorthDoor"
        if x == 4:
            return "WalkSouthDoor"
        if x == 5:
            return "WalkThroughDoor"
        if x == 6:
            return "Craft"
        if x == 7:
            return "OpenChest"
        if x == 8:
            return "ToggleDoor"

    def __init__(self, env):
        super().__init__(env)
        self.action_space = OptionSpace(9, self._desc)


    def reset(self):
        self.env.reset()
        # add tiny rotation because minecraft doesn't redraw on teleport without movement :/
        obs, _, _, _ = self.step(Turn(0.0001))
        return obs

    def get_observation(self):
        world_state = self.env._get_world_state()
        obs = self.env._get_observation(world_state)
        return obs

    def get_image(self):
        world_state = self.env._get_world_state()
        return self.env._get_video_frame(world_state)

    def get_x(self):
        world_state = self.env._get_world_state()
        obs = self.env._get_observation(world_state)
        return obs['XPos']

    def get_y(self):
        world_state = self.env._get_world_state()
        obs = self.env._get_observation(world_state)
        return obs['YPos']

    def get_z(self):
        world_state = self.env._get_world_state()
        obs = self.env._get_observation(world_state)
        return obs['ZPos']

    def get_yaw(self):
        world_state = self.env._get_world_state()
        obs = self.env._get_observation(world_state)
        return obs['Yaw']

    def get_pitch(self):
        world_state = self.env._get_world_state()
        obs = self.env._get_observation(world_state)
        return obs['Pitch']

    def step(self, action):

        if isinstance(action, int):
            return self.env.step(action)

        val = action.value
        if not isinstance(val, str):
            return self.env.step(val)
        return self.step_wrapper(val)

    def has_item(self, item, amount=1):
        world_state = self.env._get_world_state()
        obs = self.env._get_observation(world_state)
        for i in range(9):
            val = obs["Hotbar_{}_item".format(i)]
            if val == item and obs["Hotbar_{}_size".format(i)] >= amount:
                return True
        return False

    def step_wrapper(self, action):
        world_state = self.env.agent_host.peekWorldState()
        if world_state.is_mission_running:
            # if self.env._turn:
            #     if self.env._turn.can_play or True:
            # send corresponding command
            logger.debug(action)

            self.env.agent_host.sendCommand(action)
            # self.env.send_command(action)

        world_state = self.env._get_world_state()

        # Update turn state
        if world_state.number_of_observations_since_last_state > 0:
            data = json.loads(world_state.observations[-1].text)
            turn_key = data.get(u'turn_key', None)
            if turn_key is not None and turn_key != self.env._turn.key:
                self.env._turn.update(turn_key)
                # Log
        for error in world_state.errors:
            logger.warn(error.text)
        for message in world_state.mission_control_messages:
            logger.debug(message.text)
            root = ElementTree.fromstring(message.text)
            if root.tag == '{http://ProjectMalmo.microsoft.com}MissionEnded':
                for el in root.findall(
                        '{http://ProjectMalmo.microsoft.com}HumanReadableStatus'  # noqa: E501
                ):
                    logger.info("Mission ended: %s", el.text)
                    # Compute Rewards
        reward = 0
        for _reward in world_state.rewards:
            reward += _reward.getValue()
            # Get observation
        image = self.env._get_video_frame(world_state)
        # detect if done ?
        done = not world_state.is_mission_running or any(world_state.errors)

        # gather info
        info = {}
        info['has_mission_begun'] = world_state.has_mission_begun
        info['is_mission_running'] = world_state.is_mission_running
        info[
            'number_of_video_frames_since_last_state'] = world_state.number_of_video_frames_since_last_state  # noqa: E501
        info['number_of_rewards_since_last_state'] = world_state.number_of_rewards_since_last_state  # noqa: E501
        info[
            'number_of_observations_since_last_state'] = world_state.number_of_observations_since_last_state  # noqa: E501
        info['mission_control_messages'] = [msg.text for msg in world_state.mission_control_messages]  # noqa: E501
        info['observation'] = self.env._get_observation(world_state)

        if self.env.params.suppress_info:
            """
            Clear info variable to not leak in game variables
            """
            info = {}
            # Notify evaluation system, if applicable
        # marlo.CrowdAiNotifier._env_action(action)
        marlo.CrowdAiNotifier._step_reward(reward)
        if done:
            marlo.CrowdAiNotifier._episode_done()
        return image, reward, done, info

    def perturb(self):
        time.sleep(0.2)
        x, y, z, yaw, pitch = self.get_x(), self.get_y(), self.get_z(), self.get_yaw(), self.get_pitch()
        obs = self.set(x, y, z, yaw, pitch, noise=0.7)
        self.set(x, y, z, yaw, pitch, noise=0)
        return obs

    def noisy(self, x, noise):
        return np.random.normal(x, noise)

    def set(self, x, y, z, yaw, pitch, noise=0.01):

        actions = [Teleport(self.noisy(x, noise), self.noisy(y, 0), self.noisy(z, noise)),
                   Yaw(self.noisy(yaw, noise)), Pitch(self.noisy(pitch, noise)), Turn(0.0001)]
        obs = None
        for action in actions:
            obs, _, _, _ = self.step(action)
            time.sleep(0.01)
        # time.sleep(0.2)
        return obs

    def set_pitch(self, pitch, noise=0.01):

        actions = [Pitch(self.noisy(pitch, noise)), Turn(0.0001)]
        obs = None
        for action in actions:
            obs, _, _, _ = self.step(action)
            time.sleep(0.01)
        # time.sleep(0.2)
        return obs

