import json
import time

# ===== 关节列表 =====

set_joint_list = [

    {'joint_name': 'R_WRIST_Y', 'rad': 0.066444339857},
    {'joint_name': 'R_WRIST_P', 'rad': -20.531186157241},
    {'joint_name': 'R_ELBOW_R', 'rad': -1.340070808067},
    {'joint_name': 'R_SHOULDER_Y', 'rad': -1.952500297438},
    {'joint_name': 'R_SHOULDER_R', 'rad': -0.626298915913},
    {'joint_name': 'R_SHOULDER_P', 'rad': 0.342241935176},
]

drive_joint_list = [

    {'joint_name': 'R_WRIST_Y', 'rad': -0.011299818873},
    {'joint_name': 'R_WRIST_P', 'rad': -20.60070130641},
    {'joint_name': 'R_ELBOW_R', 'rad': -1.364414209864},
    {'joint_name': 'R_SHOULDER_Y', 'rad': -1.980264780151},
    {'joint_name': 'R_SHOULDER_R', 'rad': -0.560519997177},
    {'joint_name': 'R_SHOULDER_P', 'rad': 0.442999602707},
]

mode_joint_list = [

    {'joint_name': 'R_WRIST_Y', 'rad': 0.274796141277},
    {'joint_name': 'R_WRIST_P', 'rad': -20.944670868464},
    {'joint_name': 'R_ELBOW_R', 'rad': -1.136851577265},
    {'joint_name': 'R_SHOULDER_Y', 'rad': -1.543831108824},
    {'joint_name': 'R_SHOULDER_R', 'rad': -0.465282192444},
    {'joint_name': 'R_SHOULDER_P', 'rad': 0.415577797633},
]

home_joint_list = [
    {'joint_name': 'R_WRIST_Y', 'rad': 0.031161140931},
    {'joint_name': 'R_WRIST_P', 'rad': -20.577257789382},
    {'joint_name': 'R_ELBOW_R', 'rad': -1.381942977952},
    {'joint_name': 'R_SHOULDER_Y', 'rad': -2.055649114477},
    {'joint_name': 'R_SHOULDER_R', 'rad': -0.602941780228},
    {'joint_name': 'R_SHOULDER_P', 'rad': 0.314926445602},
]

start_joint_list = [

    {'joint_name': 'R_WRIST_Y', 'rad': 0.274967005474},
    {'joint_name': 'R_WRIST_P', 'rad': -20.603120933284},
    {'joint_name': 'R_ELBOW_R', 'rad': -1.499983558976},
    {'joint_name': 'R_SHOULDER_Y', 'rad': -1.492936360109},
    {'joint_name': 'R_SHOULDER_R', 'rad': -0.720539063587},
    {'joint_name': 'R_SHOULDER_P', 'rad': 0.056068136892},
]
# start_joint_list = [
#
#     {'joint_name': 'R_WRIST_Y', 'rad': 0.182186797436},
#     {'joint_name': 'R_WRIST_P', 'rad': -20.304780654951},
#     {'joint_name': 'R_ELBOW_R', 'rad': -1.535033501186},
#     {'joint_name': 'R_SHOULDER_Y', 'rad': -1.596287366447},
#     {'joint_name': 'R_SHOULDER_R', 'rad': -1.034838057889},
#     {'joint_name': 'R_SHOULDER_P', 'rad': -0.0965316264},
# ]

wyy_joint_list = [
{'joint_name': 'R_WRIST_Y', 'rad': 0.319859674658},
{'joint_name': 'R_WRIST_P', 'rad': -20.645610112769},
{'joint_name': 'R_ELBOW_R', 'rad': -1.129972394858},
{'joint_name': 'R_SHOULDER_Y', 'rad': -1.724514301329},
{'joint_name': 'R_SHOULDER_R', 'rad': -0.557216622708},
{'joint_name': 'R_SHOULDER_P', 'rad': 0.379009062563},
]


play_music_joint_list = [

    {'joint_name': 'R_WRIST_Y', 'rad': 0.00},
    {'joint_name': 'R_WRIST_P', 'rad': -20.446466942304},
    {'joint_name': 'R_ELBOW_R', 'rad': -1.205457349211},
    {'joint_name': 'R_SHOULDER_Y', 'rad': -1.959358596443},
    {'joint_name': 'R_SHOULDER_R', 'rad': -0.776803694305},
    {'joint_name': 'R_SHOULDER_P', 'rad': 0.372088113353},
]


air_joint_list = [

    {'joint_name': 'R_WRIST_Y', 'rad': 0.368498066731},
    {'joint_name': 'R_WRIST_P', 'rad': -20.806707522108},
    {'joint_name': 'R_ELBOW_R', 'rad': -1.160211560687},
    {'joint_name': 'R_SHOULDER_Y', 'rad': -1.754356682523},
    {'joint_name': 'R_SHOULDER_R', 'rad': -0.652379437044},
    {'joint_name': 'R_SHOULDER_P', 'rad': 0.325288409885},
]

empty_joint_list = [
{'joint_name': 'R_WRIST_Y', 'rad': 0.065426748642},
{'joint_name': 'R_WRIST_P', 'rad': -20.572653948527},
{'joint_name': 'R_ELBOW_R', 'rad': -1.354382583029},
{'joint_name': 'R_SHOULDER_Y', 'rad': -1.857614662799},
{'joint_name': 'R_SHOULDER_R', 'rad': -0.490473270507},
{'joint_name': 'R_SHOULDER_P', 'rad': 0.343988546965},
]


nav_joint_list = [

    {'joint_name': 'R_WRIST_Y', 'rad': 0.057895434551},
    {'joint_name': 'R_WRIST_P', 'rad': -20.659321964551},
    {'joint_name': 'R_ELBOW_R', 'rad': -1.348938659874},
    {'joint_name': 'R_SHOULDER_Y', 'rad': -1.969583869368},
    {'joint_name': 'R_SHOULDER_R', 'rad': -0.586725818719},
    {'joint_name': 'R_SHOULDER_P', 'rad': 0.35676823963},
]


air_key_joint_list = [

    {'joint_name': 'R_WRIST_Y', 'rad': 0.394130543969},
    {'joint_name': 'R_WRIST_P', 'rad': -20.514247819877},
    {'joint_name': 'R_ELBOW_R', 'rad': -1.122382227544},
    {'joint_name': 'R_SHOULDER_Y', 'rad': -1.903357855984},
    {'joint_name': 'R_SHOULDER_R', 'rad': -0.756327518825},
    {'joint_name': 'R_SHOULDER_P', 'rad': 0.422019377847},
]

def action_one():
    while True:
        client.send(start_joint_list, vel=5.0, acc=5.0)
        time.sleep(2)
        client.send(set_joint_list, vel=5.0, acc=5.0)
        client.send(start_joint_list, vel=5.0, acc=5.0)
        time.sleep(2)
        client.send(drive_joint_list, vel=5.0, acc=5.0)
        client.send(start_joint_list, vel=5.0, acc=5.0)
        client.send(mode_joint_list, vel=5.0, acc=5.0)
        client.send(start_joint_list, vel=5.0, acc=5.0)
        client.send(home_joint_list, vel=5.0, acc=5.0)

def action_two():
    while True:
        client.send(start_joint_list, vel=5.0, acc=5.0)
        time.sleep(2.0)
        client.send(wyy_joint_list, vel=5.0, acc=5.0)
        client.send(start_joint_list, vel=5.0, acc=5.0)
        time.sleep(2.0)
        client.send(play_music_joint_list, vel=3.0, acc=3.0)
        client.send(start_joint_list, vel=3.0, acc=3.0)
        time.sleep(2.0)
        client.send(play_music_joint_list, vel=3.0, acc=3.0)
        client.send(start_joint_list, vel=3.0, acc=3.0)
        time.sleep(2.0)
        client.send(home_joint_list, vel=5.0, acc=5.0)

def action_three():
    while True:
        client.send(start_joint_list, vel=5.0, acc=5.0)
        time.sleep(2.0)
        client.send(empty_joint_list, vel=5.0, acc=5.0)
        client.send(start_joint_list, vel=5.0, acc=5.0)
        time.sleep(2.0)
        client.send(empty_joint_list, vel=5.0, acc=5.0)
        client.send(start_joint_list, vel=5.0, acc=5.0)
        time.sleep(2.0)
        client.send(nav_joint_list, vel=5.0, acc=5.0)
        client.send(start_joint_list, vel=5.0, acc=5.0)
        time.sleep(7.0)
        client.send(home_joint_list, vel=5.0, acc=5.0)


# ===== 录制工具 =====
JOINT_LISTS = {
    'set_joint_list': set_joint_list,
    'drive_joint_list': drive_joint_list,
    'mode_joint_list': mode_joint_list,
    'home_joint_list': home_joint_list,
    'start_joint_list': start_joint_list,
    'wyy_joint_list': wyy_joint_list,
    'play_music_joint_list': play_music_joint_list,
    'air_joint_list': air_joint_list,
    'air_key_joint_list': air_key_joint_list,
    'empty_joint_list': empty_joint_list,
    'nav_joint_list': nav_joint_list,
}

def find_joint_list_name(lst):
    for name, value in JOINT_LISTS.items():
        if value is lst:
            return name
    return None


class StopRecording(Exception):
    pass


class RecordingClient:
    def __init__(self, max_steps=100):
        self.max_steps = max_steps
        self.steps = []
        self._send_count = 0

    def send(self, joint_list, vel=3.0, acc=3.0):
        name = find_joint_list_name(joint_list)
        self.steps.append(("move", name, vel, acc))
        self._send_count += 1

        # ① 安全兜底：防止真死循环
        if self._send_count >= self.max_steps:
            raise StopRecording

        # ② 关键：第一次回到 home 就结束一轮动作
        if name == "home_joint_list":
            raise StopRecording

    def sleep(self, seconds):
        self.steps.append(("wait", seconds))



def record_action(func, max_steps=50):
    global client

    real_client = client
    real_sleep = time.sleep

    rec_client = RecordingClient(max_steps=max_steps)

    try:
        client = rec_client
        time.sleep = rec_client.sleep
        func()
    except StopRecording:
        pass
    finally:
        client = real_client
        time.sleep = real_sleep

    return rec_client.steps


def joints_to_cmds(joint_list, default_vel=1.0):
    return [
        {
            "jointName": j["joint_name"],
            "rad": j["rad"],
            "vel": default_vel
        }
        for j in joint_list
    ]


def steps_to_json_obj(steps):
    obj = {}
    idx = 1
    for s in steps:
        kind = s[0]
        if kind == "move":
            _, list_name, vel, acc = s
            key = f"{idx}_{list_name}"
            joint_list = JOINT_LISTS[list_name]
            obj[key] = {
                "acc": acc,
                "vel": vel,
                "cmds": joints_to_cmds(joint_list, default_vel=1.0),
            }
        elif kind == "wait":
            _, seconds = s
            key = f"{idx}_wait"
            obj[key] = seconds
        idx += 1
    return obj


# ===== 主入口：录制并保存到文件 =====
if __name__ == "__main__":
    # 占位真实 client（防止 record_action 里引用失败）
    class DummyClient:
        def send(self, *args, **kwargs):
            pass

    client = DummyClient()

    # 分别录制三个动作
    steps_one = record_action(action_one, max_steps=20)
    steps_two = record_action(action_two, max_steps=20)
    steps_three = record_action(action_three, max_steps=20)

    data = {
        "action_one": steps_to_json_obj(steps_one),
        "action_two": steps_to_json_obj(steps_two),
        "action_three": steps_to_json_obj(steps_three),
    }

    # ✅ 保存到文件
    with open("actions.json", "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2, ensure_ascii=False)

    print("已保存到 actions.json")
