#!/usr/bin/env python3
"""
generate_static_transforms.py — Generate static_transforms.launch.py from extrinsics.yaml.

Usage:
  python3 generate_static_transforms.py \\
    --extrinsics /path/to/calibration/extrinsics.yaml \\
    --output /path/to/jetson/ros2_ws/src/saltybot_bringup/launch/static_transforms.launch.py
"""

import argparse
import sys
from pathlib import Path

import yaml


def load_extrinsics(yaml_path: str) -> dict:
    """Load and parse extrinsics.yaml.

    Returns the 'cameras' sub-dict with per-camera transforms.
    """
    with open(yaml_path, 'r') as f:
        data = yaml.safe_load(f)
    if 'cameras' not in data:
        raise ValueError(f'extrinsics.yaml missing "cameras" key: {yaml_path}')
    return data['cameras'], data


def generate_launch_content(cameras: dict) -> str:
    """Generate static_transforms.launch.py content from cameras dict."""
    lines = [
        '"""',
        'static_transforms.launch.py — Auto-generated by generate_static_transforms.py',
        'DO NOT EDIT manually — regenerate with:',
        '  python3 generate_static_transforms.py \\',
        '    --extrinsics /path/to/calibration/extrinsics.yaml \\',
        '    --output <this_file>',
        '"""',
        '',
        'from launch import LaunchDescription',
        'from launch_ros.actions import Node',
        '',
        '',
        'def generate_launch_description():',
        '    return LaunchDescription([',
    ]

    for cam_name, cam_data in cameras.items():
        parent = cam_data.get('parent_frame', 'base_link')
        child = cam_data.get('child_frame', f'camera_{cam_name}_optical_frame')
        tx, ty, tz = cam_data['translation']
        qx, qy, qz, qw = cam_data['rotation_quat_xyzw']

        lines += [
            f'        # {cam_name}: {parent} -> {child}',
            f'        Node(',
            f'            package=\'tf2_ros\',',
            f'            executable=\'static_transform_publisher\',',
            f'            name=\'tf_base_to_camera_{cam_name}\',',
            f'            arguments=[',
            f'                \'{tx}\', \'{ty}\', \'{tz}\',',
            f'                \'{qx}\', \'{qy}\', \'{qz}\', \'{qw}\',',
            f'                \'{parent}\', \'{child}\',',
            f'            ],',
            f'        ),',
        ]

    # Also add RealSense D435i frame at origin (identity) — placeholder
    # (actual D435i extrinsics from its own calibration)
    lines += [
        '        # RealSense D435i: base_link -> camera_color_optical_frame',
        '        # NOTE: Replace with actual D435i extrinsics from stereo calibration',
        '        Node(',
        '            package=\'tf2_ros\',',
        '            executable=\'static_transform_publisher\',',
        '            name=\'tf_base_to_realsense\',',
        '            arguments=[',
        '                \'0.0\', \'0.0\', \'0.25\',',
        '                \'0.0\', \'0.0\', \'0.0\', \'1.0\',',
        '                \'base_link\', \'camera_color_optical_frame\',',
        '            ],',
        '        ),',
        '    ])',
        '',
    ]

    return '\n'.join(lines)


def main():
    parser = argparse.ArgumentParser(
        description='Generate static_transforms.launch.py from extrinsics.yaml.',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument(
        '--extrinsics', required=True,
        help='Path to extrinsics.yaml (produced by calibrate_extrinsics.py).',
    )
    parser.add_argument(
        '--output', required=True,
        help='Output path for static_transforms.launch.py.',
    )
    args = parser.parse_args()

    extrinsics_path = args.extrinsics
    if not Path(extrinsics_path).exists():
        print(f'[ERROR] extrinsics.yaml not found: {extrinsics_path}')
        sys.exit(1)

    print(f'Loading extrinsics: {extrinsics_path}')
    cameras, full_data = load_extrinsics(extrinsics_path)

    print(f'Found {len(cameras)} camera(s): {", ".join(cameras.keys())}')

    content = generate_launch_content(cameras)

    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, 'w') as f:
        f.write(content)

    print(f'Written: {output_path}')
    print()
    print('Transforms:')
    print(f'  {"Camera":<8}  {"Parent":<12}  {"Child frame":<36}  {"Translation"}')
    print('  ' + '-' * 80)
    for cam_name, cam_data in cameras.items():
        parent = cam_data.get('parent_frame', 'base_link')
        child = cam_data.get('child_frame', f'camera_{cam_name}_optical_frame')
        tx, ty, tz = cam_data['translation']
        print(f'  {cam_name:<8}  {parent:<12}  {child:<36}  '
              f'[{tx:.3f}, {ty:.3f}, {tz:.3f}]')
    print()
    print('Next step: include static_transforms.launch.py in your main launch file:')
    print('  from launch.actions import IncludeLaunchDescription')
    print('  from launch.launch_description_sources import PythonLaunchDescriptionSource')


if __name__ == '__main__':
    main()
