#!/usr/bin/env python3
"""
G1 PASSIVE RECORDER (GHOST MODE)
--------------------------------
1. ZERO CONTROL: Does NOT send any commands.
2. COMPATIBILITY: Works perfectly alongside xr_teleoperate.
3. SAFEST OPTION: Impossible to cause balance loss.

Usage:
  python3 g1_record_only_passive.py enp3s0 --output my_teleop_move
  
"""

import time
import sys
import json
import argparse
import numpy as np
from pathlib import Path
from unitree_sdk2py.core.channel import ChannelSubscriber, ChannelFactoryInitialize
from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowState_
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowState_

G1_NUM_MOTOR = 29
DATA_DIR = Path("DataG1")
RECORD_HZ = 60.0

def ensure_unique_path(user_path: str) -> str:
    p = Path(user_path)
    if not p.suffix: p = p.with_suffix(".jsonl")
    
    if len(p.parts) > 1:
        target_dir = p.parent
        target_dir.mkdir(parents=True, exist_ok=True)
        candidate = target_dir / p.name
    else:
        DATA_DIR.mkdir(parents=True, exist_ok=True)
        candidate = DATA_DIR / p.name

    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

class PassiveRecorder:
    def __init__(self):
        self.low_state = None
        # ONLY Subscriber. NO Publisher.
        self.state_sub = ChannelSubscriber("rt/lowstate", LowState_)
        self.state_sub.Init(self.LowStateHandler, 10)
        self.first_state = False

    def LowStateHandler(self, msg: LowState_):
        self.low_state = msg
        self.first_state = True

    def Run(self, filename: str, duration: float):
        print("Waiting for robot connection...", end="", flush=True)
        while not self.first_state: time.sleep(0.1)
        print(" Connected!")
        
        print("\n" + "="*50)
        print("👻 GHOST RECORDER ACTIVE")
        print(f"   - Control:      NONE (Passive Listener)")
        print(f"   - Teleop:       Safe to use 'xr_teleoperate'")
        print(f"   - Saving to:    {filename}")
        print("="*50 + "\n")

        t_start = time.time()
        dt = 1.0 / RECORD_HZ
        samples = []

        try:
            while True:
                loop_start = time.time()
                t_curr = loop_start - t_start
                
                # Check Duration (if set)
                if duration > 0 and t_curr > duration: break

                # 1. READ STATE (Listen Only)
                current_q_all = [float(self.low_state.motor_state[i].q) for i in range(G1_NUM_MOTOR)]
                samples.append({"t": round(t_curr, 4), "q": current_q_all})

                # Timing
                time.sleep(max(0, dt - (time.time() - loop_start)))
                
                # Feedback
                if len(samples) % int(RECORD_HZ) == 0:
                     sys.stdout.write(f"\r🔴 Recording... {t_curr:.1f}s")
                     sys.stdout.flush()

        except KeyboardInterrupt:
            print("\nStopped by user.")
        
        # Save
        print(f"\n💾 Saving to {filename}...")
        with open(filename, "w") as f:
            f.write(json.dumps({"meta": {"hz": RECORD_HZ, "motors": G1_NUM_MOTOR}}) + "\n")
            for s in samples:
                f.write(json.dumps(s) + "\n")
        print(f"✅ Saved {len(samples)} samples.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("iface", help="Network interface")
    parser.add_argument("--output", default="teleop_record", help="Output filename")
    parser.add_argument("--seconds", type=float, default=0.0, help="0 = Record until Ctrl+C")
    args = parser.parse_args()

    ChannelFactoryInitialize(0, args.iface)
    final_path = ensure_unique_path(args.output)
    PassiveRecorder().Run(final_path, args.seconds)