#!/usr/bin/env python3
"""Merge taught joint lines into handeye waypoints YAML.

Input (from handeye_teach / handeye_waypoints append file):
  # pose 0
  0.123456, -0.700000, ...

Target can be either:
  - handeye/poses/<臂号>_<模式>.yaml  (/** ros__parameters)
  - handeye/calibrate.yaml           (handeye_calibrate: ros__parameters)

Usage:
  python3 scripts/merge_handeye_poses.py \\
    ~/ros2/handeye/tmp/89/taught_poses.txt \\
    ~/ros2/handeye/poses/89_eye_in_hand.yaml

  # or via wrapper:
  bash ~/ros2/scripts/run_handeye.sh merge
"""

from __future__ import annotations

import re
import sys
from pathlib import Path


def parse_poses(path: Path) -> list[list[float]]:
    lines = path.read_text(encoding="utf-8").splitlines()
    poses: list[list[float]] = []
    buf: list[float] = []

    num_re = re.compile(r"[-+]?(?:\d+\.\d*|\.\d+|\d+)(?:[eE][-+]?\d+)?")

    for raw in lines:
        line = raw.split("#", 1)[0].strip()
        if not line:
            if buf:
                if len(buf) != 7:
                    raise ValueError(f"Expected 7 joints, got {len(buf)}: {buf}")
                poses.append(buf)
                buf = []
            continue
        vals = [float(x) for x in num_re.findall(line)]
        if not vals:
            continue
        if len(vals) == 7:
            poses.append(vals)
        elif len(vals) < 7:
            buf.extend(vals)
            if len(buf) > 7:
                raise ValueError(f"Too many joint values in line: {line!r}")
        else:
            poses.append(vals[:7])

    if buf:
        if len(buf) != 7:
            raise ValueError(f"Expected 7 joints at EOF, got {len(buf)}")
        poses.append(buf)

    return poses


def format_positions_block(poses: list[list[float]], indent: str = "    ") -> str:
    rows = [f"{indent}waypoint_positions: ["]
    for i, q in enumerate(poses):
        rows.append(f"{indent}  # pose {i}")
        s = ", ".join(f"{v:.6f}" for v in q) + ","
        rows.append(f"{indent}  {s}")
    rows.append(f"{indent}]")
    return "\n".join(rows)


def merge(yaml_path: Path, poses: list[list[float]]) -> None:
    text = yaml_path.read_text(encoding="utf-8").replace("\r\n", "\n").replace("\r", "\n")
    # 探测 waypoint_positions 的缩进（通常 2 空格或 4 空格）
    m_ind = re.search(r"^([ \t]*)waypoint_positions:\s*\[", text, re.MULTILINE)
    indent = m_ind.group(1) if m_ind else "    "
    block = format_positions_block(poses, indent=indent)

    pattern = re.compile(
        rf"{re.escape(indent)}waypoint_positions:\s*\[[\s\S]*?{re.escape(indent)}\]",
        re.MULTILINE,
    )
    # 也要匹配"单行空列表"的写法：waypoint_positions: []
    one_line = re.compile(
        rf"{re.escape(indent)}waypoint_positions:\s*\[\s*\]",
        re.MULTILINE,
    )
    if pattern.search(text):
        new_text = pattern.sub(block, text, count=1)
    elif one_line.search(text):
        new_text = one_line.sub(block, text, count=1)
    else:
        raise RuntimeError("waypoint_positions block not found in YAML")

    # 顺手把 num_poses 同步成实际点数（如果原本就有这个键）
    n = len(poses)
    if re.search(r"^\s*num_poses:\s*\d+", new_text, re.MULTILINE):
        new_text = re.sub(
            r"^(\s*num_poses:\s*)\d+",
            rf"\g<1>{n}",
            new_text,
            count=1,
            flags=re.MULTILINE,
        )
    else:
        # 原本没有 num_poses，就插在 waypoint_positions 前面
        new_text = new_text.replace(
            block,
            f"{indent}num_poses: {n}\n{block}",
            1,
        )

    yaml_path.write_text(new_text, encoding="utf-8")


def main(argv: list[str]) -> int:
    if len(argv) != 3:
        print(
            "用法：merge_handeye_poses.py <taught_poses.txt> <waypoints.yaml>",
            file=sys.stderr,
        )
        return 1
    src = Path(argv[1]).expanduser()
    dst = Path(argv[2]).expanduser()
    if not src.is_file():
        print(f"Not found: {src}", file=sys.stderr)
        return 2
    if not dst.is_file():
        print(f"Not found: {dst}", file=sys.stderr)
        return 3
    poses = parse_poses(src)
    if len(poses) < 5:
        print(f"Need >=5 poses, got {len(poses)}", file=sys.stderr)
        return 4
    merge(dst, poses)
    print(f"Merged {len(poses)} poses -> {dst}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv))