#!/usr/bin/env python3
"""
Code generator: reads robot_skills.yaml and generates
- src/brain/include/brain/skill_manager.hpp
- src/brain/src/skill_manager.cpp

The generated C++ keeps the existing SkillManager API but
instantiates concrete registrars based on YAML content, avoiding
hardcoded if/else comparisons in C++.

Usage:
  python3 src/scripts/gen_skill_manager.py \
    --yaml src/brain/config/robot_skills.yaml \
    --include-dir src/brain/include/brain \
    --src-dir src/brain/src

Notes:
- Topic naming (current): only CamelCase skill name is registered (no snake_case duplicate) to avoid redundant action/service topics.
- BT node name (if used elsewhere): <SkillName>_H.
- Interface entry format: BaseName.action | BaseName.service
- C++ header includes Chinese comments.
"""
import argparse
import sys
from pathlib import Path
from string import Template
from typing import Dict

try:
    import yaml
except Exception as e:
    print("Missing dependency: PyYAML. Install with: pip install pyyaml", file=sys.stderr)
    raise

HEADER_TEMPLATE = """#pragma once

#include <rclcpp/rclcpp.hpp>
#include <yaml-cpp/yaml.h>
#include <string>
#include <unordered_map>
#include <vector>
#include <memory>

#include \"brain/action_registry.hpp\"
#include \"brain/bt_registry.hpp\"

namespace brain
{

// SkillSpec 用于描述一个技能在 YAML 中的结构体
// - name: 技能名称（唯一标识）
// - version: 版本号（可选，用于文档或兼容性控制）
// - description: 技能说明文本
// - required_skills: 该技能依赖的其它技能名列表（当前未自动生成流程，仅保留信息）
// - interfaces: 该技能暴露/使用的接口列表，格式为 \"BaseName.action\" 或 \"BaseName.service\"
struct SkillSpec
{
  std::string name;
  std::string version;
  std::string description;
  std::vector<std::string> required_skills;
  std::vector<std::string> interfaces;  // e.g. \"BaseName.action\" / \"BaseName.service\"
};

// SkillManager 负责：
// 1) 解析 robot_skills.yaml 描述文件
// 2) 自动注册其 Action/Service 接口客户端
// 3) 行为树节点注册在其它组件完成（本类不直接创建 BT action node）
class SkillManager
{
public:
  SkillManager(rclcpp::Node * node, ActionClientRegistry * action_clients, BTRegistry * bt)
  : node_(node), action_clients_(action_clients), bt_(bt) {}

  bool load_from_file(const std::string & yaml_path);

  const std::unordered_map<std::string, SkillSpec> & skills() const {return skills_;}

  // 枚举 action 技能 (skill_name, action_base)
  std::vector<std::pair<std::string, std::string>> enumerate_action_skills() const;

private:
  rclcpp::Node * node_ {nullptr};
  ActionClientRegistry * action_clients_ {nullptr};
  BTRegistry * bt_ {nullptr};
  std::unordered_map<std::string, SkillSpec> skills_;
  std::vector<rclcpp::ClientBase::SharedPtr> service_clients_; // 保存 service client 生命周期
  std::unordered_map<std::string, rclcpp::ClientBase::SharedPtr> service_client_map_; // skill name -> client

  void register_interfaces_(const SkillSpec & s);

public:
  // 获取特定类型 service client（若存在）
  template<typename T>
  std::shared_ptr<rclcpp::Client<T>> get_typed_service_client(const std::string & skill_name) const
  {
    auto it = service_client_map_.find(skill_name);
    if (it == service_client_map_.end()) {return nullptr;}
    return std::dynamic_pointer_cast<rclcpp::Client<T>>(it->second);
  }
};

} // namespace brain
"""

CPP_PREAMBLE_TEMPLATE = Template("""// skill_manager.cpp (generated)
// 本文件由 gen_skill_manager.py 自动生成，请勿手工修改。
// - 根据 robot_skills.yaml 中的技能定义，生成接口类型引用与注册逻辑。
//#include <ament_index_cpp/get_package_share_directory.hpp>
#include \"brain/skill_manager.hpp\"

#include <optional>
#include <cctype>
#include <rclcpp_action/rclcpp_action.hpp>

${includes}

${usings}

namespace brain
{

// 将字符串转为小写
static inline std::string to_lower(std::string s)
{
  for (char & c : s) {
    c = static_cast<char>(::tolower(static_cast<unsigned char>(c)));
  }
  return s;
}

// 将 TitleCase/camelCase 转换为 snake_case
static inline std::string to_snake_case(const std::string & s)
{
  std::string out;
  out.reserve(s.size() * 2);
  for (size_t i = 0; i < s.size(); ++i) {
    unsigned char uc = static_cast<unsigned char>(s[i]);
    if (::isupper(uc)) {
      if (i > 0) {out.push_back('_');}
      out.push_back(static_cast<char>(::tolower(uc)));
    } else {
      out.push_back(static_cast<char>(uc));
    }
  }
  return out;
}

struct ParsedInterface { std::string base; std::string kind; };

static std::optional<ParsedInterface> parse_interface_entry(const std::string & entry)
{
  const auto pos = entry.rfind('.');
  if (pos == std::string::npos || pos == 0 || pos + 1 >= entry.size()) {
    return std::nullopt;
  }
  ParsedInterface p;
  p.base = entry.substr(0, pos);
  p.kind = to_lower(entry.substr(pos + 1));
  return p;
}

bool SkillManager::load_from_file(const std::string & yaml_path)
{
  YAML::Node root = YAML::LoadFile(yaml_path);
  if (!root || !root.IsSequence()) {
    RCLCPP_ERROR(node_->get_logger(), "Invalid skills YAML: expected sequence at root");
    return false;
  }

  for (const auto & item : root) {
    if (!item.IsMap()) {continue;}
    SkillSpec s;
    s.name = item["name"].as<std::string>("");
    s.version = item["version"].as<std::string>("");
    s.description = item["description"].as<std::string>("");
    if (item["required_skills"]) {
      for (const auto & v : item["required_skills"]) {
        s.required_skills.push_back(v.as<std::string>());
      }
    }
    if (item["interfaces"]) {
      for (const auto & v : item["interfaces"]) {
        s.interfaces.push_back(v.as<std::string>());
      }
    }
    if (s.name.empty()) {continue;}
    skills_[s.name] = s;
    register_interfaces_(s);
  }

  RCLCPP_INFO(node_->get_logger(), "Loaded %zu skills from %s", skills_.size(), yaml_path.c_str());
  return true;
}

std::vector<std::pair<std::string, std::string>> SkillManager::enumerate_action_skills() const
{
  std::vector<std::pair<std::string, std::string>> out;
  for (const auto & kv : skills_) {
    const auto & s = kv.second;
    for (const auto & iface : s.interfaces) {
      auto p = parse_interface_entry(iface);
      if (p && p->kind == "action") { out.emplace_back(s.name, p->base); break; }
    }
  }
  return out;
}

void SkillManager::register_interfaces_(const SkillSpec & s)
{
  // 由代码生成器根据 YAML 填充的注册器表（接口基名 -> 注册逻辑）
  static const std::unordered_map<std::string, std::function<void(const std::string &, const std::string &)>> action_registrars = {
${action_registrars}
  };
  static const std::unordered_map<std::string, std::function<void(const std::string &, const std::string &)>> service_registrars = {
${service_registrars}
  };

  for (const auto & iface : s.interfaces) {
    auto parsed = parse_interface_entry(iface);
    if (!parsed) {
      RCLCPP_WARN(node_->get_logger(), "Skip invalid interface entry: %s", iface.c_str());
      continue;
    }
    // 统一策略：Action/Service 话题名直接使用技能名 (CamelCase)，不再额外注册 snake_case 变体
    const std::string topic_name = s.name;
    if (parsed->kind == "action") {
      auto it = action_registrars.find(parsed->base);
      if (it != action_registrars.end()) {
        it->second(topic_name, s.name);
      } else {
        RCLCPP_WARN(node_->get_logger(), "No action registrar for interface base '%s'", parsed->base.c_str());
      }
    } else if (parsed->kind == "service") {
      auto it = service_registrars.find(parsed->base);
      if (it != service_registrars.end()) {
        it->second(topic_name, s.name);
      } else {
        RCLCPP_WARN(node_->get_logger(), "No service registrar for interface base '%s'", parsed->base.c_str());
      }
    }
  }
}

} // namespace brain
""")


def title_to_snake(name: str) -> str:
    s = []
    for i, ch in enumerate(name):
        if ch.isupper():
            if i > 0:
                s.append('_')
            s.append(ch.lower())
        else:
            s.append(ch)
    return ''.join(s)


def load_goal_defaults(path: Path) -> Dict[str, str]:
  if not path or not path.exists():
    return {}
  try:
    data = yaml.safe_load(path.read_text(encoding='utf-8'))
    if isinstance(data, dict):
      # Expect mapping: SkillName: |\n  <C++ lines using g; must end with return g;>
      out: Dict[str, str] = {}
      for k, v in data.items():
        if isinstance(v, str):
          out[k] = v.strip() if "return" in v else (v.strip() + "\nreturn g;")
      return out
  except Exception as e:
    print(f"[warn] Failed to load goal defaults: {e}", file=sys.stderr)
  return {}


def generate(yaml_path: Path, include_dir: Path, src_dir: Path, goal_defaults_path: Path = None) -> None:
  # 1. Load goal defaults
  goal_defaults = load_goal_defaults(goal_defaults_path) if goal_defaults_path else {}

  # 2. Parse YAML
  data = yaml.safe_load(yaml_path.read_text(encoding='utf-8'))
  if not isinstance(data, list):
    raise ValueError('YAML root must be a sequence (list) of skills')

  action_bases = set()
  service_bases = set()
  # Keep skills for potential future use (currently not needed further here)
  for it in data:
    if not isinstance(it, dict):
      continue
    name = it.get('name') or ''
    if not name:
      continue
    interfaces = it.get('interfaces') or []
    for e in interfaces:
      if not isinstance(e, str):
        continue
      pos = e.rfind('.')
      if pos <= 0 or pos + 1 >= len(e):
        continue
      base = e[:pos]
      kind = e[pos+1:].lower()
      if kind == 'action':
        action_bases.add(base)
      elif kind == 'service':
        service_bases.add(base)

  # 3. Build include + using lines
  include_lines: list[str] = []
  using_lines: list[str] = []
  for base in sorted(action_bases):
    include_lines.append(f"#include \"interfaces/action/{title_to_snake(base)}.hpp\"")
    using_lines.append(f"using interfaces::action::{base};")
  for base in sorted(service_bases):
    include_lines.append(f"#include \"interfaces/srv/{title_to_snake(base)}.hpp\"")
    using_lines.append(f"using interfaces::srv::{base};")

  # 4. Helper for goal initialization (prefer skill name then base name)
  def goal_builder(skill: str, base: str) -> str:
    key = skill if skill in goal_defaults else (base if base in goal_defaults else None)
    if key:
      snippet = goal_defaults[key].replace('${BASE}', base)
      prefix = f'{base}::Goal g;'
      if '::Goal g' in snippet:
        return snippet
      return prefix + '\n' + snippet
    return f'{base}::Goal g; return g;'

  # 5. Build registrar blocks keyed only by interface base
  action_entries: list[str] = []
  for base in sorted(action_bases):
    block = (
      '    {"%s", [this](const std::string & topic, const std::string & internal_skill) {\n'
      '        if (action_clients_) {\n'
      '          action_clients_->register_client<%s>(\n'
      '            topic,\n'
      '            []() { %s },\n'
      '            nullptr,\n'
      '            [this, internal_skill](const ActionClientRegistry::GoalHandle<%s>::WrappedResult & res) {\n'
      '              if (res.code == rclcpp_action::ResultCode::SUCCEEDED && res.result && res.result->success)\n'
      '              {\n'
      '                RCLCPP_INFO(node_->get_logger(), "action succeeded: %%s", internal_skill.c_str());\n'
      '              } else {\n'
      '                RCLCPP_ERROR(node_->get_logger(), "action failed: %%s", internal_skill.c_str());\n'
      '              }\n'
      '            }\n'
      '          );\n'
      '        }\n'
      '      }},'
    ) % (base, base, goal_builder(base, base), base)
    action_entries.append(block)

  service_entries: list[str] = []
  for base in sorted(service_bases):
    block = (
      '    {"%s", [this](const std::string & topic, const std::string & internal_skill) {\n'
      '        (void)internal_skill;\n'
      '        auto client = node_->create_client<%s>(topic);\n'
      '        service_clients_.push_back(client);\n'
      '        service_client_map_[topic] = client;\n'
      '      }},'
    ) % (base, base)
    service_entries.append(block)

  # 6. Emit files
  hdr_path = include_dir / 'skill_manager.hpp'
  cpp_path = src_dir / 'skill_manager.cpp'

  hdr_code = HEADER_TEMPLATE
  cpp_code = CPP_PREAMBLE_TEMPLATE.substitute(
    includes='\n'.join(include_lines),
    usings='\n'.join(using_lines),
    action_registrars=("\n".join(action_entries) + ("\n" if action_entries else "")),
    service_registrars=("\n".join(service_entries) + ("\n" if service_entries else "")),
    camelcase_duplicates=''  # retained for template compatibility
  )

  hdr_path.parent.mkdir(parents=True, exist_ok=True)
  cpp_path.parent.mkdir(parents=True, exist_ok=True)
  hdr_path.write_text(hdr_code, encoding='utf-8')
  cpp_path.write_text(cpp_code, encoding='utf-8')
  print(f"Generated: {hdr_path}")
  print(f"Generated: {cpp_path}")


def main():
  ap = argparse.ArgumentParser()
  ap.add_argument('--yaml', dest='yaml_path', default='src/brain/config/robot_skills.yaml')
  ap.add_argument('--include-dir', dest='include_dir', default='src/brain/include/brain')
  ap.add_argument('--src-dir', dest='src_dir', default='src/brain/src')
  ap.add_argument('--goal-defaults', dest='goal_defaults', default=None,
          help='Optional YAML providing per-skill goal initialization snippets (no hard-coded names in generator).')
  args = ap.parse_args()

  yaml_path = Path(args.yaml_path)
  include_dir = Path(args.include_dir)
  src_dir = Path(args.src_dir)
  goal_defaults = Path(args.goal_defaults) if args.goal_defaults else None

  if not yaml_path.exists():
    print(f"YAML not found: {yaml_path}", file=sys.stderr)
    sys.exit(1)

  generate(yaml_path, include_dir, src_dir, goal_defaults)


if __name__ == '__main__':
    main()
