# coding=utf-8
import json
import logging
import os
import socket
import time
import sys
import subprocess
import re
import numpy as np
import cv2
import pyrealsense2 as rs

from libs.log_setting import CommonLog
from libs.auxiliary import create_folder_with_date, get_ip, popup_message
from save_poses import save_matrices_to_csv

# 把 .so 所在目录加入 Python 路径
sys.path.append("/home/lgv/cmvr/cmvr-es/cmake-build-debug/example")
from robot_wrapper import Robot

# -------------------- 全局配置 --------------------
TARGET_SERIAL = "243122075614"  # 相机序列号
EXEC_PATH = "/home/lgv/cmvr/cmvr-es/cmake-build-debug/example/solve_fk"
ROBOT_CONFIG = "/home/lgv/cmvr/cmvr-es/config/cabin_robot.xml"

logger = logging.getLogger(__name__)
logger = CommonLog(logger)

robot = Robot(ROBOT_CONFIG, "hc01")
output_dir = create_folder_with_date()  # 存储目录
count = 1


# -------------------- 工具函数 --------------------
def get_pose(joint_positions, base_link="PELVIS_S", target_link="R_FINGER_TIP", exec_path=EXEC_PATH):
    """通过 solve_fk 求解 4x4 位姿矩阵"""
    joint_str = [str(j) for j in joint_positions]
    cmd = [exec_path, base_link, target_link, *joint_str]

    try:
        res = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                             check=True, text=True)
    except subprocess.CalledProcessError as e:
        raise RuntimeError(f"solve_fk 调用失败: {e.stdout}")

    lines = res.stdout.splitlines()
    start_idx = next((i for i, l in enumerate(lines) if re.search(r"变换矩阵|T\s*:", l)), None)
    if start_idx is None or start_idx + 4 >= len(lines):
        raise RuntimeError("未找到 4x4 变换矩阵:\n" + res.stdout)

    try:
        mat = np.array([[float(x) for x in lines[start_idx + 1 + r].split()] for r in range(4)], dtype=np.float64)
        if mat.shape != (4, 4):
            raise ValueError
    except Exception:
        raise RuntimeError("矩阵解析失败:\n" + res.stdout)
    return mat


def save_pose_and_image(pose, image, index):
    """保存位姿矩阵和对应的图片"""
    try:
        # 保存位姿（追加模式 'a'）
        pose_file = os.path.join(output_dir, "RobotToolPose.csv")
        save_matrices_to_csv([pose], pose_file)  # 假设 save_matrices_to_csv 支持 mode 参数

        # 保存图片
        img_file = os.path.join(output_dir, f"{index}.jpg")
        cv2.imwrite(img_file, image)

        logger.info(f"===采集第{index}次数据！ pose 已保存到 {pose_file}, 图片保存到 {img_file}")
    except Exception as e:
        logger.error(f"保存数据失败: {e}")


def callback_and_save(frame):
    """显示视频帧，按 's' 保存，按 'q' 返回 False 退出"""
    global count
    cv2.imshow("Capture_Video2", frame)
    scaling_factor = 2.0
    cv_img = cv2.resize(frame, None, fx=scaling_factor, fy=scaling_factor,
                        interpolation=cv2.INTER_AREA)
    cv2.imshow("Capture_Video", cv_img)

    k = cv2.waitKey(30) & 0xFF
    if k == ord('s'):
        try:
            # 获取关节角 & 位姿
            js = robot.getJointQ('right')
            logger.info(f"当前关节角: {js}")

            pose = get_pose(js, base_link="PELVIS_S", target_link="R_FINGER_TIP")
            logger.info(f"获取位姿成功:\n{pose}")
            logger.info("请按 p 失能机器人，并手动移动机器人到其他位姿")

            save_pose_and_image(pose, cv_img, count)
            count += 1

        except Exception as e:
            logger.error(f"获取或保存数据失败: {e}")
    elif k == ord('q'):
        logger.info("检测到 'q'，退出程序")
        return False
    elif k == ord('o'):
        robot.torqueOn()
        logger.info("robot.torqueOn ，请勿触摸机器人")
    elif k == ord('p'):
        logger.info("robot.torqueOff")
        robot.torqueOff()


    return True




def start_pipeline(serial=TARGET_SERIAL, depth_size=(1280, 720), color_size=(1280, 720), fps: int = 30):
    ctx = rs.context()
    devices = ctx.query_devices()
    if len(devices) == 0:
        raise RuntimeError("未检测到任何 RealSense 相机")

    if not any(dev.get_info(rs.camera_info.serial_number) == serial for dev in devices):
        raise RuntimeError(f"未找到序列号为 {serial} 的 RealSense 相机")

    pipeline = rs.pipeline()
    cfg = rs.config()
    cfg.enable_device(serial)

    cfg.enable_stream(rs.stream.depth, depth_size[0], depth_size[1], rs.format.z16, fps)
    cfg.enable_stream(rs.stream.color, color_size[0], color_size[1], rs.format.bgr8, fps)

    pipeline_profile = pipeline.start(cfg)

    # 相机内参
    depth_intr = pipeline_profile.get_stream(rs.stream.depth).as_video_stream_profile().get_intrinsics()
    color_intr = pipeline_profile.get_stream(rs.stream.color).as_video_stream_profile().get_intrinsics()

    # 建议对齐到彩色流
    align = rs.align(rs.stream.color)
    # align = rs.align(rs.stream.depth)

    return pipeline, align, depth_intr, color_intr
def collect_data():
    """循环采集机器人位姿与相机图像"""
    logger.info("开始机器人数据采集程序，版本 V1.0.0")

    pipeline, align, depth_intr, color_intr = start_pipeline()
    global count
    count = 1

    try:
        while True:
            frames = pipeline.wait_for_frames()
            aligned_frames = align.process(frames)
            color_frame = aligned_frames.get_color_frame()
            depth_frame = aligned_frames.get_depth_frame()
            if not color_frame or not depth_frame:
                continue

            # 转换为 numpy
            color_image = np.asanyarray(color_frame.get_data())
            depth_image = np.asanyarray(depth_frame.get_data())

            # 调用一次 callback_and_save
            if not callback_and_save(color_image):
                break

    finally:
        pipeline.stop()
        cv2.destroyAllWindows()
        logger.info("相机已关闭，窗口已销毁")


# -------------------- 主入口 --------------------
if __name__ == '__main__':
    collect_data()