import json
import os
import copy
import shortuuid

import gym
from gym import spaces
from gym.utils import seeding

from mctextworld.action import ActionLibrary
from mctextworld.utils import *

class Env(gym.Env):
    '''
    Minecraft Simulator: This is a simulator for test the online planner in Minecraft.
    '''
    def __init__(self,task_name = "Obtain 1 diamond", init_inv = {}, task_obj = {}, biome = "plains", maximum_step = 100, action_lib = "default"):
        self.task_name = task_name
        self.init_inv = init_inv
        self.legal_actions = ["mine", "craft", "smelt"]
        self.task_obj = task_obj
        self.reward = 0
        # self.goal_lib = self.load_goal_lib()
        self.action_lib = ActionLibrary(action_lib)
        self.obs = {
            "inventory": copy.deepcopy(self.init_inv),
            "position": [0,0,0],
            "biome": biome,
        }
        self.maximum_step = maximum_step
        self.reset()

        basepath = os.path.abspath(__file__)
        folder = os.path.dirname(basepath)
        self.uuid = str(shortuuid.uuid())[:6]
        self.perturbed_action_lib_src_path = os.path.join(folder, "action_libs", "action_lib_new_base.json")
        self.perturbed_action_lib_dst_path = os.path.join(folder, "action_libs", f"action_lib_new_perturbed_{self.uuid}.json")


    def set_action_lib(self, action_lib):
        '''
        Set the action library.
        '''
        self.action_lib = ActionLibrary(action_lib)
        self.all_actions = self.action_lib.all_actions
        print(f"Action Library: {len(self.all_actions)}")

    def print_obs(self):
        print("CURRENT STATE:")
        print("-----------------------------")
        print("Inventory: ", self.obs['inventory'])
        # print("position: ", self.obs['position'])
        # print("biome: ", self.obs['biome'])
        # print("Reward: ", self.reward)
        # print("Done: ", self.done)
        print("Info: ", self.info)
        print("-----------------------------")

    def reset(self, options=None):        
        init_inv = {}
        if options is not None and 'inventory' in options:
            init_inv = options['inventory']

        self.obs = {
            "inventory": copy.deepcopy(init_inv),
            "position": [0,0,0],
            "biome": "forest",
        }
        self.reward = 0
        self.done = False
        self.info = {}

        self.curr_step = 0
        self.executed_actions = []

        # return self.obs # observation?
        return copy.deepcopy(self.obs), self.reward, self.done, copy.deepcopy(self.info)

    def close(self):
        self.reset()

    def check_done(self, inventory):
        if self.task_obj is None:
            return False
        if check_dict(inventory, self.task_obj):
            return True
        else:
            return False

    def check_inventory(self, inventory):
        new_inventory = {}
        for key, val in inventory.items():
            if val < 0:
                raise Exception('Invalid inventory!')
            elif val > 0:
                new_inventory[key] = val
            else:
                pass
        return new_inventory

    def check_max_step(self):
        if self.curr_step > self.maximum_step:
            return True
        else:
            return False

    def step(self, action: str):
        '''
        return next obs, reward, done, info
        '''
        # check maximum step
        reach_maximum_step = self.check_max_step()
        if reach_maximum_step:
            return self.obs, self.reward, self.done, {
            'action_success': False, 
            'reach_maximum_step': reach_maximum_step
            }
        # Check action validation
        act, msg = self.action_lib.check_action(self.obs['inventory'], action)
        self.curr_step += 1

        self.executed_actions.append(act)

        inventory = self.check_inventory(inventory = copy.deepcopy(self.obs['inventory']))
        self.obs['inventory'] = inventory

        if not act:
            self.info = {'action_success': False, 'message': msg}
            return copy.deepcopy(self.obs), self.reward, self.done, copy.deepcopy(self.info)

        output = act['output'] # effect
        precondition = act['precondition'] # precondition
        tool = act['tool']

        # delete precondition items
        for key, val in precondition.items():
            self.obs['inventory'][key] -= val

        # add output items
        for key, val in output.items():
            if key in self.obs['inventory'].keys():
                self.obs['inventory'][key] += val
            else:
                self.obs['inventory'][key] = val

        self.done = self.check_done(inventory = self.obs['inventory'])
        inventory = self.check_inventory(inventory = copy.deepcopy(self.obs['inventory']))
        self.obs['inventory'] = inventory

        self.info = {'action_success': True}

        output_item_name = list(output.keys())[0]
        output_qty = output[output_item_name]

        ingredients = copy.deepcopy(act['precondition'])
        ingredients.update(tool)

        required_pickaxe = max_pickaxe_level(tool)
        if act['type'] == 'mine':
            required_pickaxe = max_pickaxe_level(ingredients)
        is_crafting = act['type'] == 'craft'

        self.info.update({
            "item_name": output_item_name,
            "output_qty": output_qty,
            "ingredients": ingredients,
            "required_pickaxe": required_pickaxe,
            "is_crafting": is_crafting,
        })

        return copy.deepcopy(self.obs), self.reward, self.done, copy.deepcopy(self.info)


if __name__ == '__main__':
    env = Env(
        init_inv = {}, 
        task_obj = {'diamond': 1, 'log': 3}
        )
