import sys

sys.path.append("./")

import sapien.core as sapien
from collections import OrderedDict
import pdb
from envs import *
import yaml
import importlib
import json
import traceback
import os
import time
import inspect

current_file_path = os.path.abspath(__file__)
parent_directory = os.path.dirname(current_file_path)

SCRIPT_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "script")
CONFIGS_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "task_config")
OBJECTS_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "assets/objects")


def enrich_actors(actor_list):
    """
    Enrich the actor list by adding 'functional_points' and 'contact_points'
    from the corresponding model metadata file, and remove the 'modelname' field
    to make it suitable for prompting.

    Args:
        actor_list (dict): Dictionary of actors with metadata.

    Returns:
        dict: Enriched actor list with additional keys and without 'modelname'.
    """
    enriched_actor_list = {}

    for actor_key, actor_info in actor_list.items():
        enriched_actor = actor_info.copy()
        model_name = actor_info.get("modelname")

        if model_name is not None and model_name != "None":
            points_info_path = os.path.join(OBJECTS_PATH, model_name, "points_info.json")

            if os.path.exists(points_info_path):
                try:
                    with open(points_info_path, 'r') as f:
                        points_info = json.load(f)

                    if "functional_points" in points_info:
                        enriched_actor["functional_points"] = points_info["functional_points"]

                    if "contact_points" in points_info:
                        contact_points = points_info["contact_points"]
                        valid_contact_points = any(
                            point.get("id") is not None and (
                                not isinstance(point.get("id"), list) or len(point.get("id")) > 0
                            ) for point in contact_points
                        )
                        enriched_actor["contact_points"] = contact_points if valid_contact_points else None
                    else:
                        enriched_actor["contact_points"] = None

                except Exception as e:
                    print(f"Error reading points_info.json for {model_name}: {e}")
                    print(traceback.format_exc())
            else:
                print(f"Warning: File not found: {points_info_path}")
        else:
            print("modelname is None or invalid, skipping enrichment.")

        if "modelname" in enriched_actor:
            del enriched_actor["modelname"]

        enriched_actor_list[actor_key] = enriched_actor

    return enriched_actor_list


def class_decorator_gen(task_name):
    """
    Dynamically import and instantiate the task implementation from the code_gen module.

    Args:
        task_name (str): Name of the task.

    Returns:
        object: Instance of the task class.
    """
    envs_module = importlib.import_module(f"envs_gen.gpt_{task_name}")
    importlib.reload(envs_module)
    try:
        env_class = getattr(envs_module, f"gpt_{task_name}")
        return env_class()
    except:
        raise SystemExit("No such task")


def class_decorator_env(task_name):
    """
    Dynamically import and instantiate the task environment from the envs module.

    Args:
        task_name (str): Name of the task.

    Returns:
        object: Instance of the task class.
    """
    envs_module = importlib.import_module(f"envs.{task_name}")
    importlib.reload(envs_module)
    try:
        env_class = getattr(envs_module, task_name)
        return env_class()
    except:
        raise SystemExit("No such task")


def create_task_config(task_config_path, task_name):
    """
    Create a new task config file from the template if it doesn't exist.

    Args:
        task_config_path (str): Path to the target config file.
        task_name (str): Name of the task.
    """
    with open(os.path.join(SCRIPT_PATH, "_task_config_template.json"), "r") as file:
        task_config_template = json.load(file)

    # Modify task_name
    task_config_template["task_name"] = task_name

    # Convert field format
    if isinstance(task_config_template.get("embodiment"), str):
        task_config_template["embodiment"] = [task_config_template["embodiment"]]

    # Save as yml
    with open(task_config_path, "w") as f:
        yaml.dump(task_config_template, f, default_flow_style=False, sort_keys=False)


def get_embodiment_config(robot_file):
    """
    Load embodiment configuration from the robot folder.

    Args:
        robot_file (str): Path to the robot folder.

    Returns:
        dict: Robot configuration.
    """
    robot_config_file = os.path.join(robot_file, "config.yml")
    with open(robot_config_file, "r", encoding="utf-8") as f:
        return yaml.load(f.read(), Loader=yaml.FullLoader)


def setup_task_config(task_name):
    """
    Load or create a task configuration and set up robot embodiments.

    Args:
        task_name (str): Task name.

    Returns:
        tuple: (Task instance, task configuration dictionary)
    """
    task = class_decorator_gen(task_name)
    task_config_path = f"./task_config/{task_name}.yml"

    if not os.path.isfile(task_config_path):
        create_task_config(task_config_path, task_name)
        print(f"Task config file is missing, please check {task_config_path}")

    with open(task_config_path, "r", encoding="utf-8") as f:
        args = yaml.load(f.read(), Loader=yaml.FullLoader)

    args["domain_randomization"] = {
        "random_background": False,
        "cluttered_table": False,
        "clean_background_rate": 0.0,
        "random_head_camera_dis": 0,
        "random_table_height": 0.0,
        "random_light": False,
        "crazy_random_light_rate": 0.0,
        "random_embodiment": False,
    }

    embodiment_type = args.get("embodiment")
    embodiment_config_path = os.path.join("./task_config", "_embodiment_config.yml")
    with open(embodiment_config_path, "r", encoding="utf-8") as f:
        _embodiment_types = yaml.load(f.read(), Loader=yaml.FullLoader)

    def get_embodiment_file(embodiment_type):
        robot_file = _embodiment_types[embodiment_type]["file_path"]
        if robot_file is None:
            raise Exception("No embodiment files")
        return robot_file if os.path.isabs(robot_file) else os.path.abspath(
            os.path.join(os.path.dirname(__file__), "..", robot_file)
        )

    if len(embodiment_type) == 1:
        args["left_robot_file"] = get_embodiment_file(embodiment_type[0])
        args["right_robot_file"] = get_embodiment_file(embodiment_type[0])
        args["dual_arm_embodied"] = True
    elif len(embodiment_type) == 3:
        args["left_robot_file"] = get_embodiment_file(embodiment_type[0])
        args["right_robot_file"] = get_embodiment_file(embodiment_type[1])
        args["embodiment_dis"] = embodiment_type[2]
        args["dual_arm_embodied"] = False
    else:
        raise Exception("Embodiment items should be 1 or 3")

    args["left_embodiment_config"] = get_embodiment_config(args["left_robot_file"])
    args["right_embodiment_config"] = get_embodiment_config(args["right_robot_file"])

    args["embodiment_name"] = (
        str(embodiment_type[0]) if len(embodiment_type) == 1
        else str(embodiment_type[0]) + "+" + str(embodiment_type[1])
    )

    args["need_plan"] = True
    args["save_path"] = "./data/test"

    return task, args


def run(TASK_ENV, args, check_num=10):
    """
    Run the task in simulation to evaluate success rate.

    Args:
        TASK_ENV (object): Task environment instance.
        args (dict): Task configuration.
        check_num (int): Number of trials to run.

    Returns:
        tuple: (success rate, most common error message, error count, run records)
    """
    epid, suc_num, fail_num = 0, 0, 0

    error_list = [
        "The code can not run", "The left arm failed to grasp the object", "The right arm failed to grasp the object",
        "The target position of the object is incorrect.", "Plan execution failed",
        "Unknown error occurred during execution"
    ]
    error_num = [0, 0, 0, 0, 0, 0]
    run_records = []

    print(f"\033[34mTask name: {args['task_name']}\033[0m")
    print("\033[93m" + "[Start Testing Task Success Rate]" + "\033[0m")

    print("\n\033[92m=== play_once source code ===\033[0m")
    play_once_method = TASK_ENV.__class__.play_once
    print(inspect.getsource(play_once_method))
    print("\033[92m=== End ===\033[0m\n")

    for epid in range(check_num):
        error_id = None
        try:
            TASK_ENV.setup_demo(now_ep_num=suc_num, seed=epid, **args)
            TASK_ENV.play_once()

            if TASK_ENV.plan_success and TASK_ENV.check_success():
                print(f"simulate data episode {suc_num} success! (seed = {epid})")
                suc_num += 1
                run_records.append("success!")
            else:
                if not TASK_ENV.plan_success:
                    if hasattr(TASK_ENV, 'left_plan_success') and not TASK_ENV.lefft_plan_success:
                        error_id = 1
                        run_records.append(error_list[1])
                    elif hasattr(TASK_ENV, 'right_plan_success') and not TASK_ENV.right_plan_success:
                        error_id = 2
                        run_records.append(error_list[2])
                    else:
                        error_id = 4
                        run_records.append(error_list[4])
                else:
                    error_id = 3
                    run_records.append(error_list[3])

                print(f"simulate data episode {suc_num} fail! (seed = {epid})")
                fail_num += 1

            TASK_ENV.close()
            if args.get("render_freq"):
                TASK_ENV.viewer.close()

        except Exception as e:
            error_id = 0
            error_list[0] = str(traceback.format_exc())
            run_records.append(f"Error: {e}")
            print("-------------")
            print(f"simulate data episode {suc_num} fail! (seed = {epid})")
            print("Error:", traceback.format_exc())
            print("-------------")
            fail_num += 1
            TASK_ENV.close()
            if args.get("render_freq"):
                TASK_ENV.viewer.close()
            time.sleep(2)

        if error_id is not None:
            error_num[error_id] += 1

    if len(run_records) != check_num:
        print(f"Warning: number of records ({len(run_records)}) does not match number of trials ({check_num})")

    max_error_index = error_num.index(max(error_num)) if sum(error_num) > 0 else 5
    max_error_count = error_num[max_error_index]

    print(f'\nComplete test, success rate: {suc_num}/{check_num}')
    print(f'Error message: {error_list}')
    print(f'Run records: {run_records}')
    print(f'error_num: {error_num}')

    return suc_num / check_num, error_list[max_error_index], max_error_count, run_records
