#!/usr/bin/env python3
"""
EDITED FROM YOUR ORIGINAL EXAMPLE
---------------------------------
Now it can RECORD and REPLAY full low-level joint data (29 motors) using rt/lowstate + rt/lowcmd.

✅ Uses the SAME Unitree SDK imports/style you pasted (ChannelPublisher/Subscriber, CRC, MotionSwitcherClient).
✅ Saves in the SAME JSONL format as before:
   - First line: {"meta": {...}}
   - Then lines: {"t": <sec>, "q": [29 floats]}

✅ Saves under ./DataG1/ and auto-creates unique filenames:
   take1.jsonl -> take1(1).jsonl -> ...

Usage:
  Record 15s @100Hz:
    python3 g1_lowcmd_record_replay.py enp3s0 --record take_test.jsonl --hz 100 --seconds 15

  Replay:
    python3 g1_lowcmd_record_replay.py enp3s0 --replay take_test.jsonl --speed 1.0

  Replay loop:
    python3 g1_lowcmd_record_replay.py enp3s0 --replay take_test.jsonl --loop
"""

import time
import sys
import json
import argparse
from pathlib import Path
from typing import Optional, List

from unitree_sdk2py.core.channel import ChannelPublisher, ChannelFactoryInitialize
from unitree_sdk2py.core.channel import ChannelSubscriber, ChannelFactoryInitialize
from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowCmd_
from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowState_
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowCmd_
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowState_
from unitree_sdk2py.utils.crc import CRC
from unitree_sdk2py.utils.thread import RecurrentThread
from unitree_sdk2py.comm.motion_switcher.motion_switcher_client import MotionSwitcherClient

import numpy as np

G1_NUM_MOTOR = 29

Kp = [
    60, 60, 60, 100, 40, 40,      # legs
    60, 60, 60, 100, 40, 40,      # legs
    60, 40, 40,                   # waist
    40, 40, 40, 40,  40, 40, 40,  # arms
    40, 40, 40, 40,  40, 40, 40   # arms
]

Kd = [
    1, 1, 1, 2, 1, 1,     # legs
    1, 1, 1, 2, 1, 1,     # legs
    1, 1, 1,              # waist
    1, 1, 1, 1, 1, 1, 1,  # arms
    1, 1, 1, 1, 1, 1, 1   # arms
]

DATA_DIR = Path("DataG1")


class Mode:
    PR = 0
    AB = 1


def ensure_unique_path_in_datag1(user_path: str) -> str:
    """Save under ./DataG1 if no directory given, and uniquify filename if exists."""
    p = Path(user_path)

    if len(p.parts) > 1:
        target_dir = p.parent
        base = p.stem or "record"
        ext = p.suffix or ".jsonl"
        target_dir.mkdir(parents=True, exist_ok=True)
        candidate = target_dir / f"{base}{ext}"
    else:
        DATA_DIR.mkdir(parents=True, exist_ok=True)
        base = p.stem or "record"
        ext = p.suffix or ".jsonl"
        candidate = DATA_DIR / f"{base}{ext}"

    if not candidate.exists():
        return str(candidate)

    i = 1
    while True:
        cand2 = candidate.with_name(f"{candidate.stem}({i}){candidate.suffix}")
        if not cand2.exists():
            return str(cand2)
        i += 1


def resolve_input_path(in_jsonl: str) -> str:
    """If user gives just a filename, load from ./DataG1/filename."""
    p = Path(in_jsonl)
    if len(p.parts) == 1:
        return str(DATA_DIR / p.name)
    return str(p)


class Custom:
    def __init__(self):
        self.low_cmd = unitree_hg_msg_dds__LowCmd_()
        self.low_state: Optional[LowState_] = None
        self.mode_machine_: Optional[int] = None
        self.update_mode_machine_ = False
        self.crc = CRC()

        self.msc: Optional[MotionSwitcherClient] = None
        self.lowcmd_publisher_: Optional[ChannelPublisher] = None
        self.lowstate_subscriber: Optional[ChannelSubscriber] = None

    def Init(self):
        # Motion switcher (same as your code)
        self.msc = MotionSwitcherClient()
        self.msc.SetTimeout(5.0)
        self.msc.Init()

        status, result = self.msc.CheckMode()
        # Release any active mode (same as your code)
        while result and result.get("name"):
            self.msc.ReleaseMode()
            status, result = self.msc.CheckMode()
            time.sleep(0.5)

        # Publisher: rt/lowcmd
        self.lowcmd_publisher_ = ChannelPublisher("rt/lowcmd", LowCmd_)
        self.lowcmd_publisher_.Init()

        # Subscriber: rt/lowstate
        self.lowstate_subscriber = ChannelSubscriber("rt/lowstate", LowState_)
        self.lowstate_subscriber.Init(self.LowStateHandler, 200)

    def LowStateHandler(self, msg: LowState_):
        self.low_state = msg
        if not self.update_mode_machine_:
            self.mode_machine_ = int(msg.mode_machine)
            self.update_mode_machine_ = True

    def WaitFirstState(self, timeout_s: float = 5.0) -> bool:
        t0 = time.time()
        while time.time() - t0 < timeout_s:
            if self.low_state is not None and self.mode_machine_ is not None:
                return True
            time.sleep(0.01)
        return False

    # ---------------------- RECORD ----------------------
    def Record(self, out_jsonl: str, hz: float, seconds: float):
        if not self.WaitFirstState():
            raise SystemExit("[FATAL] No LowState received on rt/lowstate.")

        out_jsonl = ensure_unique_path_in_datag1(out_jsonl)
        dt = 1.0 / float(hz)

        print(f"[REC] {hz:.1f} Hz for {seconds:.2f}s -> {out_jsonl}")

        t0 = time.time()
        next_t = t0
        wrote = 0

        with open(out_jsonl, "w", encoding="utf-8") as f:
            meta = {
                "meta": {
                    "format": "g1_lowcmd_record_v1",
                    "dt": dt,
                    "joints": 29,
                    "created_unix": time.time(),
                    "hz_req": float(hz),
                }
            }
            f.write(json.dumps(meta) + "\n")

            try:
                while True:
                    now = time.time()
                    if now < next_t:
                        time.sleep(min(0.001, next_t - now))
                        continue
                    next_t += dt

                    if self.low_state is None:
                        continue

                    t = now - t0
                    q = [float(self.low_state.motor_state[i].q) for i in range(G1_NUM_MOTOR)]
                    f.write(json.dumps({"t": round(t, 6), "q": q}) + "\n")
                    wrote += 1

                    if seconds > 0 and t >= seconds:
                        break

            except KeyboardInterrupt:
                print("\n[REC] Ctrl+C -> stop and save.")

        print(f"[REC] Done. samples={wrote} file={out_jsonl}")

    # ---------------------- REPLAY ----------------------
    def _send_frame_q29(self, q: List[float], mode_pr: int = Mode.PR):
        if self.mode_machine_ is None:
            raise RuntimeError("mode_machine not ready (no LowState yet)")
        if self.lowcmd_publisher_ is None:
            raise RuntimeError("lowcmd publisher not initialized")
        if len(q) != G1_NUM_MOTOR:
            raise ValueError("q must be length 29")

        self.low_cmd.mode_pr = int(mode_pr)
        self.low_cmd.mode_machine = int(self.mode_machine_)

        for i in range(G1_NUM_MOTOR):
            self.low_cmd.motor_cmd[i].mode = 1  # enable
            self.low_cmd.motor_cmd[i].tau = 0.0
            self.low_cmd.motor_cmd[i].q = float(q[i])
            self.low_cmd.motor_cmd[i].dq = 0.0
            self.low_cmd.motor_cmd[i].kp = float(Kp[i])
            self.low_cmd.motor_cmd[i].kd = float(Kd[i])

        self.low_cmd.crc = self.crc.Crc(self.low_cmd)
        self.lowcmd_publisher_.Write(self.low_cmd)

    def Replay(self, in_jsonl: str, speed: float, loop: bool):
        if not self.WaitFirstState():
            raise SystemExit("[FATAL] No LowState received on rt/lowstate.")

        in_jsonl = resolve_input_path(in_jsonl)

        lines = [ln.strip() for ln in open(in_jsonl, "r", encoding="utf-8") if ln.strip()]
        if len(lines) < 2:
            raise SystemExit("[FATAL] File empty or no samples.")

        # Parse meta
        meta = json.loads(lines[0]).get("meta", {})
        # Parse samples
        samples = []
        for ln in lines[1:]:
            obj = json.loads(ln)
            if "t" in obj and "q" in obj and isinstance(obj["q"], list) and len(obj["q"]) == 29:
                samples.append(obj)
        if not samples:
            raise SystemExit("[FATAL] No valid samples found (need t + q[29]).")

        t0 = float(samples[0]["t"])
        tN = float(samples[-1]["t"])
        dur = max(0.0, tN - t0)

        print(f"[PLAY] {in_jsonl}")
        print(f"[PLAY] samples={len(samples)} recorded_dur={dur:.3f}s speed={speed} loop={loop}")
        input("[PLAY] Press Enter to start (Ctrl+C to stop)... ")

        def run_once():
            wall0 = time.time()
            base_t = float(samples[0]["t"])
            for s in samples:
                tt = (float(s["t"]) - base_t) / max(float(speed), 1e-9)
                target = wall0 + tt

                while True:
                    dtw = target - time.time()
                    if dtw <= 0:
                        break
                    time.sleep(min(0.001, dtw))

                self._send_frame_q29(s["q"], mode_pr=Mode.PR)

        try:
            while True:
                run_once()
                if not loop:
                    break
        except KeyboardInterrupt:
            print("\n[PLAY] Ctrl+C -> stop publishing.")
        finally:
            print("[PLAY] Done (stopped publishing).")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("iface", help="network iface (example: enp3s0)")
    ap.add_argument("--record", default="", help="Output jsonl (saved under ./DataG1/ if no path)")
    ap.add_argument("--replay", default="", help="Input jsonl (loaded from ./DataG1/ if no path)")
    ap.add_argument("--hz", type=float, default=200.0, help="Record Hz")
    ap.add_argument("--seconds", type=float, default=5.0, help="Record seconds (0=until Ctrl+C)")
    ap.add_argument("--speed", type=float, default=1.0, help="Replay speed (1.0=same)")
    ap.add_argument("--loop", action="store_true", help="Loop replay")
    args = ap.parse_args()

    if bool(args.record) == bool(args.replay):
        raise SystemExit("Choose exactly one: --record OUT.jsonl  OR  --replay IN.jsonl")

    print("WARNING: Please ensure there are no obstacles around the robot while running.")
    input("Press Enter to continue...")

    # DDS init (same idea as your original)
    ChannelFactoryInitialize(0, args.iface)

    custom = Custom()
    custom.Init()

    if args.record:
        custom.Record(args.record, hz=args.hz, seconds=args.seconds)
    else:
        custom.Replay(args.replay, speed=args.speed, loop=args.loop)

    # Keep process alive briefly to let DDS flush prints (optional)
    time.sleep(0.2)


if __name__ == "__main__":
    main()
