#! /usr/bin/env python3
# -*- coding: utf-8 -*-

# Copyright (c) 2023, AgiBot Inc.
# All rights reserved.

import sys

from google.protobuf.compiler import plugin_pb2 as plugin
from google.protobuf.compiler.plugin_pb2 import CodeGeneratorRequest as CodeGeneratorRequest
from google.protobuf.compiler.plugin_pb2 import CodeGeneratorResponse as CodeGeneratorResponse
from google.protobuf.descriptor_pb2 import FileDescriptorProto


class AimRTCodeGenerator:
    t_hfile: str = r"""/**
 * @file {{file_name}}.aimrt_rpc.pb.h
 * @brief This file was generated by protoc-gen-aimrt_rpc which is a self-defined pb compiler plugin, do not edit it!!!
 */
#pragma once

#include <future>

#include "aimrt_module_cpp_interface/rpc/rpc_handle.h"
#include "aimrt_module_cpp_interface/rpc/rpc_status.h"
#include "aimrt_module_cpp_interface/util/version.h"

#include "aimrt_module_cpp_interface/co/task.h"

#include "{{file_name}}.pb.h"

static_assert({{cur_gencode_version}} <= AIMRT_RUNTIME_VERSION_INT,
              "AIMRT_RUNTIME_VERSION is older than generated code version {{cur_gencode_version_str}}");
static_assert(AIMRT_MIN_GENCODE_VERSION_INT <= {{cur_gencode_version}},
              "AIMRT_MIN_GENCODE_VERSION is greater than generated code version {{cur_gencode_version_str}}");


{{namespace_begin}}
{{for service begin}}
class {{service_name}}SyncService : public aimrt::rpc::ServiceBase {
 public:
  {{service_name}}SyncService();
  ~{{service_name}}SyncService() override = default;
{{for method begin}}
  virtual aimrt::rpc::Status {{rpc_func_name}}(
      aimrt::rpc::ContextRef ctx_ref,
      const {{rpc_req_name}}& req,
      {{rpc_rsp_name}}& rsp) {
    return aimrt::rpc::Status(AIMRT_RPC_STATUS_SVR_NOT_IMPLEMENTED);
  }
{{method end}}
};
{{service end}}

{{for service begin}}
class {{service_name}}AsyncService : public aimrt::rpc::ServiceBase {
 public:
  {{service_name}}AsyncService();
  ~{{service_name}}AsyncService() override = default;
{{for method begin}}
  virtual void {{rpc_func_name}}(
      aimrt::rpc::ContextRef ctx_ref,
      const {{rpc_req_name}}& req,
      {{rpc_rsp_name}}& rsp,
      std::function<void(aimrt::rpc::Status)>&& callback) {
    callback(aimrt::rpc::Status(AIMRT_RPC_STATUS_SVR_NOT_IMPLEMENTED));
  }
{{method end}}
};
{{service end}}

{{for service begin}}
class {{service_name}}CoService : public aimrt::rpc::CoServiceBase {
 public:
  {{service_name}}CoService();
  ~{{service_name}}CoService() override = default;
{{for method begin}}
  virtual aimrt::co::Task<aimrt::rpc::Status> {{rpc_func_name}}(
      aimrt::rpc::ContextRef ctx_ref,
      const {{rpc_req_name}}& req,
      {{rpc_rsp_name}}& rsp) {
    co_return aimrt::rpc::Status(AIMRT_RPC_STATUS_SVR_NOT_IMPLEMENTED);
  }
{{method end}}
};

{{service end}}

{{for service begin}}
bool Register{{service_name}}ClientFunc(aimrt::rpc::RpcHandleRef rpc_handle_ref, std::string_view service_name);
bool Register{{service_name}}ClientFunc(aimrt::rpc::RpcHandleRef rpc_handle_ref);
{{service end}}

{{for service begin}}
class {{service_name}}SyncProxy : public aimrt::rpc::SyncProxyBase {
 public:
  explicit {{service_name}}SyncProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref);
  explicit {{service_name}}SyncProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref, std::string_view service_name);
  ~{{service_name}}SyncProxy() = default;

  static bool RegisterClientFunc(aimrt::rpc::RpcHandleRef rpc_handle_ref) {
    return Register{{service_name}}ClientFunc(rpc_handle_ref);
  }

  static bool RegisterClientFunc(aimrt::rpc::RpcHandleRef rpc_handle_ref, std::string_view service_name) {
    return Register{{service_name}}ClientFunc(rpc_handle_ref, service_name);
  }
{{for method begin}}
  aimrt::rpc::Status {{rpc_func_name}}(
      aimrt::rpc::ContextRef ctx_ref,
      const {{rpc_req_name}}& req,
      {{rpc_rsp_name}}& rsp);

  aimrt::rpc::Status {{rpc_func_name}}(
      const {{rpc_req_name}}& req,
      {{rpc_rsp_name}}& rsp) {
    return {{rpc_func_name}}(aimrt::rpc::ContextRef(), req, rsp);
  }
{{method end}}
};
{{service end}}

{{for service begin}}
class {{service_name}}AsyncProxy : public aimrt::rpc::AsyncProxyBase {
 public:
  explicit {{service_name}}AsyncProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref);
  explicit {{service_name}}AsyncProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref, std::string_view service_name);
  ~{{service_name}}AsyncProxy() = default;

  static bool RegisterClientFunc(aimrt::rpc::RpcHandleRef rpc_handle_ref) {
    return Register{{service_name}}ClientFunc(rpc_handle_ref);
  }

  static bool RegisterClientFunc(aimrt::rpc::RpcHandleRef rpc_handle_ref, std::string_view service_name) {
    return Register{{service_name}}ClientFunc(rpc_handle_ref, service_name);
  }
{{for method begin}}
  void {{rpc_func_name}}(
      aimrt::rpc::ContextRef ctx_ref,
      const {{rpc_req_name}}& req,
      {{rpc_rsp_name}}& rsp,
      std::function<void(aimrt::rpc::Status)>&& callback);

  void {{rpc_func_name}}(
      const {{rpc_req_name}}& req,
      {{rpc_rsp_name}}& rsp,
      std::function<void(aimrt::rpc::Status)>&& callback) {
    {{rpc_func_name}}(aimrt::rpc::ContextRef(), req, rsp, std::move(callback));
  }
{{method end}}
};
{{service end}}


{{for service begin}}
class {{service_name}}FutureProxy : public aimrt::rpc::FutureProxyBase {
 public:
  explicit {{service_name}}FutureProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref);
  explicit {{service_name}}FutureProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref, std::string_view service_name);
  ~{{service_name}}FutureProxy() = default;

  static bool RegisterClientFunc(aimrt::rpc::RpcHandleRef rpc_handle_ref) {
    return Register{{service_name}}ClientFunc(rpc_handle_ref);
  }

  static bool RegisterClientFunc(aimrt::rpc::RpcHandleRef rpc_handle_ref, std::string_view service_name) {
    return Register{{service_name}}ClientFunc(rpc_handle_ref, service_name);
  }
{{for method begin}}
  std::future<aimrt::rpc::Status> {{rpc_func_name}}(
      aimrt::rpc::ContextRef ctx_ref,
      const {{rpc_req_name}}& req,
      {{rpc_rsp_name}}& rsp);

  std::future<aimrt::rpc::Status> {{rpc_func_name}}(
      const {{rpc_req_name}}& req,
      {{rpc_rsp_name}}& rsp) {
    return {{rpc_func_name}}(aimrt::rpc::ContextRef(), req, rsp);
  }
{{method end}}
};
{{service end}}

{{for service begin}}
class {{service_name}}CoProxy : public aimrt::rpc::CoProxyBase {
 public:
  explicit {{service_name}}CoProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref);
  explicit {{service_name}}CoProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref, std::string_view service_name);
  ~{{service_name}}CoProxy() = default;

  static bool RegisterClientFunc(aimrt::rpc::RpcHandleRef rpc_handle_ref) {
    return Register{{service_name}}ClientFunc(rpc_handle_ref);
  }

  static bool RegisterClientFunc(aimrt::rpc::RpcHandleRef rpc_handle_ref, std::string_view service_name) {
    return Register{{service_name}}ClientFunc(rpc_handle_ref, service_name);
  }
{{for method begin}}
  aimrt::co::Task<aimrt::rpc::Status> {{rpc_func_name}}(
      aimrt::rpc::ContextRef ctx_ref,
      const {{rpc_req_name}}& req,
      {{rpc_rsp_name}}& rsp);

  aimrt::co::Task<aimrt::rpc::Status> {{rpc_func_name}}(
      const {{rpc_req_name}}& req,
      {{rpc_rsp_name}}& rsp) {
    return {{rpc_func_name}}(aimrt::rpc::ContextRef(), req, rsp);
  }
{{method end}}
};
{{service end}}

{{namespace_end}}
"""

    t_ccfile: str = r"""/**
 * @file {{file_name}}.aimrt_rpc.pb.cc
 * @brief This file was generated by protoc-gen-aimrt_rpc which is a self-defined pb compiler plugin, do not edit it!!!
 */

#include "{{file_name}}.aimrt_rpc.pb.h"

#include "aimrt_module_cpp_interface/co/inline_scheduler.h"
#include "aimrt_module_cpp_interface/co/on.h"
#include "aimrt_module_cpp_interface/co/start_detached.h"
#include "aimrt_module_cpp_interface/co/then.h"
#include "aimrt_module_protobuf_interface/util/protobuf_type_support.h"

{{namespace_begin}}
static constexpr std::string_view kRpcType = "pb";

{{for service begin}}
static constexpr std::string_view k{{service_name}}Name = "{{package_name}}.{{service_name}}";
{{service end}}

{{for service begin}}
{{service_name}}SyncService::{{service_name}}SyncService() : aimrt::rpc::ServiceBase(kRpcType, k{{service_name}}Name) {
{{for method begin}}
  {
    aimrt::rpc::ServiceFunc service_callback(
        [this](const aimrt_rpc_context_base_t* ctx, const void* req, void* rsp, aimrt_function_base_t* result_callback_ptr) {
          aimrt::rpc::ServiceCallback result_callback(result_callback_ptr);

          aimrt::rpc::ContextRef ctx_ref(ctx);

          auto status = {{rpc_func_name}}(
              ctx_ref,
              *static_cast<const {{rpc_req_name}}*>(req),
              *static_cast<{{rpc_rsp_name}}*>(rsp));

          result_callback(status.Code());
        });
    RegisterServiceFunc(
        "{{rpc_func_name}}",
        nullptr,
        aimrt::GetProtobufMessageTypeSupport<{{rpc_req_name}}>(),
        aimrt::GetProtobufMessageTypeSupport<{{rpc_rsp_name}}>(),
        std::move(service_callback));
  }
{{method end}}
}
{{service end}}

{{for service begin}}
{{service_name}}AsyncService::{{service_name}}AsyncService() : aimrt::rpc::ServiceBase(kRpcType, k{{service_name}}Name) {
{{for method begin}}
  {
    aimrt::rpc::ServiceFunc service_callback(
        [this](const aimrt_rpc_context_base_t* ctx, const void* req, void* rsp, aimrt_function_base_t* result_callback_ptr) {
          auto result_callback_func_ptr = std::make_shared<aimrt::rpc::ServiceCallback>(result_callback_ptr);

          aimrt::rpc::ContextRef ctx_ref(ctx);

          {{rpc_func_name}}(
              ctx_ref,
              *static_cast<const {{rpc_req_name}}*>(req),
              *static_cast<{{rpc_rsp_name}}*>(rsp),
              [result_callback_func_ptr{std::move(result_callback_func_ptr)}](aimrt::rpc::Status status) {
                (*result_callback_func_ptr)(status.Code());
              });
        });
    RegisterServiceFunc(
        "{{rpc_func_name}}",
        nullptr,
        aimrt::GetProtobufMessageTypeSupport<{{rpc_req_name}}>(),
        aimrt::GetProtobufMessageTypeSupport<{{rpc_rsp_name}}>(),
        std::move(service_callback));
  }
{{method end}}
}
{{service end}}

{{for service begin}}
{{service_name}}CoService::{{service_name}}CoService() : aimrt::rpc::CoServiceBase(kRpcType, k{{service_name}}Name) {
{{for method begin}}
  {
    aimrt::rpc::ServiceFunc service_callback(
        [this](const aimrt_rpc_context_base_t* ctx, const void* req, void* rsp, aimrt_function_base_t* result_callback_ptr) {
          auto handle_ptr = std::make_unique<const aimrt::rpc::CoRpcHandle>(
              [this](aimrt::rpc::ContextRef ctx_ref, const void* req_ptr, void* rsp_ptr)
                  -> aimrt::co::Task<aimrt::rpc::Status> {
                return {{rpc_func_name}}(
                    ctx_ref,
                    *static_cast<const {{rpc_req_name}}*>(req_ptr),
                    *static_cast<{{rpc_rsp_name}}*>(rsp_ptr));
              });

          aimrt::rpc::ServiceCallback result_callback(result_callback_ptr);

          aimrt::rpc::ContextRef ctx_ref(ctx);

          auto* ptr = handle_ptr.get();
          aimrt::co::StartDetached(
              aimrt::co::On(
                  aimrt::co::InlineScheduler(),
                  filter_mgr_.InvokeRpc(*ptr, ctx_ref, req, rsp)) |
              aimrt::co::Then(
                  [handle_ptr{std::move(handle_ptr)}, result_callback{std::move(result_callback)}](aimrt::rpc::Status status) {
                    result_callback(status.Code());
                  }));
        });
    RegisterServiceFunc(
        "{{rpc_func_name}}",
        nullptr,
        aimrt::GetProtobufMessageTypeSupport<{{rpc_req_name}}>(),
        aimrt::GetProtobufMessageTypeSupport<{{rpc_rsp_name}}>(),
        std::move(service_callback));
  }
{{method end}}
}
{{service end}}

{{for service begin}}
bool Register{{service_name}}ClientFunc(aimrt::rpc::RpcHandleRef rpc_handle_ref, std::string_view service_name) {
{{for method begin}}
  if (!(rpc_handle_ref.RegisterClientFunc(
          kRpcType,
          service_name,
          "{{rpc_func_name}}",
          nullptr,
          aimrt::GetProtobufMessageTypeSupport<{{rpc_req_name}}>(),
          aimrt::GetProtobufMessageTypeSupport<{{rpc_rsp_name}}>())))
    return false;
{{method end}}
  return true;
}
bool Register{{service_name}}ClientFunc(aimrt::rpc::RpcHandleRef rpc_handle_ref) {
  return Register{{service_name}}ClientFunc(rpc_handle_ref, k{{service_name}}Name);
}
{{service end}}

{{for service begin}}
{{service_name}}SyncProxy::{{service_name}}SyncProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref)
    : aimrt::rpc::SyncProxyBase(rpc_handle_ref, kRpcType, k{{service_name}}Name) {}

{{service_name}}SyncProxy::{{service_name}}SyncProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref, std::string_view service_name)
    : aimrt::rpc::SyncProxyBase(rpc_handle_ref, kRpcType, service_name) {}

{{for method begin}}
aimrt::rpc::Status {{service_name}}SyncProxy::{{rpc_func_name}}(
    aimrt::rpc::ContextRef ctx_ref,
    const {{rpc_req_name}}& req,
    {{rpc_rsp_name}}& rsp) {
  const std::string& full_func_name = aimrt::rpc::GetFullFuncName(rpc_type_, service_name_, "{{rpc_func_name}}");
  return Invoke(full_func_name, ctx_ref, req, rsp);
}
{{method end}}
{{service end}}

{{for service begin}}
{{service_name}}AsyncProxy::{{service_name}}AsyncProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref)
    : aimrt::rpc::AsyncProxyBase(rpc_handle_ref, kRpcType, k{{service_name}}Name) {}

{{service_name}}AsyncProxy::{{service_name}}AsyncProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref, std::string_view service_name)
    : aimrt::rpc::AsyncProxyBase(rpc_handle_ref, kRpcType, service_name) {}

{{for method begin}}
void {{service_name}}AsyncProxy::{{rpc_func_name}}(
    aimrt::rpc::ContextRef ctx_ref,
    const {{rpc_req_name}}& req,
    {{rpc_rsp_name}}& rsp,
    std::function<void(aimrt::rpc::Status)>&& callback) {
  const std::string& full_func_name = aimrt::rpc::GetFullFuncName(rpc_type_, service_name_, "{{rpc_func_name}}");
  Invoke(full_func_name, ctx_ref, req, rsp, std::move(callback));
}
{{method end}}
{{service end}}

{{for service begin}}
{{service_name}}FutureProxy::{{service_name}}FutureProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref)
    : aimrt::rpc::FutureProxyBase(rpc_handle_ref, kRpcType, k{{service_name}}Name) {}

{{service_name}}FutureProxy::{{service_name}}FutureProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref, std::string_view service_name)
    : aimrt::rpc::FutureProxyBase(rpc_handle_ref, kRpcType, service_name) {}

{{for method begin}}
std::future<aimrt::rpc::Status> {{service_name}}FutureProxy::{{rpc_func_name}}(
    aimrt::rpc::ContextRef ctx_ref,
    const {{rpc_req_name}}& req,
    {{rpc_rsp_name}}& rsp) {
  const std::string& full_func_name = aimrt::rpc::GetFullFuncName(rpc_type_, service_name_, "{{rpc_func_name}}");
  return Invoke(full_func_name, ctx_ref, req, rsp);
}
{{method end}}
{{service end}}

{{for service begin}}
{{service_name}}CoProxy::{{service_name}}CoProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref)
    : aimrt::rpc::CoProxyBase(rpc_handle_ref, kRpcType, k{{service_name}}Name) {}

{{service_name}}CoProxy::{{service_name}}CoProxy(aimrt::rpc::RpcHandleRef rpc_handle_ref, std::string_view service_name)
    : aimrt::rpc::CoProxyBase(rpc_handle_ref, kRpcType, service_name) {}

{{for method begin}}
aimrt::co::Task<aimrt::rpc::Status> {{service_name}}CoProxy::{{rpc_func_name}}(
    aimrt::rpc::ContextRef ctx_ref,
    const {{rpc_req_name}}& req,
    {{rpc_rsp_name}}& rsp) {
  const std::string& full_func_name = aimrt::rpc::GetFullFuncName(rpc_type_, service_name_, "{{rpc_func_name}}");
  co_return co_await Invoke(full_func_name, ctx_ref, req, rsp);
}
{{method end}}
{{service end}}

{{namespace_end}}
"""

    class MethodNode:
        def __init__(self):
            self.kv = {}

    class ServiceNode:
        def __init__(self):
            self.kv = {}
            self.method_vec = []

    class PackageNode:
        def __init__(self):
            self.kv = {}
            self.service_vec = []

    @staticmethod
    def gen_method_code(temp, method_node) -> str:
        result = temp
        for key, value in method_node.kv.items():
            result = result.replace(key, value)
        return result

    @staticmethod
    def gen_service_code(temp, service_node) -> str:
        result = temp

        for key, value in service_node.kv.items():
            result = result.replace(key, value)

        method_begin_flag = "{{for method begin}}\n"
        method_end_flag = "{{method end}}\n"

        cur_pos = 0
        while True:
            begin_pos = result.find(method_begin_flag, cur_pos)
            if begin_pos == -1:
                break
            end_pos = result.find(method_end_flag, begin_pos + len(method_begin_flag))
            if end_pos == -1:
                break

            cur_temp = result[begin_pos + len(method_begin_flag): end_pos]
            cur_result = ""
            for node in service_node.method_vec:
                cur_result += AimRTCodeGenerator.gen_method_code(cur_temp, node)

            result = result[:begin_pos] + cur_result + result[end_pos + len(method_end_flag):]
            cur_pos = begin_pos + len(cur_result)

        return result

    @staticmethod
    def gen_package_code(temp, package_node) -> str:
        result = temp

        for key, value in package_node.kv.items():
            result = result.replace(key, value)

        service_begin_flag = "{{for service begin}}\n"
        service_end_flag = "{{service end}}\n"

        cur_pos = 0
        while True:
            begin_pos = result.find(service_begin_flag, cur_pos)
            if begin_pos == -1:
                break
            end_pos = result.find(service_end_flag, begin_pos + len(service_begin_flag))
            if end_pos == -1:
                break

            cur_temp = result[begin_pos + len(service_begin_flag): end_pos]
            cur_result = ""
            for node in package_node.service_vec:
                cur_result += AimRTCodeGenerator.gen_service_code(cur_temp, node)

            result = result[:begin_pos] + cur_result + result[end_pos + len(service_end_flag):]
            cur_pos = begin_pos + len(cur_result)

        return result

    @staticmethod
    def gen_name_space_str(ns: str) -> str:
        return ns.replace(".", "::")

    @staticmethod
    def gen_namespace_begin_str(ns: str) -> str:
        namespace_vec: list[str] = ns.split(".")
        result: str = ""
        for itr in namespace_vec:
            result += "namespace " + itr + " {\n"
        return result

    @staticmethod
    def gen_namespace_end_str(ns: str) -> str:
        namespace_vec: list[str] = ns.split(".")
        namespace_vec.reverse()
        result: str = ""
        for itr in namespace_vec:
            result += "}  // namespace " + itr + "\n"
        return result

    def generate(self, request: CodeGeneratorRequest) -> CodeGeneratorResponse:
        """Generate code for the given request"""
        response: CodeGeneratorResponse = CodeGeneratorResponse()
        for proto_file in request.proto_file:
            if proto_file.name not in request.file_to_generate:
                continue

            # Generate code for each file
            file_name: str = proto_file.name
            package_name: str = proto_file.package

            package_node = AimRTCodeGenerator.PackageNode()
            package_node.kv["{{cur_gencode_version}}"] = "1000000"
            package_node.kv["{{cur_gencode_version_str}}"] = "1.0.0"
            package_node.kv["{{file_name}}"] = file_name.replace(".proto", "")
            package_node.kv["{{package_name}}"] = package_name
            package_node.kv["{{namespace_begin}}"] = self.gen_namespace_begin_str(package_name)
            package_node.kv["{{namespace_end}}"] = self.gen_namespace_end_str(package_name)

            for ii in range(0, len(proto_file.service)):
                service = proto_file.service[ii]

                service_node = AimRTCodeGenerator.ServiceNode()
                service_node.kv["{{service_name}}"] = service.name

                for jj in range(0, len(service.method)):
                    method = service.method[jj]

                    method_node = AimRTCodeGenerator.MethodNode()

                    method_node.kv["{{rpc_func_name}}"] = method.name
                    method_node.kv["{{rpc_req_name}}"] = self.gen_name_space_str(method.input_type)
                    method_node.kv["{{rpc_rsp_name}}"] = self.gen_name_space_str(method.output_type)

                    service_node.method_vec.append(method_node)

                package_node.service_vec.append(service_node)

            hfile: CodeGeneratorResponse.File = CodeGeneratorResponse.File()
            hfile.name = file_name.replace(".proto", ".aimrt_rpc.pb.h")
            hfile.content = AimRTCodeGenerator.gen_package_code(self.t_hfile, package_node)
            response.file.append(hfile)

            ccfile: CodeGeneratorResponse.File = CodeGeneratorResponse.File()
            ccfile.name = file_name.replace(".proto", ".aimrt_rpc.pb.cc")
            ccfile.content = AimRTCodeGenerator.gen_package_code(self.t_ccfile, package_node)
            response.file.append(ccfile)

        return response


if __name__ == "__main__":
    request: CodeGeneratorRequest = CodeGeneratorRequest.FromString(sys.stdin.buffer.read())

    aimrt_code_generator: AimRTCodeGenerator = AimRTCodeGenerator()

    response: CodeGeneratorResponse = aimrt_code_generator.generate(request)

    sys.stdout.buffer.write(response.SerializeToString())
