#!/usr/bin/env python3
"""Convert Flexiv Elements .traj (text) -> arm_application YAML waypoints."""

from __future__ import annotations

import argparse
import math
import re
import sys
from pathlib import Path
from typing import Any


def parse_flexiv_traj(text: str) -> list[dict[str, Any]]:
    """Parse Elements Studio scatter-point .traj file."""
    blocks = re.split(r'pattern_params\s*\{', text)[1:]
    waypoints: list[dict[str, Any]] = []
    for block in blocks:
        name_m = re.search(r'pattern_name:\s*"([^"]+)"', block)
        vel_m = re.search(r'velocity:\s*"([^"]+)"', block)
        joints: dict[str, float] = {}
        for j in range(1, 8):
            m = re.search(rf'joint{j}:\s*"([^"]+)"', block)
            if m:
                joints[f'joint{j}'] = math.radians(float(m.group(1)))
        if len(joints) != 7:
            continue
        positions = [joints[f'joint{i}'] for i in range(1, 8)]
        waypoints.append({
            'name': name_m.group(1) if name_m else f'WPT{len(waypoints)}',
            'positions': positions,
            'velocity': float(vel_m.group(1)) if vel_m else 0.2,
        })
    return waypoints


def assign_times(waypoints: list[dict[str, Any]], min_seg: float = 3.0) -> None:
    """Cumulative time_from_start from joint-space segment length."""
    t = 0.0
    prev = None
    for wp in waypoints:
        pos = wp['positions']
        if prev is None:
            t += min_seg
        else:
            delta = max(abs(a - b) for a, b in zip(pos, prev))
            # ~4 s per radian max joint move, floored at min_seg
            t += max(min_seg, delta * 4.0)
        wp['time'] = round(t, 2)
        prev = pos


def to_yaml(
    waypoints: list[dict[str, Any]],
    *,
    joint_prefix: str = 'Rizon4-063989_',
    velocity_scale: float | None = None,
) -> str:
    if velocity_scale is None:
        velocity_scale = waypoints[0].get('velocity', 0.2) if waypoints else 0.2
    times = [wp['time'] for wp in waypoints]
    flat: list[float] = []
    for wp in waypoints:
        flat.extend(wp['positions'])
    pos_line = ', '.join(f'{v:.6f}' for v in flat)
    time_line = ', '.join(str(t) for t in times)
    names = ', '.join(f'  # {wp.get("name", "")}' if i == 0 else '' for i, wp in enumerate(waypoints))
    _ = names  # names kept in comments below
    lines = [
        '# Auto-converted from Flexiv Elements .traj (joints: deg -> rad)',
        '# ROS params: flat waypoint_positions + waypoint_times (no nested dicts)',
        '# Edit velocity_scale / loop_count as needed (no rebuild).',
        'run_joint_trajectory:',
        '  ros__parameters:',
        f'    joint_prefix: "{joint_prefix}"',
        '    joint_names: [joint1, joint2, joint3, joint4, joint5, joint6, joint7]',
        f'    velocity_scale: {velocity_scale}',
        '    loop_count: 1                # 1=once, 0=forever, N=repeat N',
        '    pause_between_loops: 0.5',
        f'    waypoint_times: [{time_line}]',
        '    waypoint_positions: [',
    ]
    for i, wp in enumerate(waypoints):
        pos = ', '.join(f'{v:.6f}' for v in wp['positions'])
        lines.append(f'      {pos},  # {wp.get("name", f"WPT{i}")}')
    lines.append('    ]')
    return '\n'.join(lines) + '\n'


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument('input', type=Path, help='Flexiv Elements .traj file')
    ap.add_argument('-o', '--output', type=Path, help='Output YAML path')
    ap.add_argument('--prefix', default='Rizon4-063989_', help='Joint name prefix (site SN + _)')
    ap.add_argument('--min-seg', type=float, default=3.0, help='Min seconds per segment')
    args = ap.parse_args()

    text = args.input.read_text(encoding='utf-8')
    wps = parse_flexiv_traj(text)
    if not wps:
        print('ERROR: no waypoints parsed', file=sys.stderr)
        return 1
    assign_times(wps, min_seg=args.min_seg)
    yaml_text = to_yaml(wps, joint_prefix=args.prefix)

    out = args.output
    if out:
        out.parent.mkdir(parents=True, exist_ok=True)
        out.write_text(yaml_text, encoding='utf-8')
        print(f'Wrote {len(wps)} waypoints -> {out}')
    else:
        print(yaml_text)
    return 0


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