from typing import Literal
from .transforms import _tolist
import numpy as np
import sapien


class ArmTag:
    _instances = {}

    def __new__(cls, value):
        if isinstance(value, ArmTag):
            return value
        if isinstance(value, str) and value in ["left", "right"]:
            value = value.lower()
            if value in cls._instances:
                return cls._instances[value]
            instance = super().__new__(cls)
            cls._instances[value] = instance
            return instance
        raise ValueError(f"Invalid arm tag: {value}. Must be 'left' or 'right'.")

    def __init__(self, value):
        if isinstance(value, str):
            self.arm = value.lower()

    @property
    def opposite(self):
        return ArmTag("right") if self.arm == "left" else ArmTag("left")

    def __eq__(self, other):
        if isinstance(other, ArmTag):
            return self.arm == other.arm
        if isinstance(other, str):
            return self.arm == other.lower()
        return False

    def __hash__(self):
        return hash(self.arm)

    def __repr__(self):
        return f"ArmTag('{self.arm}')"

    def __str__(self):
        return self.arm


class Action:
    arm_tag: ArmTag
    action: Literal["move", "gripper"]
    target_pose: list = None
    target_gripper_pos: float = None

    def __init__(
        self,
        arm_tag: ArmTag | Literal["left", "right"],
        action: Literal["move", "open", "close", "gripper"],
        target_pose: sapien.Pose | list | np.ndarray = None,
        target_gripper_pos: float = None,
        **args,
    ):
        self.arm_tag = ArmTag(arm_tag)
        if action != "move":
            if action == "open":
                self.action = "gripper"
                self.target_gripper_pos = (target_gripper_pos if target_gripper_pos is not None else 1.0)
            elif action == "close":
                self.action = "gripper"
                self.target_gripper_pos = (target_gripper_pos if target_gripper_pos is not None else 0.0)
            elif action == "gripper":
                self.action = "gripper"
            else:
                raise ValueError(f"Invalid action: {action}. Must be 'open' or 'close'.")
            assert (self.target_gripper_pos is not None), "target_gripper_pos cannot be None for gripper action."
        else:
            self.action = "move"
            assert (target_pose is not None), "target_pose cannot be None for move action."
            self.target_pose = _tolist(target_pose)
        self.args = args

    def __str__(self):
        result = f"{self.arm_tag}: {self.action}"
        if self.action == "move":
            result += f"({self.target_pose})"
        else:
            result += f"({self.target_gripper_pos})"
        if self.args:
            result += f"    {self.args}"
        return result
