#!/usr/bin/env python3
"""
G1 REPLAY TRIGGER (R2 + X) - SINGLE FILE
---------------------------------------
- Hold R2 + X to replay DataG1/photo_G3.jsonl (default).
- While playing: ignore any additional R2+X until replay finishes and returns home.
- Cancel combo: R2 + L1 cancels replay immediately and returns arms to home.

Usage:
  python3 g1_replay_trigger_r2x.py enp3s0
  python3 g1_replay_trigger_r2x.py enp3s0 --input photo_G3.jsonl --home arm_home.jsonl --speed 1.0
"""

import time
import sys
import json
import argparse
import struct
from pathlib import Path

from unitree_sdk2py.core.channel import ChannelPublisher, 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 Logger import Logs

# ---------------- Logger (exactly as requested) ----------------
controller_logs = Logs()
controller_logs.LogEngine("G1_Logs", "g1_replay_trigger_r2x.log")

# ---------------- Constants ----------------
G1_NUM_MOTOR = 29
ENABLE_ARM_SDK_INDEX = 29

DATA_DIR = Path("DataG1")
REPLAY_HZ = 60.0

# --- GAINS (same as your replay) ---
KP_HIGH = 300.0  # Core/Legs
KD_HIGH = 3.0
KP_LOW  = 80.0   # Arms/Ankles
KD_LOW  = 3.0
KP_WRIST = 40.0
KD_WRIST = 1.5

WEAK_MOTORS  = [4, 10, 15, 16, 17, 18, 22, 23, 24, 25]
WRIST_MOTORS = [19, 20, 21, 26, 27, 28]

# ---------------- Helpers ----------------
def resolve_input_path(in_path: str) -> str:
    p = Path(in_path)
    if len(p.parts) == 1:
        return str(DATA_DIR / p.name)
    return str(p)

def load_home_pose(home_path: str):
    """Reads the last frame of arm_home.jsonl to get the target pose."""
    path = resolve_input_path(home_path)
    try:
        last_valid_q = None
        with open(path, 'r') as f:
            for line in f:
                d = json.loads(line)
                if 'q' in d and len(d['q']) == G1_NUM_MOTOR:
                    last_valid_q = d['q']
        if last_valid_q:
            controller_logs.print_and_log(f"✅ Loaded Home Pose from {path}", "info")
            return last_valid_q
        else:
            controller_logs.print_and_log(f"⚠️  Warning: {path} found but contained no valid 'q' data.", "warning")
    except FileNotFoundError:
        controller_logs.print_and_log(f"⚠️  Warning: Home file {path} not found.", "warning")

    controller_logs.print_and_log("⚠️  Using Default Home (Arms at 0.0)", "warning")
    return [0.0] * G1_NUM_MOTOR

# ---------------- Wireless Controller Parser (embedded) ----------------
class unitreeRemoteController:
    def __init__(self):
        self.Lx = 0; self.Rx = 0; self.Ry = 0; self.Ly = 0
        self.L1 = 0; self.L2 = 0; self.R1 = 0; self.R2 = 0
        self.A = 0; self.B = 0; self.X = 0; self.Y = 0
        self.Up = 0; self.Down = 0; self.Left = 0; self.Right = 0
        self.Select = 0; self.F1 = 0; self.F3 = 0; self.Start = 0

    def parse_botton(self, data1, data2):
        self.R1 = (data1 >> 0) & 1; self.L1 = (data1 >> 1) & 1
        self.Start = (data1 >> 2) & 1; self.Select = (data1 >> 3) & 1
        self.R2 = (data1 >> 4) & 1; self.L2 = (data1 >> 5) & 1
        self.F1 = (data1 >> 6) & 1; self.F3 = (data1 >> 7) & 1
        self.A = (data2 >> 0) & 1; self.B = (data2 >> 1) & 1
        self.X = (data2 >> 2) & 1; self.Y = (data2 >> 3) & 1
        self.Up = (data2 >> 4) & 1; self.Right = (data2 >> 5) & 1
        self.Down = (data2 >> 6) & 1; self.Left = (data2 >> 7) & 1

    def parse_key(self, data):
        offsets = [4, 8, 12, 20]  # Lx, Rx, Ry, Ly
        self.Lx, self.Rx, self.Ry, self.Ly = [struct.unpack('<f', data[o:o+4])[0] for o in offsets]

    def parse(self, remoteData):
        self.parse_key(remoteData)
        self.parse_botton(remoteData[2], remoteData[3])

    def get_state(self):
        return self.__dict__.copy()

# ---------------- Replay Engine ----------------
class ReplayWithHome:
    def __init__(self, watchdog_timeout=0.25, watchdog_disable_after=1.0):
        self.low_state = None
        self.low_cmd = unitree_hg_msg_dds__LowCmd_()
        self.crc = CRC()

        self.arm_pub = ChannelPublisher("rt/arm_sdk", LowCmd_)
        self.arm_pub.Init()

        self.state_sub = None
        self.first_state = False

        # watchdog
        self.last_state_time = 0.0
        self.watchdog_timeout = float(watchdog_timeout)
        self.watchdog_disable_after = float(watchdog_disable_after)

        # controller
        self.remote = unitreeRemoteController()
        self.controller_state = self.remote.get_state()

        self.is_playing = False

    def InitStateSubscriber(self, topic: str):
        self.state_sub = ChannelSubscriber(topic, LowState_)
        self.state_sub.Init(self.LowStateHandler, 10)

    def LowStateHandler(self, msg: LowState_):
        self.low_state = msg
        self.first_state = True
        self.last_state_time = time.time()
        try:
            self.remote.parse(msg.wireless_remote)
            self.controller_state = self.remote.get_state()
        except Exception:
            pass

    def StateFresh(self) -> bool:
        return (time.time() - self.last_state_time) < self.watchdog_timeout

    def CancelRequested(self) -> bool:
        """Cancel combo: R2 + L1"""
        s = self.controller_state
        return bool(s.get("R2", 0) and s.get("L1", 0))

    def SendFrame(self, arm_target_q, body_lock_q):
        self.low_cmd.motor_cmd[ENABLE_ARM_SDK_INDEX].q = 1.0

        for i in range(G1_NUM_MOTOR):
            self.low_cmd.motor_cmd[i].mode = 1
            self.low_cmd.motor_cmd[i].dq = 0
            self.low_cmd.motor_cmd[i].tau = 0

            if i >= 15:
                self.low_cmd.motor_cmd[i].q = arm_target_q[i]
            else:
                self.low_cmd.motor_cmd[i].q = body_lock_q[i]

            if i in WEAK_MOTORS:
                self.low_cmd.motor_cmd[i].kp = KP_LOW
                self.low_cmd.motor_cmd[i].kd = KD_LOW
            elif i in WRIST_MOTORS:
                self.low_cmd.motor_cmd[i].kp = KP_WRIST
                self.low_cmd.motor_cmd[i].kd = KD_WRIST
            else:
                self.low_cmd.motor_cmd[i].kp = KP_HIGH
                self.low_cmd.motor_cmd[i].kd = KD_HIGH

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

    def DisableSDK(self):
        controller_logs.print_and_log("🔌 Disabling SDK...", "info")
        self.low_cmd.motor_cmd[ENABLE_ARM_SDK_INDEX].q = 0.0
        self.low_cmd.crc = self.crc.Crc(self.low_cmd)
        for _ in range(10):
            self.arm_pub.Write(self.low_cmd)
            time.sleep(0.02)

    def ReturnArmsHome(self, last_arm_q, body_lock_q, home_q, home_steps=180):
        """Always go home smoothly (arms only), body remains locked."""
        controller_logs.print_and_log("🏡 Returning arms to HOME...", "info")
        for k in range(home_steps):
            stale_for = (time.time() - self.last_state_time)
            if stale_for > self.watchdog_disable_after:
                controller_logs.print_and_log("🛑 WATCHDOG: connection lost during home. Disabling SDK.", "error")
                self.DisableSDK()
                return

            if not self.StateFresh():
                controller_logs.print_and_log("⚠️ WATCHDOG: state stale during home. Holding last pose...", "warning")
                self.SendFrame(last_arm_q, body_lock_q)
                time.sleep(1.0 / REPLAY_HZ)
                continue

            alpha = k / home_steps
            interp_q = list(last_arm_q)
            for j in range(15, 29):
                interp_q[j] = (1-alpha)*last_arm_q[j] + alpha*home_q[j]
            self.SendFrame(interp_q, body_lock_q)
            time.sleep(1.0 / REPLAY_HZ)

        controller_logs.print_and_log("✅ Home Reached.", "info")

    def RunReplay(self, filename: str, home_filename: str, speed: float):
        self.is_playing = True
        cancel_requested = False

        try:
            controller_logs.print_and_log(f"🎬 Triggered replay: {filename}", "info")

            controller_logs.print_and_log("Waiting for robot...", "info")
            while not self.first_state:
                time.sleep(0.05)
            controller_logs.print_and_log("✅ Robot Connected!", "info")

            home_q = load_home_pose(home_filename)
            full_body_lock_q = [self.low_state.motor_state[i].q for i in range(G1_NUM_MOTOR)]

            frames = []
            with open(filename, 'r') as f:
                for line in f:
                    d = json.loads(line)
                    if 'q' in d:
                        frames.append(d)

            if not frames:
                controller_logs.print_and_log("❌ No frames found in input file.", "error")
                return

            controller_logs.print_and_log(f"🟢 Ready to play {len(frames)} frames.", "info")
            controller_logs.print_and_log("🔒 Body LOCKED. Arms replay then return HOME.", "info")
            controller_logs.print_and_log("🛑 Cancel combo: Hold R2 + L1", "info")

            # Move to start
            controller_logs.print_and_log("Moving to start...", "info")
            file_start_q = frames[0]['q']
            last_played_q = file_start_q  # will update

            steps = 60
            for k in range(steps):
                if self.CancelRequested():
                    controller_logs.print_and_log("🛑 CANCEL: R2+L1 detected أثناء الانتقال للبداية.", "warning")
                    cancel_requested = True
                    break

                if not self.StateFresh():
                    controller_logs.print_and_log("⚠️ WATCHDOG: stale while moving to start. Holding...", "warning")
                    self.SendFrame(last_played_q, full_body_lock_q)
                    time.sleep(1.0 / REPLAY_HZ)
                    continue

                alpha = k / steps
                interp_q = list(full_body_lock_q)
                for j in range(15, 29):
                    interp_q[j] = (1-alpha)*full_body_lock_q[j] + alpha*file_start_q[j]
                self.SendFrame(interp_q, full_body_lock_q)
                last_played_q = interp_q
                time.sleep(1.0 / REPLAY_HZ)

            # Play replay (unless canceled)
            if not cancel_requested:
                controller_logs.print_and_log("▶️ Playing...", "info")
                play_elapsed = 0.0
                last_real = time.time()

                while True:
                    if self.CancelRequested():
                        controller_logs.print_and_log("🛑 CANCEL: R2+L1 detected أثناء التشغيل.", "warning")
                        cancel_requested = True
                        break

                    stale_for = (time.time() - self.last_state_time)
                    if stale_for > self.watchdog_disable_after:
                        controller_logs.print_and_log("🛑 WATCHDOG: connection lost. Disabling SDK.", "error")
                        self.DisableSDK()
                        return

                    if not self.StateFresh():
                        controller_logs.print_and_log("⚠️ WATCHDOG: timeout! Holding last pose...", "warning")
                        self.SendFrame(last_played_q, full_body_lock_q)
                        time.sleep(1.0 / REPLAY_HZ)
                        continue

                    now_real = time.time()
                    dt_real = now_real - last_real
                    last_real = now_real
                    play_elapsed += dt_real * speed

                    target_frame = None
                    for fr in frames:
                        if fr['t'] - frames[0]['t'] >= play_elapsed:
                            target_frame = fr
                            break
                    if target_frame is None:
                        break

                    self.SendFrame(target_frame['q'], full_body_lock_q)
                    last_played_q = target_frame['q']
                    time.sleep(1.0 / REPLAY_HZ)

            # Always go home (normal finish OR cancel)
            self.ReturnArmsHome(last_played_q, full_body_lock_q, home_q, home_steps=180)
            self.DisableSDK()

        except Exception as e:
            controller_logs.print_and_log(f"❌ Exception in replay: {e}", "error")
            try:
                self.DisableSDK()
            except Exception:
                pass
        finally:
            self.is_playing = False

# ---------------- Trigger Loop (R2 + X) ----------------
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("iface", help="Network interface")
    parser.add_argument("--input", default="photo_G3.jsonl", help="Input recording file (default: photo_G3.jsonl)")
    parser.add_argument("--home", default="arm_home.jsonl", help="Home pose file")
    parser.add_argument("--speed", type=float, default=1.0)
    parser.add_argument("--state_topic", default="rt/lowstate", help="Lowstate topic (default: rt/lowstate)")
    parser.add_argument("--watchdog", type=float, default=0.25, help="Watchdog timeout seconds")
    parser.add_argument("--watchdog_disable_after", type=float, default=1.0, help="Disable SDK if stale longer than this")
    args = parser.parse_args()

    ChannelFactoryInitialize(0, args.iface)

    engine = ReplayWithHome(
        watchdog_timeout=args.watchdog,
        watchdog_disable_after=args.watchdog_disable_after
    )
    engine.InitStateSubscriber(args.state_topic)

    target_path = resolve_input_path(args.input)

    controller_logs.print_and_log("✅ G1 R2+X Trigger Ready", "info")
    controller_logs.print_and_log(f"🎯 Trigger file: {target_path}", "info")
    controller_logs.print_and_log("🎮 Hold R2 + X to replay", "info")
    controller_logs.print_and_log("🛑 Cancel while playing: Hold R2 + L1", "info")

    # Rising-edge trigger + require release before next trigger
    combo_prev = False

    while True:
        time.sleep(0.01)

        if not engine.first_state:
            continue

        s = engine.controller_state

        # Trigger combo: R2 + X
        combo_now = bool(s.get("R2", 0) and s.get("X", 0))

        # ✅ IGNORE while playing (and also prevent retrigger until released)
        if engine.is_playing:
            combo_prev = combo_now
            continue

        # Trigger only on rising edge (must release and press again)
        if combo_now and not combo_prev:
            controller_logs.print_and_log("[TRIGGER] R2 + X detected", "info")
            engine.RunReplay(target_path, args.home, args.speed)
            # After RunReplay returns, if user still holding, combo_prev will get updated next loop
            # and won’t retrigger until they release and press again.

        combo_prev = combo_now

if __name__ == "__main__":
    main()