#!/usr/bin/env python3
"""Chinese phase banners for LeRobot recording; no changes to robot control."""
import sys


def show_phase(message: str) -> None:
    if message.startswith("Recording episode "):
        number = int(message.rsplit(" ", 1)[1]) + 1
        title, color = f"● 正在录制｜数据集第 {number} 条", "32"
        lines = [
            "现在操作主臂：抓起积木 → 放入白盒 → 松开 → 退开",
            "完成：按 → 结束本条，进入复位（最后一条将保存退出）",
            "失败：按 ← 标记重录，然后按提示复位",
            "停止：按 Esc；当前片段可能保存，之后检查完整性",
        ]
    elif message == "Reset the environment":
        title, color = "● 复位阶段｜当前不录入示范", "33"
        lines = [
            "把积木取出放回起点，并通过主臂让从臂回到起始姿态",
            "注意：从臂仍跟随主臂！",
            "准备好：按 → 继续；超时也会自动继续",
            "上一条失败：按 ← 标记重录；按 Esc 结束采集",
        ]
    elif message == "Re-record episode":
        title, color = "● 已丢弃本条｜即将重录", "33"
        lines = ["等待绿色“正在录制”提示后，再开始抓取。"]
    elif message == "Stop recording":
        title, color = "● 正在结束采集｜请等待保存与清理", "36"
        lines = ["不要关闭终端；完成后会返回命令行。"]
    elif message == "Exiting":
        title, color = "● 采集程序已退出", "36"
        lines = ["请检查实际保存条数及视频；提前结束的片段可能不完整。"]
    else:
        return
    use_color = sys.stdout.isatty()
    start = f"\033[1;{color}m" if use_color else ""
    end = "\033[0m" if use_color else ""
    print("\n" + start + "━" * 60, flush=True)
    print(title, flush=True)
    print("━" * 60, flush=True)
    for line in lines:
        print("  " + line, flush=True)
    print("━" * 60 + end + "\n", flush=True)
    if use_color:
        print("\a", end="", flush=True)


def main() -> None:
    if sys.argv[1:] == ["--preview"]:
        for message in ["Recording episode 0", "Reset the environment", "Re-record episode", "Stop recording", "Exiting"]:
            show_phase(message)
        return
    from lerobot.scripts import lerobot_record

    original_log_say = lerobot_record.log_say

    def log_with_banner(message, *args, **kwargs):
        result = original_log_say(message, *args, **kwargs)
        show_phase(message)
        return result

    lerobot_record.log_say = log_with_banner
    lerobot_record.main()


if __name__ == "__main__":
    main()