#!/usr/bin/env python3
"""
Generate a C++ header with YAML -> Goal/Request converters
based on src/brain/config/robot_skills.yaml and a sibling
'hivecore_robot_interfaces' package containing .action/.srv files.

Output: src/scripts/payload_converters.hpp

Usage:
  python3 gen_payload_converters.py [--skills path/to/robot_skills.yaml] [--interfaces /path/to/hivecore_robot_interfaces] [--out path]

If arguments omitted the script assumes:
  skills: ../src/brain/config/robot_skills.yaml (relative to repo root)
  interfaces: ../hivecore_robot_interfaces (sibling folder)
  out: src/scripts/payload_converters.hpp

This generator is conservative: it only auto-generates handling for primitive
and array primitives. For complex ROS message types it will emit a TODO comment.
"""

import argparse
import os
import re
import sys
import textwrap
import yaml

PRIMITIVE_MAP = {
    'int8': 'int8_t',
    'int16': 'int16_t',
    'float32': 'float',
    'float64': 'double',
    'float': 'double',
    'double': 'double',
    'uint8': 'uint8_t',
    'uint16': 'uint16_t',
    'int32': 'int32_t',
    'int64': 'int64_t',
    'uint32': 'uint32_t',
    'uint64': 'uint64_t',
    'int': 'int32_t',
    'string': 'std::string',
    'bool': 'bool',
    'time': 'rclcpp::Time',
}

HEADER_PREAMBLE = '''// GENERATED FILE - do not edit by hand
// Generated by src/scripts/gen_payload_converters.py
#pragma once

#include <yaml-cpp/yaml.h>
#include <string>
#include <vector>
#include <cstdint>
'''

NAMESPACE_BEGIN = '''
namespace brain {

template<typename T>
T from_yaml(const YAML::Node & n);

'''

HEADER_POSTAMBLE = '\n} // namespace brain\n'


def to_snake(s: str) -> str:
    out = ''
    for i,ch in enumerate(s):
        if ch.isupper():
            if i>0:
                out += '_'
            out += ch.lower()
        else:
            out += ch
    return out


def parse_action_file(path: str):
    """Parse an .action file, return dict with goal_fields: list of (type,name)
    """
    with open(path, 'r', encoding='utf-8') as f:
        content = f.read()
    parts = content.split('\n---\n')
    goal_part = parts[0] if parts else ''
    fields = []
    for line in goal_part.splitlines():
        line = line.strip()
        if not line or line.startswith('#'):
            continue
        # lines like: float32 move_pitch_degree
        m = re.match(r"^([A-Za-z0-9_/\[\]]+)\s+([A-Za-z0-9_]+).*$", line)
        if m:
            ty = m.group(1)
            name = m.group(2)
            fields.append((ty, name))
    return fields


def parse_srv_file(path: str):
    # request before ---
    with open(path, 'r', encoding='utf-8') as f:
        content = f.read()
    parts = content.split('\n---\n')
    req_part = parts[0] if parts else ''
    fields = []
    for line in req_part.splitlines():
        line = line.strip()
        if not line or line.startswith('#'):
            continue
        m = re.match(r"^([A-Za-z0-9_/\[\]]+)\s+([A-Za-z0-9_]+).*$", line)
        if m:
            ty = m.group(1)
            name = m.group(2)
            fields.append((ty, name))
    return fields


def find_file_case_insensitive(dirname: str, filename: str):
    """Return the path to a file in dirname matching filename case-insensitively, or None."""
    if not os.path.isdir(dirname):
        return None
    target = filename.lower()
    for e in os.listdir(dirname):
        if e.lower() == target:
            return os.path.join(dirname, e)
    return None


def make_cpp_type(ros_type: str):
    # arrays: float32[]
    is_array = ros_type.endswith('[]')
    base = ros_type[:-2] if is_array else ros_type
    if '/' in base:
        # complex message type
        cpp = None
    else:
        cpp = PRIMITIVE_MAP.get(base)
    return is_array, cpp, base


def snake_to_camel(s: str) -> str:
    # move_waist -> MoveWaist
    parts = re.split(r'[_\-]', s)
    return ''.join(p.capitalize() for p in parts if p)


def resolve_interface_subdir(root: str, sub: str):
    """Return the first existing directory for common layouts: root/sub, root/src/sub, root/interfaces/sub."""
    candidates = [
        os.path.join(root, sub),
        os.path.join(root, 'src', sub),
        os.path.join(root, 'interfaces', sub),
    ]
    for c in candidates:
        if os.path.isdir(c):
            return c
    # default to primary location even if missing
    return candidates[0]


def gen_converter_action(base: str, fields, include_path: str):
    # base: MoveWaist
    cpp = []
    cpp.append(f'// {base}::Goal: ' + ', '.join([n for _,n in fields]))
    cpp.append('template<>')
    cpp.append(f'inline interfaces::action::{base}::Goal from_yaml<interfaces::action::{base}::Goal>(const YAML::Node & n)')
    cpp.append('{')
    cpp.append(f'  interfaces::action::{base}::Goal g;')
    cpp.append('  if (n && n.IsMap()) {')
    has_data_array = any(n == 'data_array' for _, n in fields)
    has_data_length = any(n == 'data_length' for _, n in fields)
    for ty,name in fields:
        is_array, cpp_type, base_type = make_cpp_type(ty)
        if is_array and cpp_type:
            cpp.append(f'    if (n["{name}"]) {{')
            cpp.append(f'      auto vec = n["{name}"].as<std::vector<{cpp_type}>>();')
            cpp.append(f'      g.{name}.assign(vec.begin(), vec.end());')
            cpp.append('    }')
        elif not is_array and cpp_type:
            # map to as<type>()
            cpp.append(f'    if (n["{name}"]) g.{name} = n["{name}"].as<{cpp_type}>();')
        else:
            # complex types: handle some known messages
            # normalize geometry type names like geometry_msgs/msg/PoseStamped -> geometry_msgs/PoseStamped
            bt = base_type.replace('::', '/')
            bt = bt.replace('/msg/', '/')
            bt = bt.replace('/Msg/', '/')
            if bt in ("geometry_msgs/PoseStamped", "geometry_msgs/PoseStamped") and not is_array:
                cpp.append('    // geometry_msgs::msg::PoseStamped')
                cpp.append('    if (n) {')
                cpp.append('      if (n["frame_id"]) g.{0}.header.frame_id = n["frame_id"].as<std::string>();'.format(name))
                cpp.append('      if (n["stamp_sec"]) g.{0}.header.stamp.sec = n["stamp_sec"].as<int32_t>();'.format(name))
                cpp.append('      if (n["stamp_nanosec"]) g.{0}.header.stamp.nanosec = n["stamp_nanosec"].as<uint32_t>();'.format(name))
                cpp.append('      if (auto p = n["position"]) {')
                cpp.append('        if (p["x"]) g.{0}.pose.position.x = p["x"].as<double>();'.format(name))
                cpp.append('        if (p["y"]) g.{0}.pose.position.y = p["y"].as<double>();'.format(name))
                cpp.append('        if (p["z"]) g.{0}.pose.position.z = p["z"].as<double>();'.format(name))
                cpp.append('      }')
                cpp.append('      if (auto q = n["orientation"]) {')
                cpp.append('        if (q["x"]) g.{0}.pose.orientation.x = q["x"].as<double>();'.format(name))
                cpp.append('        if (q["y"]) g.{0}.pose.orientation.y = q["y"].as<double>();'.format(name))
                cpp.append('        if (q["z"]) g.{0}.pose.orientation.z = q["z"].as<double>();'.format(name))
                cpp.append('        if (q["w"]) g.{0}.pose.orientation.w = q["w"].as<double>();'.format(name))
                cpp.append('      }')
                cpp.append('    }')
            elif bt in ("geometry_msgs/TwistStamped", "geometry_msgs/TwistStamped") and not is_array:
                cpp.append('    // geometry_msgs::msg::TwistStamped')
                cpp.append('    if (n) {')
                cpp.append('      if (n["frame_id"]) g.{0}.header.frame_id = n["frame_id"].as<std::string>();'.format(name))
                cpp.append('      if (n["stamp_sec"]) g.{0}.header.stamp.sec = n["stamp_sec"].as<int32_t>();'.format(name))
                cpp.append('      if (n["stamp_nanosec"]) g.{0}.header.stamp.nanosec = n["stamp_nanosec"].as<uint32_t>();'.format(name))
                cpp.append('      if (auto lin = n["linear"]) {')
                cpp.append('        if (lin["x"]) g.{0}.twist.linear.x = lin["x"].as<double>();'.format(name))
                cpp.append('        if (lin["y"]) g.{0}.twist.linear.y = lin["y"].as<double>();'.format(name))
                cpp.append('        if (lin["z"]) g.{0}.twist.linear.z = lin["z"].as<double>();'.format(name))
                cpp.append('      }')
                cpp.append('      if (auto ang = n["angular"]) {')
                cpp.append('        if (ang["x"]) g.{0}.twist.angular.x = ang["x"].as<double>();'.format(name))
                cpp.append('        if (ang["y"]) g.{0}.twist.angular.y = ang["y"].as<double>();'.format(name))
                cpp.append('        if (ang["z"]) g.{0}.twist.angular.z = ang["z"].as<double>();'.format(name))
                cpp.append('      }')
                cpp.append('    }')
            else:
                cpp.append(f'    // TODO: parse field `{name}` of ROS type `{ty}` into g.{name} (complex type)')
    # post-processing example: if both data_array and data_length exist, set length from array size
    if has_data_array and has_data_length:
        cpp.append('  if (!g.data_array.empty()) {')
        cpp.append('    using DL = decltype(g.data_length);')
        cpp.append('    g.data_length = static_cast<DL>(g.data_array.size());')
        cpp.append('  }')
    cpp.append('  }')
    cpp.append('  return g;')
    cpp.append('}\n')
    return '\n'.join(cpp)


def gen_converter_srv(base: str, fields, include_path: str):
    cpp = []
    cpp.append(f'// {base}::Request: ' + ', '.join([n for _,n in fields]))
    cpp.append('template<>')
    cpp.append(f'inline interfaces::srv::{base}::Request from_yaml<interfaces::srv::{base}::Request>(const YAML::Node & n)')
    cpp.append('{')
    cpp.append(f'  interfaces::srv::{base}::Request r;')
    cpp.append('  if (n && n.IsMap()) {')
    for ty,name in fields:
        is_array, cpp_type, base_type = make_cpp_type(ty)
        if is_array and cpp_type:
            cpp.append(f'    if (n["{name}"]) {{')
            cpp.append(f'      auto vec = n["{name}"].as<std::vector<{cpp_type}>>();')
            cpp.append(f'      r.{name}.assign(vec.begin(), vec.end());')
            cpp.append('    }')
        elif not is_array and cpp_type:
            cpp.append(f'    if (n["{name}"]) r.{name} = n["{name}"].as<{cpp_type}>();')
        else:
            bt = base_type.replace('::', '/')
            bt = bt.replace('/msg/', '/')
            bt = bt.replace('/Msg/', '/')
            if bt in ("geometry_msgs/PoseStamped", "geometry_msgs/PoseStamped") and not is_array:
                cpp.append('    // geometry_msgs::msg::PoseStamped')
                cpp.append('    if (n) {')
                cpp.append('      if (n["frame_id"]) r.{0}.header.frame_id = n["frame_id"].as<std::string>();'.format(name))
                cpp.append('      if (n["stamp_sec"]) r.{0}.header.stamp.sec = n["stamp_sec"].as<int32_t>();'.format(name))
                cpp.append('      if (n["stamp_nanosec"]) r.{0}.header.stamp.nanosec = n["stamp_nanosec"].as<uint32_t>();'.format(name))
                cpp.append('      if (auto p = n["position"]) {')
                cpp.append('        if (p["x"]) r.{0}.pose.position.x = p["x"].as<double>();'.format(name))
                cpp.append('        if (p["y"]) r.{0}.pose.position.y = p["y"].as<double>();'.format(name))
                cpp.append('        if (p["z"]) r.{0}.pose.position.z = p["z"].as<double>();'.format(name))
                cpp.append('      }')
                cpp.append('      if (auto q = n["orientation"]) {')
                cpp.append('        if (q["x"]) r.{0}.pose.orientation.x = q["x"].as<double>();'.format(name))
                cpp.append('        if (q["y"]) r.{0}.pose.orientation.y = q["y"].as<double>();'.format(name))
                cpp.append('        if (q["z"]) r.{0}.pose.orientation.z = q["z"].as<double>();'.format(name))
                cpp.append('        if (q["w"]) r.{0}.pose.orientation.w = q["w"].as<double>();'.format(name))
                cpp.append('      }')
                cpp.append('    }')
            elif bt in ("geometry_msgs/TwistStamped", "geometry_msgs/TwistStamped") and not is_array:
                cpp.append('    // geometry_msgs::msg::TwistStamped')
                cpp.append('    if (n) {')
                cpp.append('      if (n["frame_id"]) r.{0}.header.frame_id = n["frame_id"].as<std::string>();'.format(name))
                cpp.append('      if (n["stamp_sec"]) r.{0}.header.stamp.sec = n["stamp_sec"].as<int32_t>();'.format(name))
                cpp.append('      if (n["stamp_nanosec"]) r.{0}.header.stamp.nanosec = n["stamp_nanosec"].as<uint32_t>();'.format(name))
                cpp.append('      if (auto lin = n["linear"]) {')
                cpp.append('        if (lin["x"]) r.{0}.twist.linear.x = lin["x"].as<double>();'.format(name))
                cpp.append('        if (lin["y"]) r.{0}.twist.linear.y = lin["y"].as<double>();'.format(name))
                cpp.append('        if (lin["z"]) r.{0}.twist.linear.z = lin["z"].as<double>();'.format(name))
                cpp.append('      }')
                cpp.append('      if (auto ang = n["angular"]) {')
                cpp.append('        if (ang["x"]) r.{0}.twist.angular.x = ang["x"].as<double>();'.format(name))
                cpp.append('        if (ang["y"]) r.{0}.twist.angular.y = ang["y"].as<double>();'.format(name))
                cpp.append('        if (ang["z"]) r.{0}.twist.angular.z = ang["z"].as<double>();'.format(name))
                cpp.append('      }')
                cpp.append('    }')
            else:
                cpp.append(f'    // TODO: parse field `{name}` of ROS type `{ty}` into r.{name} (complex type)')
    cpp.append('  }')
    cpp.append('  return r;')
    cpp.append('}\n')
    return '\n'.join(cpp)


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--skills', default='src/brain/config/robot_skills.yaml')
    p.add_argument('--interfaces', default='../hivecore_robot_interfaces')
    p.add_argument('--out', default='src/brain/include/brain/payload_converters.hpp')
    args = p.parse_args()

    repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..'))
    skills_path = args.skills if os.path.isabs(args.skills) else os.path.join(repo_root, args.skills)
    interfaces_root = args.interfaces if os.path.isabs(args.interfaces) else os.path.join(repo_root, args.interfaces)
    out_path = args.out if os.path.isabs(args.out) else os.path.join(repo_root, args.out)

    if not os.path.exists(skills_path):
        print('ERROR: skills file not found:', skills_path, file=sys.stderr)
        sys.exit(2)

    print('Using skills:', skills_path)
    print('Using interfaces root:', interfaces_root)
    print('Writing output to:', out_path)

    with open(skills_path, 'r', encoding='utf-8') as f:
        skills = yaml.safe_load(f)

    includes = set()
    converters = []
    added_actions = set()
    added_srvs = set()

    for item in skills:
        if not isinstance(item, dict):
            continue
        name = item.get('name')
        interfaces = item.get('interfaces', [])
        for iface in interfaces:
            token = None
            if isinstance(iface, str):
                token = iface
            elif isinstance(iface, dict):
                token = iface.get('name')
            if not token:
                continue
            # token like BaseName.action
            if '.' not in token:
                continue
            base, kind = token.split('.', 1)
            kind = kind.lower()
            header_snake = to_snake(base)
            if kind.startswith('action'):
                includes.add(f'"interfaces/action/{header_snake}.hpp"')
                # locate .action file in common subdirs
                action_dir = resolve_interface_subdir(interfaces_root, 'action')
                tried_paths = []
                action_file = os.path.join(action_dir, base + '.action')
                tried_paths.append(action_file)
                if not os.path.exists(action_file):
                    camel = snake_to_camel(base)
                    if camel != base:
                        action_file_camel = os.path.join(action_dir, camel + '.action')
                        tried_paths.append(action_file_camel)
                        if os.path.exists(action_file_camel):
                            action_file = action_file_camel
                if not os.path.exists(action_file):
                    found = find_file_case_insensitive(action_dir, base + '.action')
                    if found:
                        action_file = found
                if not os.path.exists(action_file):
                    # final fallback to snake name in case-insensitive form
                    snake_path = os.path.join(action_dir, header_snake + '.action')
                    tried_paths.append(snake_path)
                    if os.path.exists(snake_path):
                        action_file = snake_path
                if not os.path.exists(action_file):
                    print('WARN: action file not found for', base, 'checked', tried_paths[-1])
                    fields = []
                else:
                    fields = parse_action_file(action_file)
                converters.append(gen_converter_action(base, fields, header_snake))
                added_actions.add(base)
            elif kind.startswith('srv') or kind.startswith('service'):
                includes.add(f'"interfaces/srv/{header_snake}.hpp"')
                srv_dir = resolve_interface_subdir(interfaces_root, 'srv')
                tried_paths = []
                srv_file = os.path.join(srv_dir, base + '.srv')
                tried_paths.append(srv_file)
                if not os.path.exists(srv_file):
                    # try snake_case -> CamelCase fallback
                    camel = snake_to_camel(base)
                    if camel != base:
                        srv_file_camel = os.path.join(srv_dir, camel + '.srv')
                        tried_paths.append(srv_file_camel)
                        if os.path.exists(srv_file_camel):
                            srv_file = srv_file_camel
                    if srv_file == os.path.join(srv_dir, base + '.srv'):
                        found = find_file_case_insensitive(srv_dir, base + '.srv')
                        if found:
                            srv_file = found
                        else:
                            srv_file = os.path.join(srv_dir, header_snake + '.srv')
                            tried_paths.append(srv_file)
                if not os.path.exists(srv_file):
                    print('WARN: srv file not found for', base, 'checked', tried_paths[-1])
                    fields = []
                else:
                    fields = parse_srv_file(srv_file)
                converters.append(gen_converter_srv(base, fields, header_snake))
                added_srvs.add(base)
            # end of iface loop
            pass

    # Second pass: also generate for all action/srv files in interfaces even if not referenced in skills
    action_dir = resolve_interface_subdir(interfaces_root, 'action')
    if os.path.isdir(action_dir):
        for fname in os.listdir(action_dir):
            if not fname.endswith('.action'):
                continue
            base_name = os.path.splitext(fname)[0]
            # preserve filename's CamelCase as type, only use to_snake for header path
            base_type = base_name
            if base_type in added_actions:
                continue
            header_snake = to_snake(base_type)
            includes.add(f'"interfaces/action/{header_snake}.hpp"')
            fields = parse_action_file(os.path.join(action_dir, fname))
            converters.append(gen_converter_action(base_type, fields, header_snake))
            added_actions.add(base_type)

    srv_dir = resolve_interface_subdir(interfaces_root, 'srv')
    if os.path.isdir(srv_dir):
        for fname in os.listdir(srv_dir):
            if not fname.endswith('.srv'):
                continue
            base_name = os.path.splitext(fname)[0]
            base_type = base_name
            if base_type in added_srvs:
                continue
            header_snake = to_snake(base_type)
            includes.add(f'"interfaces/srv/{header_snake}.hpp"')
            fields = parse_srv_file(os.path.join(srv_dir, fname))
            converters.append(gen_converter_srv(base_type, fields, header_snake))
            added_srvs.add(base_type)

    # assemble header
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    with open(out_path, 'w', encoding='utf-8') as outf:
        outf.write(HEADER_PREAMBLE)
        if includes:
            for inc in sorted(includes):
                outf.write(f'#include {inc}\n')
            outf.write('\n')
        outf.write(NAMESPACE_BEGIN)
        outf.write('\n'.join(converters))
        outf.write(HEADER_POSTAMBLE)

    print('Done.')

if __name__ == '__main__':
    main()
