#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
_mock_server.py — 仅用于在电脑上测试协议，不需要 ROS。

它模拟板子端：返回一张合成图像、一片合成点云（3 m 处、像素 800,648 附近），
以及一个会随时间变化的位姿，用来验证 session_client 的抓帧、点选、回传、Δ 显示。
"""
import io
import threading
import time
import zlib

import numpy as np
from PIL import Image
import pynng

import calib
import target_math
import wire

CAL_W, CAL_H = calib.IMAGE_WIDTH, calib.IMAGE_HEIGHT
U0, V0 = 800.0, 648.0
DEPTH = 3.0


def build_cloud():
    pts = []
    for du in range(-12, 13, 4):
        for dv in range(-12, 13, 4):
            ray = calib.pixel_to_ray(U0 + du, V0 + dv)
            p_cam = ray * DEPTH
            pts.append(calib.cam_to_lidar(p_cam)[0])
    pts = np.asarray(pts, dtype=np.float32)
    # 加一些无效点，模拟真实点云的零值
    zeros = np.zeros((97, 3), dtype=np.float32)
    return np.concatenate([pts, zeros], axis=0)


def build_image():
    arr = np.zeros((CAL_H, CAL_W, 3), dtype=np.uint8)
    arr[:, :, 2] = 40
    u, v = int(round(U0)), int(round(V0))
    arr[v - 6:v + 6, u - 6:u + 6] = (255, 255, 255)
    b = io.BytesIO()
    Image.fromarray(arr).save(b, format='JPEG', quality=85)
    return b.getvalue()


class MockServer:
    def __init__(self, port=5599):
        self.port = port
        self.xyz = build_cloud()
        self.image = build_image()
        self.target = None
        self.index = None
        self.sessions = {}
        self.stop_event = threading.Event()
        self.thread = None

    def start(self):
        self.thread = threading.Thread(target=self._serve, daemon=True)
        self.thread.start()

    def stop(self):
        self.stop_event.set()
        if self.thread:
            self.thread.join(timeout=2)

    def _odom(self):
        return {
            'stamp_ns': time.time_ns(),
            'position': [0.0, 0.0, 0.0],
            'orientation_xyzw': [0.0, 0.0, 0.0, 1.0],
            'frame_id': 'odom',
            'child_frame_id': 'imu',
            'delta_ms': 1.0,
        }

    def _handle(self, hdr):
        typ = hdr.get('type')
        if typ == 'hello':
            return {'ok': True, 'type': 'hello', 'node': 'mock',
                    'calib_width': CAL_W, 'calib_height': CAL_H,
                    'board_can_compute': True,
                    'compute_modes': ['board', 'pc'],
                    'topics': {'image': 'mock', 'cloud': 'mock', 'odom': 'mock'}}, b''
        if typ == 'capture':
            sid = (hdr.get('session') or {}).get('id')
            self.sessions[sid] = {'target': None}
            return {'ok': True, 'type': 'frame', 'image_format': 'jpeg',
                    'image_bytes': len(self.image),
                    'image_width': CAL_W, 'image_height': CAL_H,
                    'calib_width': CAL_W, 'calib_height': CAL_H,
                    'image_stamp_ns': time.time_ns(),
                    'cloud_stamp_ns': time.time_ns(),
                    'stamp_delta_ms': 1.2,
                    'n_points': int(self.xyz.shape[0]),
                    'board_can_compute': True,
                    'odom': self._odom()}, self.image
        if typ == 'click':
            sid = (hdr.get('session') or {}).get('id')
            s = self.sessions.get(sid)
            if s is None:
                return {'ok': False, 'type': 'target', 'error': '先 capture'}, b''
            if self.index is None:
                self.index = target_math.TargetIndex(self.xyz)
            u = float(hdr['u'])
            v = float(hdr['v'])
            res = self.index.locate(u, v,
                                    radius=float(hdr.get('radius', 10.0)),
                                    fit_plane=bool(hdr.get('fit_plane', True)))
            out = {'ok': True, 'type': 'target', 'found': bool(res['found']),
                   'u': u, 'v': v, 'n_cand': int(res['n_cand']),
                   'plane_used': bool(res['plane_used']),
                   'multi_layer': bool(res['multi_layer'])}
            if res['found']:
                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,
                })
                pm = np.asarray(res['p_lidar'], dtype=np.float64)
                out['p_map'] = [float(x) for x in pm]
                out['odom'] = self._odom()
                s['target'] = pm
            return out, b''
        if typ == 'cloud':
            body = zlib.compress(self.xyz.astype('<f4').tobytes(), 1)
            return {'ok': True, 'type': 'cloud', 'encoding': 'zlib_f32',
                    'n_points': int(self.xyz.shape[0]),
                    'shape': [int(self.xyz.shape[0]), 3],
                    'raw_bytes': int(self.xyz.nbytes),
                    'zlib_bytes': len(body)}, body
        if typ == 'set_target':
            sid = (hdr.get('session') or {}).get('id')
            s = self.sessions.get(sid)
            if s 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(pm, dtype=np.float64)
            return {'ok': True, 'type': 'target',
                    'target': [float(v) for v in s['target']]}, b''
        if typ == 'status':
            sid = (hdr.get('session') or {}).get('id')
            s = self.sessions.get(sid)
            target = None if s is None else s.get('target')
            out = {'ok': True, 'type': 'status',
                   'has_target': target is not None,
                   'valid': True, 'yaw_deg': 0.0, 'odom_age_s': 0.01,
                   'odom_frame': 'odom', 'odom_child': 'imu'}
            # 模拟：人把板子往前多走了 0.25 m
            pos = np.array([0.2507, -0.0209, 0.0144])
            out['pos'] = [float(v) for v in pos]
            if target is not None:
                d = pos - target
                out['target'] = [float(v) for v in target]
                out['delta'] = [float(v) for v in d]
            return out, b''
        if typ == 'reset':
            sid = (hdr.get('session') or {}).get('id')
            s = self.sessions.get(sid)
            if s is not None:
                s['target'] = None
            return {'ok': True, 'type': 'reset'}, b''
        return {'ok': False, 'type': typ, 'error': 'unknown type'}, b''

    def _serve(self):
        with pynng.Rep0(listen='tcp://127.0.0.1:%d' % self.port,
                        recv_timeout=200) as rep:
            rep.recv_max_size = 32 * 1024 * 1024
            while not self.stop_event.is_set():
                try:
                    data = rep.recv()
                except pynng.Timeout:
                    continue
                except pynng.Closed:
                    break
                hdr, _ = wire.unpack(data)
                reply, blob = self._handle(hdr)
                reply['session'] = hdr.get('session') or {}
                reply['server_ns'] = time.time_ns()
                rep.send(wire.pack(reply, blob))


if __name__ == '__main__':
    srv = MockServer()
    srv.start()
    print('mock server on 127.0.0.1:5599 (Ctrl+C 退出)')
    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        srv.stop()