import copy
import json
import logging
import os
import random
import shutil
import shortuuid
from typing import Dict

import hydra.core.hydra_config
from rich.logging import RichHandler


PICKAXE_TO_INT = {
    'wooden_pickaxe': 1,
    'stone_pickaxe': 2,
    'iron_pickaxe': 3,
}

INT_TO_PICKAXE = {
    1: 'wooden_pickaxe',
    2: 'stone_pickaxe',
    3: 'iron_pickaxe',
}


def is_belief_correct(believed_dependency, action_cond_list):
    correct_flag = False
    for cond_dict in action_cond_list:
        type = cond_dict['type']

        precondition = cond_dict['precondition']
        tool = cond_dict['tool']

        merged_dict = {**precondition, **tool}
        if type == 'smelt':
            if 'coal' in believed_dependency:
                del believed_dependency['coal']
            if 'coals' in believed_dependency:
                del believed_dependency['coals']
            if 'coal' in merged_dict:
                del merged_dict['coal']
            if 'coals' in merged_dict:
                del merged_dict['coals']

        if 'coal' in merged_dict:
            merged_dict['coals'] = merged_dict['coal']
            del merged_dict['coal']
        if 'oak_log' in merged_dict:
            merged_dict['logs'] = merged_dict['oak_log']
            del merged_dict['oak_log']
        if 'oak_planks' in merged_dict:
            merged_dict['planks'] = merged_dict['oak_planks']
            del merged_dict['oak_planks']

        if merged_dict == believed_dependency:
            correct_flag = True
            break

    return correct_flag


def create_new_action_lib_json(env, action_lib):

    perturb_requirements = False
    perturb_operations = False

    perturb_iron_tool = False # "craft_iron_sword", "craft_iron_pickaxe"
    perturb_stone_tool = False # "craft_furnace", "craft_stone_axe", "craft_stone_pickaxe"
    perturb_wooden_tool = False # "craft_wooden_axe", "craft_wooden_pickaxe"

    if 'op' in action_lib:
        perturb_operations = True
    elif 'req' in action_lib:
        perturb_requirements = True
    elif 'both' in action_lib:
        perturb_operations = True
        perturb_requirements = True
    else:
        raise ValueError("Invalid action_lib name!")

    if 'level1' in action_lib:
        perturb_iron_tool = True
    elif 'level2' in action_lib:
        perturb_iron_tool = True
        perturb_stone_tool = True
    elif 'level3' in action_lib:
        perturb_iron_tool = True
        perturb_stone_tool = True
        perturb_wooden_tool = True

    # copy json file of self.perturbed_action_lib_base_path into self.perturbed_action_lib_dst_path
    shutil.copy(env.perturbed_action_lib_src_path, env.perturbed_action_lib_dst_path)
    with open(env.perturbed_action_lib_dst_path, 'r') as f:
        prev_action_lib_data = json.load(f)

    new_add = prev_action_lib_data["add"]
    if perturb_requirements:
        if perturb_iron_tool:
            replaced_req_iron_ingot = random.choice(['iron_nugget', 'cobblestone'])
            new_add["craft_iron_sword"] = [{"output": {"iron_sword": 1}, "type": "craft", "precondition": {"stick": 1, replaced_req_iron_ingot: 2}, "tool": {"crafting_table": 1}, "abbr": "craft_iron_sword"}]
            new_add["craft_iron_pickaxe"] = [{"output": {"iron_pickaxe": 1}, "type": "craft", "precondition": {"stick": 2, replaced_req_iron_ingot: 3}, "tool": {"crafting_table": 1}, "abbr": "craft_iron_pickaxe"}]
        if perturb_stone_tool:
            replaced_req_cobblestone = random.choice(['oak_planks', 'oak_log'])
            new_add["craft_furnace"] = [{"output": {"furnace": 1}, "type": "craft", "precondition": {replaced_req_cobblestone: 8}, "tool": {"crafting_table": 1}, "abbr": "craft_furnace"}]
            new_add["craft_stone_axe"] = [{"output": {"stone_axe": 1}, "type": "craft", "precondition": {"stick": 2, replaced_req_cobblestone: 3}, "tool": {"crafting_table": 1}, "abbr": "craft_stone_axe"}]
            new_add["craft_stone_pickaxe"] = [{"output": {"stone_pickaxe": 1}, "type": "craft", "precondition": {"stick": 2, replaced_req_cobblestone: 3}, "tool": {"crafting_table": 1}, "abbr": "craft_stone_pickaxe"}]
        if perturb_wooden_tool:
            replaced_req_oak_planks = 'oak_log'
            new_add["craft_wooden_axe"] = [{"output": {"wooden_axe": 1}, "type": "craft", "precondition": {"stick": 2, replaced_req_oak_planks: 3}, "tool": {"crafting_table": 1}, "abbr": "craft_wooden_axe"}]
            new_add["craft_wooden_pickaxe"] = [{"output": {"wooden_pickaxe": 1}, "type": "craft", "precondition": {"stick": 2, replaced_req_oak_planks: 3}, "tool": {"crafting_table": 1}, "abbr": "craft_wooden_pickaxe"}]

    if perturb_operations:
        if perturb_iron_tool:
            replaced_op_craft = random.choice(['smelt', 'mine'])

            updated_iron_sword = copy.deepcopy(new_add["craft_iron_sword"][0])
            updated_iron_sword['type'] = replaced_op_craft
            updated_iron_sword['abbr'] = f"{replaced_op_craft}_iron_sword"
            new_add[f"{replaced_op_craft}_iron_sword"] = [updated_iron_sword]
            del new_add["craft_iron_sword"]

            updated_iron_pickaxe = copy.deepcopy(new_add["craft_iron_pickaxe"][0])
            updated_iron_pickaxe['type'] = replaced_op_craft
            updated_iron_pickaxe['abbr'] = f"{replaced_op_craft}_iron_pickaxe"
            new_add[f"{replaced_op_craft}_iron_pickaxe"] = [updated_iron_pickaxe]
            del new_add["craft_iron_pickaxe"]

        if perturb_stone_tool:
            replaced_op_craft = random.choice(['smelt', 'mine'])

            updated_furnace = copy.deepcopy(new_add["craft_furnace"][0])
            updated_furnace['type'] = replaced_op_craft
            updated_furnace['abbr'] = f"{replaced_op_craft}_furnace"
            new_add[f"{replaced_op_craft}_furnace"] = [updated_furnace]
            del new_add["craft_furnace"]

            updated_stone_axe = copy.deepcopy(new_add["craft_stone_axe"][0])
            updated_stone_axe['type'] = replaced_op_craft
            updated_stone_axe['abbr'] = f"{replaced_op_craft}_stone_axe"
            new_add[f"{replaced_op_craft}_stone_axe"] = [updated_stone_axe]
            del new_add["craft_stone_axe"]

            updated_stone_pickaxe = copy.deepcopy(new_add["craft_stone_pickaxe"][0])
            updated_stone_pickaxe['type'] = replaced_op_craft
            updated_stone_pickaxe['abbr'] = f"{replaced_op_craft}_stone_pickaxe"
            new_add[f"{replaced_op_craft}_stone_pickaxe"] = [updated_stone_pickaxe]
            del new_add["craft_stone_pickaxe"]

        if perturb_wooden_tool:
            replaced_op_craft = random.choice(['smelt', 'mine'])
            updated_wooden_axe = copy.deepcopy(new_add["craft_wooden_axe"][0])
            updated_wooden_axe['type'] = replaced_op_craft
            updated_wooden_axe['abbr'] = f"{replaced_op_craft}_wooden_axe"
            new_add[f"{replaced_op_craft}_wooden_axe"] = [updated_wooden_axe]
            del new_add["craft_wooden_axe"]

            updated_wooden_pickaxe = copy.deepcopy(new_add["craft_wooden_pickaxe"][0])
            updated_wooden_pickaxe['type'] = replaced_op_craft
            updated_wooden_pickaxe['abbr'] = f"{replaced_op_craft}_wooden_pickaxe"
            new_add[f"{replaced_op_craft}_wooden_pickaxe"] = [updated_wooden_pickaxe]
            del new_add["craft_wooden_pickaxe"]

    new_action_lib_data = copy.deepcopy(prev_action_lib_data)
    new_action_lib_data["add"] = new_add

    with open(env.perturbed_action_lib_dst_path, "w") as f:
        json.dump(new_action_lib_data, f, indent=2)

    return f'action_lib_new_perturbed_{env.uuid}'



def get_action_lib_patch(action_lib_patch_name):
    basepath = os.path.abspath(__file__)
    folder = os.path.dirname(basepath)
    action_lib_patch_file = f"{action_lib_patch_name}.json" 
    action_lib_patch_path = os.path.join(folder, "action_libs", action_lib_patch_file)

    with open(action_lib_patch_path, 'r') as f:
        action_lib_patch = json.load(f)

    return action_lib_patch


def max_pickaxe_level(inventory: Dict[str, int]) -> int:
    pickaxes_levels = [PICKAXE_TO_INT[item] for item in inventory if item in PICKAXE_TO_INT]
    pickaxes_levels.append(0)

    return max(pickaxes_levels)

def change_textworld_item_name(item_name: str) -> str:
    if "log" in item_name:
        return "logs"
    elif "planks" in item_name:
        return "planks"
    elif item_name == "coal" or item_name == "coals":
        return "coals"
    else:
        return item_name

def language_action_to_subgoal(action, waypoint) -> str:
    subgoal = {
        "task": action,
        "goal": [waypoint, 1],
    }
    subgoal_str = json.dumps(subgoal)

    return subgoal, subgoal_str

def change_log_planks_coal_recipe(_recipe_data):

    item_name = _recipe_data['item_name']
    if "log" in _recipe_data['item_name']:
        item_name = "logs"
    elif "planks" in _recipe_data['item_name']:
        item_name = "planks"
    elif _recipe_data['item_name'] == "coal" or _recipe_data['item_name'] == "coals":
        item_name = "coals"
    else:
        item_name = _recipe_data['item_name']

    ingredients = {}
    for k, v in _recipe_data['ingredients'].items():
        if "log" in k:
            ingredients["logs"] = v
        elif "planks" in k:
            ingredients["planks"] = v
        elif k == "coal" or k == "coals":
            ingredients["coals"] = v
        else:
            ingredients[k] = v

    recipe_data = {
        "item_name": item_name,
        "output_qty": _recipe_data['output_qty'],
        "ingredients": ingredients,
        "required_pickaxe": _recipe_data['required_pickaxe'],
        "is_crafting": _recipe_data['is_crafting'],
    }

    return recipe_data


def get_logger(name) -> logging.Logger:
    hydra_path = hydra.core.hydra_config.HydraConfig.get().runtime.output_dir
    job_name = hydra.core.hydra_config.HydraConfig.get().job.name

    file_handler = logging.FileHandler(os.path.join(hydra_path, f"{job_name}.log"))
    rich_handler = RichHandler(markup=True)

    logging.basicConfig(format="[%(asctime)s]: %(message)s", datefmt="%Y-%m-%d %H:%M:%S")
    log = logging.getLogger(name)
    log.setLevel(logging.INFO)
    log.addHandler(rich_handler)
    log.addHandler(file_handler)
    log.propagate = False

    return log


# check if all items in inventory
def check_dict(inventory:dict, needed_items:dict):
    for i in needed_items.keys():
        if i in inventory.keys(): # check item exist
            if needed_items[i] > inventory[i]: # check item enough
                return False
        else:
            return False
    return True

# add group to action
def add_group_to_action(old_path, path):
    with open(old_path, 'r') as j:
        old_action_lib = json.load(j)
    new_action_lib = {}
    for action in old_action_lib.keys():
        new_action_lib[action] = []
        for cand_action in old_action_lib[action]:
            cand_action['group'] = action.split("_")[0]
            new_action_lib[action].append(cand_action)
    with open(path, 'w') as j:
        json.dump(new_action_lib, j)
        
def create_task_json(old_path, path):
    with open(old_path, 'r') as j:
        old_task_lib = json.load(j)
    new_task_lib = []
    for task in old_task_lib:
        task["description"] = "Obtain a " + task["task"].replace("_", " ")
        new_task_lib.append(task)
    with open(path, 'w') as j:
        json.dump(old_task_lib, j)

def add_mine_coal_to_plans(old_path, path):
    with open(old_path, "r") as j:
        old_plans = json.load(j)
    new_plans = {}
    for key in old_plans:
        actions = old_plans[key]
        new_plans[key]=[]
        for action in actions:
            if(action["type"]=="smelt"):
                sum_goal = 0
                for item in action["goal"]:
                    sum_goal += action["goal"][item]
                smelt_action ={
                    "goal": {
                        "coal": 1
                    },
                    "type": "mine",
                    "text": "coal"
                }
                for i in range(sum_goal):
                    new_plans[key].append(smelt_action)
            new_plans[key].append(action)
    with open(path, "w") as j:
        json.dump(new_plans, j)
# create_task_json("jarvis_task.json", "tasks.json")
# add_mine_coal_to_plans("jarvis_plans.json", "plans.json")



