#!/usr/bin/env python3
"""拍一张顶视 D435，把 k0 用眼在手外标定变到左臂基座。

stdout 只打一行，给 run_all.sh --d435-only 接厂家 MoveL：

    K0_BASE x y z

x y z 单位米，坐标系是这条臂的基座（与 RDK 的 WORLD 同一套）。
变换就是 handeye/results/<臂>_d435_eye_to_hand.yaml 的 T_4x4：

    p_基座 = T_cam_to_base · p_相机

k0 的 z 是工件表面，不是工具尖该去的高度。运动脚本会保持当前尖高，只用这里的 xy。
没有可信 k0 就退出，不用检测框中心冒充。
"""

from __future__ import annotations

import sys
import time
from pathlib import Path

import numpy as np
import rclpy
from rclpy.node import Node
from rclpy.qos import qos_profile_sensor_data
from sensor_msgs.msg import CameraInfo, Image

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "scripts" / "lib"))

from export_site_env import load_site, perception_overrides  # noqa: E402
from arm_application.geom_ik import load_handeye_T  # noqa: E402
from arm_application.overhead_oneshot import (  # noqa: E402
    depth_raw_to_m,
    oneshot_k0_base,
    select_oneshot_det,
)
from arm_application.perception_node import _bgr_copy  # noqa: E402
from arm_application.yolo_detector import build_detector  # noqa: E402


def _decode_depth(msg: Image):
    w, h = int(msg.width), int(msg.height)
    enc = (msg.encoding or "").lower()
    if enc in ("16uc1", "mono16"):
        raw = np.frombuffer(msg.data, dtype=np.uint16)
        if raw.size < w * h:
            return None
        return raw[: w * h].reshape(h, w)
    if enc == "32fc1":
        raw = np.frombuffer(msg.data, dtype=np.float32)
        if raw.size < w * h:
            return None
        return raw[: w * h].reshape(h, w)
    return None


class OneShot(Node):
    def __init__(self, topics: dict) -> None:
        super().__init__("d435_k0_once")
        self.color = None
        self.depth = None
        self.info = None
        self.create_subscription(Image, topics["image"], self._on_color, qos_profile_sensor_data)
        self.create_subscription(Image, topics["depth"], self._on_depth, qos_profile_sensor_data)
        self.create_subscription(CameraInfo, topics["info"], self._on_info, qos_profile_sensor_data)

    def _on_color(self, msg: Image) -> None:
        self.color = msg

    def _on_depth(self, msg: Image) -> None:
        self.depth = msg

    def _on_info(self, msg: CameraInfo) -> None:
        self.info = msg

    def ready(self) -> bool:
        return self.color is not None and self.depth is not None and self.info is not None


def main() -> int:
    perc = perception_overrides(load_site())
    if not perc.get("overhead_oneshot_enable"):
        print("site 没有接上 D435 一次性（overhead_oneshot_enable）", file=sys.stderr)
        return 2
    handeye = str(perc.get("overhead_handeye_file") or "")
    topics = {
        "image": str(perc.get("overhead_image_topic") or ""),
        "depth": str(perc.get("overhead_depth_topic") or ""),
        "info": str(perc.get("overhead_info_topic") or ""),
    }
    if not handeye or not all(topics.values()):
        print(f"眼在手外话题或标定文件没接上：{handeye} {topics}", file=sys.stderr)
        return 2
    T = load_handeye_T(handeye)
    print(f"手眼 {handeye}", file=sys.stderr)
    print(f"图 {topics['image']}", file=sys.stderr)

    weights = str(perc.get("weights") or "")
    detector = build_detector(
        weights,
        backend=str(perc.get("backend") or "auto"),
        conf=float(perc.get("conf") or 0.1),
        iou=float(perc.get("iou") or 0.45),
        device=str(perc.get("device") or "0"),
        imgsz=int(perc.get("imgsz") or 640),
    )

    rclpy.init()
    node = OneShot(topics)
    t0 = time.monotonic()
    try:
        while rclpy.ok() and not node.ready():
            rclpy.spin_once(node, timeout_sec=0.2)
            if time.monotonic() - t0 > 15.0:
                print("15s 内没有凑齐彩色图、深度、内参", file=sys.stderr)
                return 3
        color, depth, info = node.color, node.depth, node.info
    finally:
        node.destroy_node()
        if rclpy.ok():
            rclpy.shutdown()

    bgr = _bgr_copy(color)
    dets = detector.detect(bgr)
    min_conf = float(perc["overhead_oneshot_conf"])
    det = select_oneshot_det(dets, min_conf)
    if det is None:
        print(f"画面里没有 conf≥{min_conf:.2f} 的工件", file=sys.stderr)
        return 4
    raw = _decode_depth(depth)
    depth_m = depth_raw_to_m(raw, depth.encoding or "")
    if depth_m is None:
        print("深度无法转成米", file=sys.stderr)
        return 4
    k = info.k
    got = oneshot_k0_base(
        det,
        depth_m,
        fx=float(k[0]),
        fy=float(k[4]),
        cx=float(k[2]),
        cy=float(k[5]),
        T_cam_to_base=T,
        min_z_m=float(perc["overhead_min_depth_m"]),
        max_z_m=float(perc["overhead_max_depth_m"]),
        kpt_index=int(perc.get("keypoint_index") or 0),
        kpt_min_conf=float(perc.get("keypoint_min_conf") or 0.30),
        min_det_conf=min_conf,
        color_wh=(int(color.width), int(color.height)),
        depth_wh=(int(depth_m.shape[1]), int(depth_m.shape[0])),
    )
    if got is None:
        print("k0 或深度不可用，不发点", file=sys.stderr)
        return 4
    p_base, is_k0 = got
    if not is_k0:
        print(
            f"顶视没有可信 k0，只有框心 ({p_base[0]:.3f},{p_base[1]:.3f},{p_base[2]:.3f})，拒绝",
            file=sys.stderr,
        )
        return 4
    print(
        f"k0 基座 ({p_base[0]:.4f}, {p_base[1]:.4f}, {p_base[2]:.4f}) m  conf={float(det.conf):.2f}",
        file=sys.stderr,
    )
    print(f"K0_BASE {p_base[0]:.6f} {p_base[1]:.6f} {p_base[2]:.6f}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())