import math
from typing import Dict, List, Literal

import jax
import jax.numpy as jnp
import numpy as np
from craftax.craftax.constants import *
from craftax.craftax.craftax_state import EnvState
from craftax.craftax_classic.renderer import render_craftax_pixels
from craftax.craftax_env import make_craftax_env_from_name

from craftaxlm.classic.metadata import (
    CRAFTAX_CLASSIC_ACHIEVEMENTS,
    CRAFTAX_CLASSIC_ACTION_MAPPING,
)
from craftaxlm.classic.state import (
    CraftaxClassicState,
    CraftaxRecorder,
    mob_id_to_name,
    player_chars,
    render_craftax_classic_text_custom,
)
from craftaxlm.shared import CraftaxBaseACI


class CraftaxClassicACI(CraftaxBaseACI):
    formatting: Literal["md", "xml"] = "md"
    map_format: Literal["full", "compact"] = "full"
    vantage: int = 3

    def __init__(
        self, formatting="md", map_format="full", vantage=3, verbose=False, seed=None
    ):
        self.formatting = formatting
        self.map_format = map_format
        self.vantage = vantage
        self.verbose = verbose
        self.achievement_deltas = []
        super().__init__(seed=seed)

    def make_env(self):
        return make_craftax_env_from_name(
            "Craftax-Classic-Symbolic-v1", auto_reset=False
        )

    def create_starting_obs(self):
        step_info = self.create_step_info(self.state, 0.0, False)
        return {
            "state_image": step_info["state_image"],
            "state_text": step_info["state_text"],
            "reward": 0.0,
            "done": False,
        }

    def map_action_string_to_int(self, action_string: str) -> int:
        return CRAFTAX_CLASSIC_ACTION_MAPPING.get(action_string.lower(), 0)

    def get_achievements(self, state):
        return {
            "achievements": {
                k: state.achievements[i]
                for i, k in CRAFTAX_CLASSIC_ACHIEVEMENTS.items()
            }
        }

    def get_achievement_delta(self, achievements):
        """
        Calculate which achievements have been newly unlocked.

        Args:
            achievements: Dictionary of current achievements

        Returns:
            Dictionary of newly unlocked achievements or None if no new achievements
        """
        if not hasattr(self, "last_achievements"):
            self.last_achievements = {
                "achievements": {k: 0 for k in achievements["achievements"].keys()}
            }
            return {}

        # Find new achievements (changed from 0 to 1)
        delta = {}
        for k, v in achievements["achievements"].items():
            if v > 0 and self.last_achievements["achievements"].get(k, 0) == 0:
                delta[k] = str(v)

        # Update last_achievements
        self.last_achievements = achievements

        # Return None if no new achievements
        if not delta:
            return {}

        return {"new_achievements": delta}

    def create_step_info(self, state, reward, done):
        craftaxlm_state = render_craftax_classic_text_custom(
            state, vantage=self.vantage
        )
        return {
            "state_image": craftaxlm_state.image,
            "state_text": craftaxlm_state.render_to_text_simple(
                verbose=self.verbose,
                formatting=self.formatting,
                map_format=self.map_format,
            ),
            "reward": float(reward),
            "done": bool(done),
        }

    def accept_tool_call(self, tool_call):
        """
        Accepts an OpenAI SDK-like tool call and executes the corresponding action.

        Args:
            tool_call: A dictionary with the structure similar to OpenAI tool calls:
                {
                    "type": "function",
                    "function": {
                        "name": "take_action",
                        "arguments": {
                            "action": "up" | "down" | "left" | "right" | etc.
                        }
                    },
                    "id": "call_abc123" # optional tool call ID
                }

        Returns:
            Dictionary formatted as an OpenAI tool message response:
            {
                "role": "tool",
                "content": JSON string with state information,
                "tool_call_id": The ID of the tool call being responded to
            }
        """
        state_before = self.state

        # Extract action from tool call
        action = None
        tool_call_id = None

        if isinstance(tool_call, str):
            # Handle legacy string action format
            action = tool_call
        elif isinstance(tool_call, dict):
            # Extract action from OpenAI SDK-like tool call format
            if "id" in tool_call:
                tool_call_id = tool_call["id"]

            if "function" in tool_call:
                function_data = tool_call["function"]
                if "arguments" in function_data:
                    arguments = function_data["arguments"]
                    # Handle both string and dict arguments
                    if isinstance(arguments, str):
                        import json

                        try:
                            arguments = json.loads(arguments)
                        except json.JSONDecodeError:
                            action = arguments

                    if isinstance(arguments, dict) and "action" in arguments:
                        action = arguments["action"]

        if action is None:
            raise ValueError(f"Could not extract action from tool call: {tool_call}")

        # Map action string to integer and execute
        action_int = self.map_action_string_to_int(action)
        _, state, reward, done, info = self.env.step(
            self.rngs[2], state_before, action_int, self.env_params
        )
        self.state = state

        # Calculate achievement deltas
        achievements = self.get_achievements(state)
        achievement_delta = self.get_achievement_delta(achievements)
        assert achievement_delta is not None
        # Store achievement deltas if the attribute exists
        # if hasattr(self, "achievement_deltas"):
        self.achievement_deltas.append(achievement_delta)

        # Print if verbose mode is enabled
        # if achievement_delta and hasattr(self, "verbose") and self.verbose:
        print(achievement_delta)

        step_info = self.create_step_info(state, reward, done)
        updates, expectation = get_action_result_vs_expectation(
            state_before, state, action
        )

        # Format response as an OpenAI tool message
        import base64
        import io
        import json

        import numpy as np
        from PIL import Image

        # Convert numpy array to base64 for OpenAI compatibility
        state_image = step_info.pop("state_image")  # Remove image from step_info

        # Convert numpy array to PIL Image and then to bytes
        pil_image = Image.fromarray(state_image)
        buffer = io.BytesIO()
        pil_image.save(buffer, format="PNG")

        # Get base64 encoded string
        img_str = base64.b64encode(buffer.getvalue()).decode("utf-8")

        # Format as OpenAI compatible image object
        image_obj = {
            "type": "image_url",
            "image_url": {"url": f"data:image/png;base64,{img_str}"},
        }

        response_content = {
            "direct_results": updates,
            "expected_results": expectation,
            "state": step_info,
            "state_image": image_obj,
        }

        # Include achievement delta if available
        if achievement_delta:
            response_content["achievement_delta"] = achievement_delta

        response = {
            "role": "tool",
            "content": json.dumps(response_content),
        }

        # Add tool_call_id if it was provided
        if tool_call_id:
            response["tool_call_id"] = tool_call_id

        return response

    def step_action(self, action_string, tool_call_id=None):
        """
        Helper method to perform an action using the OpenAI tool call format.

        Args:
            action_string: String representation of the action to take
            tool_call_id: Optional ID for the tool call

        Returns:
            Dictionary formatted as an OpenAI tool message response
        """
        # Create a tool call in OpenAI format
        tool_call = {
            "type": "function",
            "function": {"name": "take_action", "arguments": {"action": action_string}},
        }

        # Add tool_call_id if provided
        if tool_call_id:
            tool_call["id"] = tool_call_id

        # Process the tool call
        return self.accept_tool_call(tool_call)


def render_craftax_text(state):
    """Renders the CraftAX Classic state as text representation."""
    obs_dim_array = jnp.array([OBS_DIM[0], OBS_DIM[1]], dtype=jnp.int32)

    # Map
    padded_grid = jnp.pad(
        state.map,
        (MAX_OBS_DIM + 2, MAX_OBS_DIM + 2),
        constant_values=BlockType.OUT_OF_BOUNDS.value,
    )

    tl_corner = state.player_position - obs_dim_array // 2 + MAX_OBS_DIM + 2
    map_view = jax.lax.dynamic_slice(padded_grid, tl_corner, OBS_DIM)

    # Convert blocks to text representation
    block_chars = {
        BlockType.INVALID.value: ".",
        BlockType.OUT_OF_BOUNDS.value: "#",
        BlockType.GRASS.value: "G",
        BlockType.WATER.value: "~",
        BlockType.STONE.value: "S",
        BlockType.TREE.value: "T",
        BlockType.WOOD.value: "W",
        BlockType.PATH.value: "P",
        BlockType.COAL.value: "C",
        BlockType.IRON.value: "I",
        BlockType.DIAMOND.value: "D",
        BlockType.CRAFTING_TABLE.value: "B",  # B for crafting Bench
        BlockType.FURNACE.value: "F",
        BlockType.SAND.value: "_",
        BlockType.LAVA.value: "L",
        BlockType.PLANT.value: "p",
        BlockType.RIPE_PLANT.value: "r",
        BlockType.WALL.value: "X",
        BlockType.DARKNESS.value: " ",
        BlockType.WALL_MOSS.value: "M",
        BlockType.STALAGMITE.value: "^",
        BlockType.SAPPHIRE.value: "$",
        BlockType.RUBY.value: "R",
        BlockType.CHEST.value: "H",
        BlockType.FOUNTAIN.value: "O",
        BlockType.FIRE_GRASS.value: "f",
        BlockType.ICE_GRASS.value: "i",
        BlockType.GRAVEL.value: "g",
        BlockType.FIRE_TREE.value: "t",
        BlockType.ICE_SHRUB.value: "s",
        BlockType.ENCHANTMENT_TABLE_FIRE.value: "E",
        BlockType.ENCHANTMENT_TABLE_ICE.value: "e",
        BlockType.NECROMANCER.value: "N",
        BlockType.GRAVE.value: "v",
        BlockType.GRAVE2.value: "V",
        BlockType.GRAVE3.value: "u",
        BlockType.NECROMANCER_VULNERABLE.value: "n",
    }

    # Create text grid - swap x and y to match game coordinates
    height, width = map_view.shape
    text_grid = []
    for y in range(height):
        row = []
        for x in range(width):
            block_value = int(map_view[y, x])
            row.append(block_chars.get(block_value, "?"))
        text_grid.append(row)

    # Add player at center
    center_y, center_x = obs_dim_array // 2
    player_direction = int(state.player_direction)
    # print("Player direction: ", player_direction)
    text_grid[center_y][center_x] = player_chars.get(player_direction, "P")

    # print("Player char: ", text_grid[center_y][center_x])
    # Add mobs - ensure coordinates are properly aligned
    def add_mob_to_grid(mob_positions, mob_masks, symbol):
        for pos, mask in zip(mob_positions, mob_masks):
            if not mask:
                continue
            # Calculate local position relative to player
            local_pos = pos - state.player_position + obs_dim_array // 2
            # Check if mob is within visible area
            if (local_pos >= 0).all() and (local_pos < obs_dim_array).all():
                y, x = int(local_pos[0]), int(local_pos[1])
                text_grid[y][x] = symbol

    add_mob_to_grid(state.zombies.position, state.zombies.mask, "Z")
    add_mob_to_grid(state.cows.position, state.cows.mask, "c")
    add_mob_to_grid(state.skeletons.position, state.skeletons.mask, "K")
    add_mob_to_grid(state.arrows.position, state.arrows.mask, "a")

    # Convert to string - join rows in correct order
    text_map = "\n".join("".join(row) for row in text_grid)

    return text_map


compact_map_translation_full = {
    # Basic blocks
    "G": "Grass",
    "~": "Water",
    "S": "Stone",
    "T": "Tree",
    "W": "Wood",
    "P": "Path",
    "C": "Coal",
    "I": "Iron",
    "D": "Diamond",
    "B": "Crafting Bench",
    "F": "Furnace",
    "E": "Enchantment Table Fire",
    "e": "Enchantment Table Ice",
    # Additional blocks
    "_": "Sand",
    "L": "Lava",
    "p": "Plant",
    "r": "Ripe Plant",
    "X": "Wall",
    " ": "Darkness",
    "M": "Wall Moss",
    "^": "Stalagmite",
    "$": "Sapphire",
    "R": "Ruby",
    "H": "Chest",
    "O": "Fountain",
    # Special terrain
    "f": "Fire Grass",
    "i": "Ice Grass",
    "g": "Gravel",
    "t": "Fire Tree",
    "s": "Ice Shrub",
    # Necromancer and graves
    "N": "Necromancer",
    "n": "Necromancer Vulnerable",
    "v": "Grave",
    "V": "Grave2",
    "u": "Grave3",
    # Mobs
    "Z": "Zombie",
    "c": "Cow",
    "K": "Skeleton",
    "a": "Arrow",
    # Boundaries
    ".": "Invalid",
    "#": "Out of Bounds",
}

compact_map_translation_classic = {
    # Basic blocks
    "G": "Grass",
    "~": "Water",
    "S": "Stone",
    "T": "Tree",
    "W": "Wood",
    "P": "Path",
    "C": "Coal",
    "I": "Iron",
    "D": "Diamond",
    "B": "Crafting Bench",
    "F": "Furnace",
    "E": "Enchantment Table Fire",
    "e": "Enchantment Table Ice",
    # Additional blocks
    "_": "Sand",
    "L": "Lava",
    "p": "Plant",
    "r": "Ripe Plant",
    "X": "Wall",
    # Mobs
    "Z": "Zombie",
    "c": "Cow",
    "K": "Skeleton",
    "a": "Arrow",
    # Boundaries
    ".": "Invalid",
    "#": "Out of Bounds",
}


def render_craftax_classic_text_exp(
    state: EnvState, enable_recording=False, vantage=3
) -> CraftaxClassicState:
    print("Exp")
    map_data = []
    for x in range(state.map.shape[0]):
        for y in range(state.map.shape[1]):
            if (
                not max(
                    abs(x - state.player_position[0]), abs(y - state.player_position[1])
                )
                <= 4
            ):
                continue
            tile = {
                "position": {
                    "x": x - state.player_position[0],
                    "y": y - state.player_position[1],
                },
                "visible": max(
                    abs(x - state.player_position[0]), abs(y - state.player_position[1])
                )
                <= 4,
                "block": BlockType(state.map[x, y]).name.lower(),
            }
            if state.mob_map[x, y].max() > 0.5:
                tile["mob"] = mob_id_to_name(state.mob_map[x, y].argmax())
            map_data.append(tile)
    text_map = render_craftax_text(state)

    inventory_data = {
        "resources": {
            "wood": int(state.inventory.wood),
            "stone": int(state.inventory.stone),
            "coal": int(state.inventory.coal),
            "iron": int(state.inventory.iron),
            "diamond": int(state.inventory.diamond),
            "sapling": int(state.inventory.sapling),
        },
        "tools": {},
    }

    inventory_data["tools"]["pickaxe"] = {
        "wood": int(state.inventory.wood_pickaxe),
        "stone": int(state.inventory.stone_pickaxe),
        "iron": int(state.inventory.iron_pickaxe),
    }

    inventory_data["tools"]["sword"] = {
        "wood": int(state.inventory.wood_sword),
        "stone": int(state.inventory.stone_sword),
        "iron": int(state.inventory.iron_sword),
    }

    player_data = {
        "health": int(state.player_health),
        "food": int(state.player_food),
        "drink": int(state.player_drink),
        "energy": int(state.player_energy),
        "direction_facing": Action(state.player_direction).name.lower(),
    }

    environment_data = {
        "light_level": float(state.light_level),
        "is_sleeping": bool(state.is_sleeping),
        "floor": 0,  # Hardcode to 0 since player_level isn't in EnvState
    }

    def to_json_friendly(data):
        if isinstance(data, dict):
            return {key: to_json_friendly(value) for key, value in data.items()}
        elif isinstance(data, (list, tuple)):
            return [to_json_friendly(item) for item in data]
        elif isinstance(data, (jnp.ndarray, np.ndarray)):
            if data.ndim == 0:
                return data.item()
            return [to_json_friendly(item) for item in data]
        elif isinstance(data, (jnp.bool_, np.bool_)):
            return bool(data)
        elif isinstance(data, (jnp.integer, np.integer)):
            return int(data)
        elif isinstance(data, (jnp.floating, np.floating)):
            return float(data)
        else:
            return data

    pixels = render_craftax_pixels(state, BLOCK_PIXEL_SIZE_HUMAN)
    # Convert from JAX array to numpy and ensure uint8 format
    frame = np.array(pixels).astype(np.uint8)

    craftax_state = CraftaxClassicState(
        map_full=to_json_friendly(map_data),
        map_compact=text_map,
        inventory=to_json_friendly(inventory_data),
        player=to_json_friendly(player_data),
        environment=to_json_friendly(environment_data),
        image=frame,
        recorder=None if not enable_recording else CraftaxRecorder(),
        vantage=vantage,
    )

    # Record frame if recording is enabled
    if enable_recording:
        craftax_state.recorder.record_frame(state)

    return craftax_state


def canonical_crafter_score_classic(achievements_by_run: List[Dict]):
    assert all(
        isinstance(achievements, dict)
        and all(
            k in achievements
            for k in [
                "Collect Wood",
                "Place Table",
                "Eat Cow",
                "Collect Sapling",
                "Collect Drink",
                "Make Wood Pickaxe",
                "Make Wood Sword",
                "Place Plant",
                "Defeat Zombie",
                "Collect Stone",
                "Place Stone",
                "Eat Plant",
                "Defeat Skeleton",
                "Make Stone Pickaxe",
                "Make Stone Sword",
                "Wake Up",
                "Place Furnace",
                "Collect Coal",
                "Collect Iron",
                "Collect Diamond",
                "Make Iron Pickaxe",
                "Make Iron Sword",
            ]
        )
        and all(isinstance(v, int) for v in achievements.values())
        for achievements in achievements_by_run
    )
    successes_by_achievement = {
        achievement: sum(1 for run in achievements_by_run if run[achievement] > 0)
        for achievement in achievements_by_run[0].keys()
    }
    success_rates = [
        count / len(achievements_by_run) for count in successes_by_achievement.values()
    ]

    # Compute the Crafter score: S = exp((1/N)*sum(ln(1 + s_i))) - 1, where s_i is the success rate for achievement i.
    N = len(success_rates)
    return math.exp(sum(math.log(1 + s) for s in success_rates) / N) - 1


CRAFTAX_CLASSIC_ACTION_MAPPING = {
    "noop": 0,
    "left": 1,
    "right": 2,
    "up": 3,
    "down": 4,
    "do": 5,
    "sleep": 6,
    "place_stone": 7,
    "place_table": 8,
    "place_furnace": 9,
    "place_plant": 10,
    "make_wood_pickaxe": 11,
    "make_stone_pickaxe": 12,
    "make_iron_pickaxe": 13,
    "make_wood_sword": 14,
    "make_stone_sword": 15,
    "make_iron_sword": 16,
}


def get_action_expectations():
    """
    Returns a dictionary containing the required and optional expected results
    for each valid action in CraftAx Classic, using parametric inventory format.
    """
    expectations = {
        "noop": {"required": [], "optional": []},
        "left": {"required": [], "optional": ["turn_left"]},
        "right": {"required": [], "optional": ["turn_right"]},
        "up": {"required": [], "optional": ["turn_up"]},
        "down": {"required": [], "optional": ["turn_down"]},
        "do": {
            "required": [],
            "optional": [
                "stone_k->[k+1]",
                "wood_k->[k+1]",
                "coal_k->[k+1]",
                "iron_k->[k+1]",
                "diamond_k->[k+1]",
                "stone_removed",
                "tree_removed",
                "coal_removed",
                "iron_removed",
                "diamond_removed",
                "zombie_removed",
                "skeleton_removed",
            ],
        },
        "sleep": {"required": [], "optional": ["time_night->day", "health_low->full"]},
        "place_stone": {"required": ["stone_k->[k-1]", "stone_added"], "optional": []},
        "place_table": {
            "required": ["wood_k->[k-2]", "crafting_table_added"],
            "optional": [],
        },
        "place_furnace": {
            "required": ["stone_k->[k-1]", "furnace_added"],
            "optional": [],
        },
        "place_plant": {"required": [], "optional": []},
        "make_wood_pickaxe": {
            "required": ["wood_k->[k-1]", "wood_pickaxe_k->[k+1]"],
            "optional": [],
        },  # update these?
        "make_stone_pickaxe": {
            "required": ["stone_k->[k-1]", "wood_k->[k-1]", "stone_pickaxe_k->[k+1]"],
            "optional": [],
        },
        "make_iron_pickaxe": {
            "required": ["iron_k->[k-1]", "wood_k->[k-1]", "iron_pickaxe_k->[k+1]"],
            "optional": [],
        },
        "make_wood_sword": {
            "required": ["wood_k->[k-1]", "wood_sword_k->[k+1]"],
            "optional": [],
        },
        "make_stone_sword": {
            "required": ["stone_k->[k-1]", "wood_k->[k-1]", "stone_sword_k->[k+1]"],
            "optional": [],
        },
        "make_iron_sword": {
            "required": ["iron_k->[k-1]", "wood_k->[k-1]", "iron_sword_k->[k+1]"],
            "optional": [],
        },
    }

    return expectations


def get_action_expectation(action):
    """
    Returns the expected results for a specific action.

    Args:
        action: String name or integer ID of the action

    Returns:
        Dictionary with 'required' and 'optional' effects
    """
    expectations = get_action_expectations()

    # Handle case where action is passed as an integer ID
    if isinstance(action, int):
        for name, id in CRAFTAX_CLASSIC_ACTION_MAPPING.items():
            if id == action:
                action = name
                break

    if action in expectations:
        return expectations[action]
    else:
        return {"required": [], "optional": []}


def get_action_result(state_before, state_after):
    # updates = {}

    # Position
    updates = []
    player_position_before = state_before.player_position
    player_position_after = state_after.player_position
    # print("Player position before: ", player_position_before)
    # print("Player position after: ", player_position_after)
    if not all(player_position_before[i] == player_position_after[i] for i in range(2)):
        if player_position_before[0] < player_position_after[0]:
            updates.append("move_down")
        elif player_position_before[0] > player_position_after[0]:
            updates.append("move_up")
        elif player_position_before[1] < player_position_after[1]:
            updates.append("move_right")
        elif player_position_before[1] > player_position_after[1]:
            updates.append("move_left")
    else:
        pass
        # print("Player position is the same")

    # Direction
    direction_dict = {
        3: "up",
        4: "down",
        1: "left",
        2: "right",
    }
    player_direction_before = state_before.player_direction
    player_direction_after = state_after.player_direction
    if player_direction_before != player_direction_after:
        print(player_direction_after)
        updates.append("turn_" + direction_dict[int(player_direction_after)])

    # Inventory
    craftax_state_before = render_craftax_classic_text_custom(state_before, vantage=3)
    craftax_state_after = render_craftax_classic_text_custom(state_after, vantage=3)
    # print("Keys: ",craftax_state_before.inventory)
    # for rt in craftax_state_before.inventory.keys():
    for k in craftax_state_before.inventory["resources"].keys():
        if (
            craftax_state_before.inventory["resources"][k]
            != craftax_state_after.inventory["resources"][k]
        ):
            updates.append(
                k
                + "_"
                + str(craftax_state_before.inventory["resources"][k])
                + "->"
                + str(craftax_state_after.inventory["resources"][k])
            )
    for tool_type in craftax_state_before.inventory["tools"].keys():
        for tool_variety in craftax_state_before.inventory["tools"][tool_type].keys():
            if (
                craftax_state_before.inventory["tools"][tool_type][tool_variety]
                != craftax_state_after.inventory["tools"][tool_type][tool_variety]
            ):
                updates.append(
                    tool_variety
                    + "_"
                    + tool_type
                    + "_"
                    + str(
                        craftax_state_before.inventory["tools"][tool_type][tool_variety]
                    )
                    + "->"
                    + str(
                        craftax_state_after.inventory["tools"][tool_type][tool_variety]
                    )
                )
    # Map
    # If these are removed, show it: stone, tree, coal, iron, diamond, zombie, skeleton
    # If these are added, show it: crafting_table, furnace
    if all(player_position_before[i] == player_position_after[i] for i in range(2)):
        for tile in craftax_state_before.map_full:
            after_tile = next(
                (
                    t
                    for t in craftax_state_after.map_full
                    if t["position"]["x"] == tile["position"]["x"]
                    and t["position"]["y"] == tile["position"]["y"]
                ),
                None,
            )
            if tile["block"] in [
                "stone",
                "tree",
                "coal",
                "iron",
                "diamond",
                "zombie",
                "skeleton",
            ]:
                if tile != after_tile:
                    updates.append(tile["block"] + "_removed")
        for tile in craftax_state_after.map_full:
            before_tile = next(
                (
                    t
                    for t in craftax_state_before.map_full
                    if t["position"]["x"] == tile["position"]["x"]
                    and t["position"]["y"] == tile["position"]["y"]
                ),
                None,
            )
            if tile["block"] in ["crafting_table", "furnace"]:
                if tile != before_tile:
                    updates.append(tile["block"] + "_added")
    return updates


def get_action_result_vs_expectation(state_before, state_after, action):
    updates = get_action_result(state_before, state_after)
    expectation = get_action_expectation(action)
    return updates, expectation


def ensure_json_serializable(obj):
    """Convert NumPy arrays and other non-serializable objects to JSON-compatible types."""
    import numpy as np

    if isinstance(obj, dict):
        return {k: ensure_json_serializable(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [ensure_json_serializable(item) for item in obj]
    elif isinstance(obj, np.ndarray):
        # Convert NumPy array to Python native type
        return obj.item() if obj.size == 1 else obj.tolist()
    elif isinstance(obj, np.bool_):
        return bool(obj)
    elif hasattr(obj, "dtype") and np.isscalar(obj):
        # Handle other NumPy scalar types
        return obj.item()
    else:
        return obj


if __name__ == "__main__":
    import json

    craftax_aci = CraftaxClassicACI(
        formatting="md", map_format="full", verbose=True, seed=0
    )
    craftax_aci.reset()

    # Example 1: Using direct step calls (old way)
    print("EXAMPLE 1: Using direct step calls")
    state_before = craftax_aci.state
    action_int = craftax_aci.map_action_string_to_int("down")
    print("Action int: ", action_int)
    _, state, reward, done, info = craftax_aci.env.step(
        craftax_aci.rngs[2], state_before, action_int, craftax_aci.env_params
    )
    updates, expected_updates = get_action_result_vs_expectation(
        state_before, state, action_int
    )
    craftax_aci.state = state  # Make sure to update the state
    print("Updates:", updates)
    print("Expected:", expected_updates)

    # Example 2: Using the new step_action method
    print("\nEXAMPLE 2: Using step_action method")
    response = craftax_aci.step_action("down", tool_call_id="action_123")
    print("Tool Role:", response["role"])
    print("Tool Call ID:", response.get("tool_call_id"))

    content = json.loads(response["content"])
    print("Achievement Delta:", content.get("achievement_delta"))

    # Let's try to trigger an achievement
    print("\nTrying to trigger an achievement...")

    # Find a tree and try to collect wood (multiple steps to demonstrate)
    # First, look around for a tree
    directions = ["up", "right", "down", "left"]
    for direction in directions:
        print(f"\nChecking {direction} direction...")
        response = craftax_aci.step_action(direction)
        content = json.loads(response["content"])

        # See if we're facing a tree
        state_text = content["state"]["state_text"]
        if "Tree" in state_text and "Object_you_are_facing" in state_text:
            print(f"Found a tree! Attempting to collect wood...")
            # Use the 'do' action to collect wood
            response = craftax_aci.step_action("do")
            content = json.loads(response["content"])
            print("Action results:", content["direct_results"])

            # Check if we got an achievement
            if content.get("achievement_delta"):
                print("Achievement unlocked:", content["achievement_delta"])
            break

    # Example 3: Using accept_tool_call with OpenAI format
    print("\nEXAMPLE 3: Using accept_tool_call with OpenAI format")
    tool_call = {
        "type": "function",
        "function": {"name": "take_action", "arguments": {"action": "up"}},
        "id": "call_abc123",
    }
    response = craftax_aci.accept_tool_call(tool_call)
    print("Tool Role:", response["role"])
    print("Tool Call ID:", response.get("tool_call_id"))

    # Parse and display content for better readability
    import json

    content = json.loads(response["content"])
    print("\nResults:")
    print("Direct Results:", content["direct_results"])
    print("Expected Results:", content["expected_results"])
    print("Achievement Delta:", content.get("achievement_delta"))

    # Verify image format is correct
    image_data = content["state_image"]
    print("\nImage Data Format:")
    print(f"Type: {image_data['type']}")
    print(f"URL prefix: {image_data['image_url']['url'][:30]}...")

    # Optional: Save the image to a file to verify it works
    import base64

    try:
        # Extract base64 data
        base64_data = image_data["image_url"]["url"].split(",")[1]
        # Decode and save
        with open("test_output_image.png", "wb") as f:
            f.write(base64.b64decode(base64_data))
        print("\nImage saved to test_output_image.png")
    except Exception as e:
        print(f"Error saving image: {e}")

    # Show all accumulated achievement deltas
    print("\nAll Achievement Deltas:")
    print("Length of achievement deltas: ", len(craftax_aci.achievement_deltas))
    for i, delta in enumerate(craftax_aci.achievement_deltas):
        if delta:
            print(f"Step {i+1}:", delta)
    assert isinstance(craftax_aci.achievement_deltas, list)
    assert len(craftax_aci.achievement_deltas) > 0
    print("Achievement deltas: ", craftax_aci.achievement_deltas)
