#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
session_server_node.py — 板子端：会话式「抓帧 / 回传图像 / 回传点云 / 误差跟踪」

一次运行 = 一个会话。电脑端每一条请求、板子每一条应答都带 session
（id / client_ip / board_ip / started_ns），会话之间状态隔离。

pynng REQ/REP 协议（默认 tcp://0.0.0.0:5555）：

  hello       握手；返回标定分辨率、是否支持板子端计算
  capture     等一帧同步的（图像 + 点云 + 位姿），返回图像字节；
              本会话记住这一帧（点云、位姿），并清掉旧目标
  click       默认：电脑只发 (u, v)，板子用本地点云算
              p_cam / p_lidar / p_map，并把它记成目标（点云不出板）
  cloud       可选（--compute pc）：返回 capture 那一帧的点云
  set_target  可选（--compute pc）：电脑算好的 p_map 回传
  status      返回当前 pos / yaw / target / Δ（只给坐标，不给距离）
  reset       清除本会话的目标

投影优化：
  click 第一次会建一个「像素格子 -> 点索引」倒排表（每帧一次），
  之后同一帧的点击只查周围几个格子，不再全量投影，平面拟合也只在
  局部候选里做。格子大小 --index-cell（默认 32 px）。

用法（板子）：
  source /opt/ros/jazzy/setup.bash
  python3 session_server_node.py --port 5555

常用参数：
  --image-topic /odin1/image/compressed   （默认走压缩图，直接转发，不用编码）
  --no-compressed                          （强制用原始 /odin1/image，需要 cv2 或 Pillow）
  --cloud-topic /odin1/cloud_raw
  --odom-topic  /odin1/odometry
"""
import argparse
import threading
import time
import zlib
from collections import deque

import numpy as np
import rclpy
from rclpy.node import Node
from rclpy.qos import (QoSProfile, ReliabilityPolicy, HistoryPolicy,
                       qos_profile_sensor_data)
from sensor_msgs.msg import Image, CompressedImage, PointCloud2
from nav_msgs.msg import Odometry
from message_filters import Subscriber, ApproximateTimeSynchronizer

import pynng

import wire

try:
    import calib
    CALIB_W = int(calib.IMAGE_WIDTH)
    CALIB_H = int(calib.IMAGE_HEIGHT)
except Exception:  # 没放 calib.py 也能跑，只是分辨率用默认值
    CALIB_W, CALIB_H = 1600, 1296

# 板子端本地算 p_cam/p_lidar/p_map（默认模式）。没有 target_math.py 时
# 自动降级为"电脑端算 + 点云回传"。
try:
    import target_math
    _HAS_TARGET_MATH = True
except Exception:
    target_math = None
    _HAS_TARGET_MATH = False

# sensor_msgs/PointField.datatype -> numpy dtype
_PF = {1: np.int8, 2: np.uint8, 3: np.int16, 4: np.uint16,
       5: np.int32, 6: np.uint32, 7: np.float32, 8: np.float64}
_CH = {'bgr8': 3, 'rgb8': 3, 'bgra8': 4, 'rgba8': 4, 'mono8': 1}


def stamp_ns(h):
    return int(h.stamp.sec) * 1_000_000_000 + int(h.stamp.nanosec)


def unpack_cloud(msg):
    """按字段名/偏移解析 PointCloud2（兼容 padding 与非对齐字段）。"""
    names, formats, offsets = [], [], []
    for f in msg.fields:
        if f.datatype not in _PF:
            continue
        names.append(f.name)
        formats.append(_PF[f.datatype])
        offsets.append(f.offset)
    dt = np.dtype({'names': names, 'formats': formats,
                   'offsets': offsets, 'itemsize': msg.point_step})
    arr = np.frombuffer(msg.data, dtype=dt, count=msg.width * msg.height)
    return arr, names


def image_to_array(msg):
    ch = _CH.get(msg.encoding)
    if ch is None:
        return None
    row = msg.step if msg.step else msg.width * ch
    buf = np.frombuffer(msg.data, np.uint8).reshape(msg.height, row)
    img = buf[:, :msg.width * ch].reshape(msg.height, msg.width, ch)
    return img.copy()


def quat_to_yaw(x, y, z, w):
    return float(np.arctan2(2.0 * (w * z + x * y),
                            1.0 - 2.0 * (y * y + z * z)))


def quat_to_rot(q):
    x, y, z, w = [float(v) for v in q]
    return np.array([
        [1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)],
        [2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)],
        [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)],
    ], dtype=np.float64)


class SessionServer(Node):
    def __init__(self, args):
        super().__init__('session_server')
        self.args = args
        self._stop = False
        self._rep = None

        self._lock = threading.Lock()
        self._cv = threading.Condition(self._lock)
        self._latest = None
        self._odom_hist = deque(maxlen=400)
        self._sessions = {}

        self._setup_subs()

        self.get_logger().info(
            '会话服务: image=%s cloud=%s odom=%s'
            % (args.image_topic, args.cloud_topic, args.odom_topic))
        self.get_logger().info(
            '标定分辨率 %dx%d，等待 %s 的第一帧…' % (CALIB_W, CALIB_H, args.image_topic))

    # ------------------------------------------------------------------ 订阅
    def _setup_subs(self):
        img_type = CompressedImage if self.args.image_topic.rstrip('/').endswith('/compressed') else Image
        a = Subscriber(self, img_type, self.args.image_topic,
                       qos_profile=qos_profile_sensor_data)
        b = Subscriber(self, PointCloud2, self.args.cloud_topic,
                       qos_profile=qos_profile_sensor_data)
        self.sync = ApproximateTimeSynchronizer([a, b],
                                                queue_size=self.args.queue,
                                                slop=self.args.slop)
        self.sync.registerCallback(self.on_pair)

        # 里程计：best-effort 订阅既能接 reliable 也能接 best-effort 的发布者
        odom_qos = QoSProfile(depth=100,
                              reliability=ReliabilityPolicy.BEST_EFFORT,
                              history=HistoryPolicy.KEEP_LAST)
        self.create_subscription(Odometry, self.args.odom_topic,
                                 self.on_odom, odom_qos)

    def on_odom(self, msg):
        p = msg.pose.pose.position
        q = msg.pose.pose.orientation
        rec = {
            'stamp_ns': stamp_ns(msg.header),
            'recv_mono': time.monotonic(),
            'pos': (float(p.x), float(p.y), float(p.z)),
            'quat': (float(q.x), float(q.y), float(q.z), float(q.w)),
            'frame_id': msg.header.frame_id,
            'child_frame_id': msg.child_frame_id,
        }
        with self._lock:
            self._odom_hist.append(rec)

    def on_pair(self, img_msg, pc_msg):
        """图像 + 点云时间对齐后，把这一帧放进 latest。"""
        arr, names = unpack_cloud(pc_msg)
        if not all(k in names for k in ('x', 'y', 'z')):
            self.get_logger().warn('点云没有 x/y/z 字段，实际：%s' % names)
            return

        xyz = np.stack([arr['x'], arr['y'], arr['z']], axis=-1).astype(np.float32)
        conf = np.ascontiguousarray(arr['confidence']) if 'confidence' in names else None
        cloud_stamp = stamp_ns(pc_msg.header)
        img_stamp = stamp_ns(img_msg.header)

        if isinstance(img_msg, CompressedImage):
            frame = {
                'kind': 'compressed',
                'blob': bytes(img_msg.data),
                'format': str(img_msg.format or 'jpeg'),
                'width': None,
                'height': None,
            }
        else:
            img = image_to_array(img_msg)
            if img is None:
                self.get_logger().warn('不支持的图像编码：%s' % img_msg.encoding)
                return
            frame = {
                'kind': 'raw',
                'img': img,
                'format': None,
                'width': int(img_msg.width),
                'height': int(img_msg.height),
                'encoding': img_msg.encoding,
            }

        frame.update({
            'xyz': xyz,
            'conf': conf,
            'cloud_stamp_ns': cloud_stamp,
            'img_stamp_ns': img_stamp,
            'cloud_frame_id': pc_msg.header.frame_id,
            'recv_mono': time.monotonic(),
        })

        with self._cv:
            self._latest = frame
            self._cv.notify_all()

        self.get_logger().debug(
            '同步帧: 图像 %s, 点云 %d 点, 时间差 %.2f ms'
            % (frame['kind'], xyz.shape[0], (cloud_stamp - img_stamp) / 1e6))

    # ------------------------------------------------------------------ 帧工具
    def wait_for_frame(self, after_mono, timeout):
        """等一帧 recv 时间 >= after_mono 的帧；超时返回当前最新帧。"""
        deadline = time.monotonic() + timeout
        with self._cv:
            while True:
                frame = self._latest
                if frame is not None and frame['recv_mono'] >= after_mono:
                    return dict(frame)
                left = deadline - time.monotonic()
                if left <= 0:
                    return dict(frame) if frame is not None else None
                self._cv.wait(min(left, 0.2))

    def _nearest_odom_locked(self, stamp_ns_value):
        if not self._odom_hist:
            return None
        return min(self._odom_hist, key=lambda r: abs(r['stamp_ns'] - stamp_ns_value))

    def frame_odom(self, frame):
        with self._lock:
            od = self._nearest_odom_locked(frame['cloud_stamp_ns'])
        if od is None:
            return None
        return {
            'stamp_ns': int(od['stamp_ns']),
            'position': [float(v) for v in od['pos']],
            'orientation_xyzw': [float(v) for v in od['quat']],
            'frame_id': od['frame_id'],
            'child_frame_id': od['child_frame_id'],
            'delta_ms': (od['stamp_ns'] - frame['cloud_stamp_ns']) / 1e6,
        }

    def encode_image(self, frame):
        """返回 (bytes, format)。压缩图直接转发；原图才编码。"""
        if frame.get('kind') == 'compressed':
            return frame['blob'], frame.get('format') or 'jpeg'

        img = frame['img']
        try:
            import cv2
            ok, buf = cv2.imencode(
                '.jpg', img,
                [int(cv2.IMWRITE_JPEG_QUALITY), int(self.args.jpeg_quality)])
            if ok:
                return buf.tobytes(), 'jpeg'
        except Exception:
            pass

        try:
            import io
            from PIL import Image as PILImage
            enc = frame.get('encoding', 'bgr8')
            rgb = img[..., ::-1] if enc.startswith('bgr') else img[..., :3]
            b = io.BytesIO()
            PILImage.fromarray(np.ascontiguousarray(rgb)).save(
                b, format='JPEG', quality=int(self.args.jpeg_quality))
            return b.getvalue(), 'jpeg'
        except Exception as exc:
            raise RuntimeError('原图编码失败，需要 OpenCV 或 Pillow：%s' % exc)

    # ------------------------------------------------------------------ 会话
    def _prune(self):
        now = time.monotonic()
        for sid in list(self._sessions):
            if now - self._sessions[sid]['last'] > self.args.session_ttl:
                del self._sessions[sid]

    def handle_hello(self, hdr):
        sid = hdr['session']['id']
        with self._lock:
            self._sessions.setdefault(sid, {
                'session': hdr['session'], 'frame': None,
                'target': None, 'target_info': None,
                'last': time.monotonic(),
            })
            self._sessions[sid]['last'] = time.monotonic()
        return {
            'ok': True, 'type': 'hello',
            'node': self.get_name(),
            'calib_width': CALIB_W, 'calib_height': CALIB_H,
            'board_can_compute': bool(_HAS_TARGET_MATH),
            'compute_modes': (['board', 'pc'] if _HAS_TARGET_MATH else ['pc']),
            'topics': {'image': self.args.image_topic,
                       'cloud': self.args.cloud_topic,
                       'odom': self.args.odom_topic},
        }

    def handle_capture(self, hdr):
        self._prune()
        t_req = time.monotonic()
        frame = self.wait_for_frame(t_req - self.args.max_frame_age,
                                    self.args.frame_timeout)
        if frame is None:
            return {'ok': False, 'type': 'frame',
                    'error': '还没有收到同步的图像+点云'}, b''

        blob, fmt = self.encode_image(frame)
        sid = hdr['session']['id']
        with self._lock:
            self._sessions[sid] = {
                'session': hdr['session'], 'frame': frame,
                'target': None, 'target_info': None,
                'last': time.monotonic(),
            }

        rep = {
            'ok': True, 'type': 'frame',
            'image_format': fmt,
            'image_bytes': len(blob),
            'image_width': frame.get('width'),
            'image_height': frame.get('height'),
            'calib_width': CALIB_W, 'calib_height': CALIB_H,
            'image_stamp_ns': int(frame['img_stamp_ns']),
            'cloud_stamp_ns': int(frame['cloud_stamp_ns']),
            'stamp_delta_ms': (frame['cloud_stamp_ns'] - frame['img_stamp_ns']) / 1e6,
            'n_points': int(frame['xyz'].shape[0]),
            'board_can_compute': bool(_HAS_TARGET_MATH),
            'odom': self.frame_odom(frame),
        }
        return rep, blob

    def handle_click(self, hdr):
        """板子端算：像素 -> p_cam / p_lidar / p_map，并把它记成目标。

        这是默认模式：点云不出板，电脑只发 (u, v)。
        """
        if not _HAS_TARGET_MATH:
            return {'ok': False, 'type': 'target',
                    'error': '板子没装 target_math.py，请电脑端用 --compute pc'}, b''

        sid = hdr['session']['id']
        with self._lock:
            s = self._sessions.get(sid)
            frame = None if s is None else s.get('frame')
        if frame is None:
            return {'ok': False, 'type': 'target', 'error': '先 capture'}, b''

        u = float(hdr['u'])
        v = float(hdr['v'])
        radius = float(hdr.get('radius', self.args.radius))
        fit = bool(hdr.get('fit_plane', True))

        # 每帧只建一次索引；同一帧后续点击只查局部格子
        with self._lock:
            index = None if s is None else s.get('index')
        if index is None:
            index = target_math.TargetIndex(frame['xyz'], conf=frame.get('conf'),
                                            cell=self.args.index_cell)
            with self._lock:
                if s is not None:
                    s['index'] = index

        res = index.locate(u, v, radius=radius, fit_plane=fit)
        out = {
            'ok': True, 'type': 'target',
            'found': bool(res['found']),
            'u': u, 'v': v,
            'n_valid': int(res['n_valid']),
            'n_cand': int(res['n_cand']),
            'plane_used': bool(res['plane_used']),
            'multi_layer': bool(res['multi_layer']),
        }
        if not res['found']:
            return out, b''

        out.update({
            'p_cam': [float(x) for x in res['p_cam']],
            'p_lidar': [float(x) for x in res['p_lidar']],
            'hit_u': float(res['hit_u']),
            'hit_v': float(res['hit_v']),
            'px_err': float(res['px_err']),
            'confidence': (None if res.get('confidence') is None
                           else float(res['confidence'])),
        })

        od = self.frame_odom(frame)
        if od is None:
            out['p_map'] = None
            out['error'] = '本帧没有里程计位姿'
            return out, b''

        R = quat_to_rot(od['orientation_xyzw'])
        t = np.asarray(od['position'], dtype=np.float64)
        pm = R @ np.asarray(res['p_lidar'], dtype=np.float64) + t

        out['p_map'] = [float(x) for x in pm]
        out['odom'] = od
        with self._lock:
            s['target'] = pm
            s['target_info'] = out
            s['last'] = time.monotonic()
        return out, b''

    def handle_cloud(self, hdr):
        sid = hdr['session']['id']
        with self._lock:
            s = self._sessions.get(sid)
            frame = None if s is None else s.get('frame')
        if frame is None:
            return {'ok': False, 'type': 'cloud', 'error': '先 capture'}, b''

        xyz = np.asarray(frame['xyz'], dtype='<f4')
        body = zlib.compress(xyz.tobytes(), 1)
        rep = {
            'ok': True, 'type': 'cloud',
            'encoding': 'zlib_f32',
            'n_points': int(xyz.shape[0]),
            'shape': [int(xyz.shape[0]), 3],
            'raw_bytes': int(xyz.nbytes),
            'zlib_bytes': int(len(body)),
            'cloud_stamp_ns': int(frame['cloud_stamp_ns']),
        }
        return rep, body

    def handle_set_target(self, hdr):
        sid = hdr['session']['id']
        with self._lock:
            s = self._sessions.get(sid)
            if s is None or s.get('frame') is None:
                return {'ok': False, 'type': 'target', 'error': '先 capture'}, b''
            pm = hdr.get('p_map')
            if not pm or len(pm) != 3:
                return {'ok': False, 'type': 'target',
                        'error': 'p_map 需要 3 个数'}, b''
            s['target'] = np.asarray([float(pm[0]), float(pm[1]), float(pm[2])],
                                     dtype=np.float64)
            s['target_info'] = {k: hdr.get(k) for k in ('u', 'v', 'p_cam', 'p_lidar', 'p_map')}
            s['last'] = time.monotonic()
            target = [float(v) for v in s['target']]
        return {'ok': True, 'type': 'target', 'target': target}, b''

    def handle_status(self, hdr):
        sid = hdr['session']['id']
        with self._lock:
            s = self._sessions.get(sid)
            od = self._odom_hist[-1] if self._odom_hist else None
            target = None if s is None else s.get('target')

        out = {'ok': True, 'type': 'status', 'has_target': target is not None}
        if od is None:
            out.update({'valid': False, 'error': '还没有里程计'})
            return out, b''

        age = time.monotonic() - od['recv_mono']
        pos = np.asarray(od['pos'], dtype=np.float64)
        out.update({
            'valid': bool(age <= self.args.odom_timeout),
            'pos': [float(v) for v in pos],
            'yaw_deg': float(np.rad2deg(quat_to_yaw(*od['quat']))),
            'odom_age_s': float(age),
            'odom_frame': od['frame_id'],
            'odom_child': od['child_frame_id'],
        })
        if target is not None:
            d = pos - np.asarray(target, dtype=np.float64)
            out['target'] = [float(v) for v in target]
            out['delta'] = [float(v) for v in d]   # 只给坐标差，不给距离
        return out, b''

    def handle_reset(self, hdr):
        sid = hdr['session']['id']
        with self._lock:
            s = self._sessions.get(sid)
            if s is not None:
                s['target'] = None
                s['target_info'] = None
                s['last'] = time.monotonic()
        return {'ok': True, 'type': 'reset'}, b''

    def dispatch(self, typ, hdr):
        if typ == 'hello':
            return self.handle_hello(hdr), b''
        if typ == 'capture':
            return self.handle_capture(hdr)
        if typ == 'click':
            return self.handle_click(hdr)
        if typ == 'cloud':
            return self.handle_cloud(hdr)
        if typ == 'set_target':
            return self.handle_set_target(hdr)
        if typ == 'status':
            return self.handle_status(hdr)
        if typ == 'reset':
            return self.handle_reset(hdr)
        return {'ok': False, 'type': typ, 'error': 'unknown type'}, b''

    # ------------------------------------------------------------------ 服务循环
    def serve(self):
        addr = 'tcp://0.0.0.0:%d' % self.args.port
        self._rep = pynng.Rep0(listen=addr)
        self._rep.recv_max_size = 32 * 1024 * 1024
        self._rep.recv_timeout = 500
        self.get_logger().info('pynng REP 监听 %s' % addr)

        while not self._stop:
            try:
                data = self._rep.recv()
            except pynng.Timeout:
                continue
            except pynng.Closed:
                break
            except Exception as exc:
                self.get_logger().error('recv 失败: %s' % exc)
                continue

            hdr = {}
            try:
                hdr, _ = wire.unpack(data)
                typ = hdr.get('type', '?')
                sess = hdr.get('session') or {}
                sid = sess.get('id', '?')
                reply, out_blob = self.dispatch(typ, hdr)
                reply['session'] = sess          # 原样带回，含 IP
                reply['server_ns'] = time.time_ns()
                self._rep.send(wire.pack(reply, out_blob))
                self.get_logger().info(
                    '[%s] %-10s from %s -> ok=%s'
                    % (sid[:8], typ, sess.get('client_ip', '?'), reply.get('ok')))
            except Exception as exc:
                try:
                    self._rep.send(wire.pack({
                        'ok': False, 'type': hdr.get('type', 'error'),
                        'error': str(exc),
                        'session': hdr.get('session') or {},
                        'server_ns': time.time_ns(),
                    }))
                except Exception:
                    pass
                self.get_logger().error('处理请求失败: %s' % exc)

    def stop(self):
        self._stop = True
        if self._rep is not None:
            try:
                self._rep.close()
            except Exception:
                pass


def parse_args():
    ap = argparse.ArgumentParser(description='会话式抓帧/定位/误差服务')
    ap.add_argument('--port', type=int, default=5555)
    ap.add_argument('--image-topic', default='/odin1/image/compressed')
    ap.add_argument('--cloud-topic', default='/odin1/cloud_raw')
    ap.add_argument('--odom-topic', default='/odin1/odometry')
    ap.add_argument('--no-compressed', action='store_true',
                    help='强制用原始 /odin1/image（需要 cv2 或 Pillow）')
    ap.add_argument('--slop', type=float, default=0.08, help='图像/点云允许时间差(秒)')
    ap.add_argument('--queue', type=int, default=20)
    ap.add_argument('--max-frame-age', type=float, default=0.35,
                    help='抓帧时允许复用多旧的帧(秒)')
    ap.add_argument('--frame-timeout', type=float, default=2.0,
                    help='等新帧的最长时间(秒)')
    ap.add_argument('--jpeg-quality', type=int, default=85)
    ap.add_argument('--radius', type=float, default=10.0,
                    help='click 请求没带半径时的默认候选像素半径')
    ap.add_argument('--index-cell', type=int, default=32,
                    help='像素格子大小(px)，用于每帧的投影索引')
    ap.add_argument('--odom-timeout', type=float, default=1.0,
                    help='里程计超过这个时间没更新就标记 invalid(秒)')
    ap.add_argument('--session-ttl', type=float, default=600.0)
    args = ap.parse_args()
    if args.no_compressed:
        args.image_topic = '/odin1/image'
    return args


def main():
    args = parse_args()
    rclpy.init()
    node = SessionServer(args)
    th = threading.Thread(target=node.serve, daemon=True)
    th.start()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.stop()
        node.destroy_node()
        if rclpy.ok():
            rclpy.shutdown()


if __name__ == '__main__':
    main()