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

# Copyright (c) 2024 The AimRT Authors.
# AimRT is licensed under Mulan PSL v2.

import os
import sys

t_pyfile: str = r"""# This file was generated by ros2-aimrt_py_rpc-gen-code-tool, do not edit it!!!

from typing import overload

import aimrt_py
import aimrt_py.aimrt_python_runtime_ros2 as aimrt_py_ros2
import {{package_name}}.srv

assert (
    {{cur_gencode_version}} <= aimrt_py.AimRT_RUNTIME_VERSION_INT
), "AimRT_RUNTIME_VERSION is older than generated code version {{cur_gencode_version_str}}"
assert (
    aimrt_py.AIMRT_MIN_GENCODE_VERSION_INT <= {{cur_gencode_version}}
), "AIMRT_MIN_GENCODE_VERSION is greater than generated code version {{cur_gencode_version_str}}"


if {{package_name}}.srv.{{service_name}}._TYPE_SUPPORT is None:
    {{package_name}}.srv.{{service_name}}.__import_type_support__()


class {{service_name}}Service(aimrt_py.ServiceBase):
    def __init__(self):
        super().__init__("ros2", "{{package_name}}/srv")

        # {{service_name}}
        {{service_name}}_req_aimrt_ts = aimrt_py_ros2.PyRos2TypeSupport({{package_name}}.srv.{{service_name}}.Request)
        {{service_name}}_req_aimrt_ts.SetTypeName(aimrt_py.GetRos2MessageTypeName({{package_name}}.srv.{{service_name}}.Request))
        {{service_name}}_req_aimrt_ts.SetSerializationTypesSupportedList(["ros2"])

        {{service_name}}_rsp_aimrt_ts = aimrt_py_ros2.PyRos2TypeSupport({{package_name}}.srv.{{service_name}}.Response)
        {{service_name}}_rsp_aimrt_ts.SetTypeName(aimrt_py.GetRos2MessageTypeName({{package_name}}.srv.{{service_name}}.Response))
        {{service_name}}_rsp_aimrt_ts.SetSerializationTypesSupportedList(["ros2"])

        if not aimrt_py.check_is_valid_srv_type({{package_name}}.srv.{{service_name}}):
            raise RuntimeError("The service type provided is not valid")

        aimrt_py_ros2.Ros2RegisterServiceFunc(
            self,
            "{{service_name}}",
            {{package_name}}.srv.{{service_name}},
            {{service_name}}_req_aimrt_ts,
            {{package_name}}.srv.{{service_name}}.Request,
            {{service_name}}_rsp_aimrt_ts,
            {{package_name}}.srv.{{service_name}}.Response,
            self.{{service_name}})

    def {{service_name}}(self, ctx_ref, req) -> tuple[aimrt_py.RpcStatus, {{package_name}}.srv.{{service_name}}.Response]:
        return (aimrt_py.RpcStatus(aimrt_py.RpcStatusRetCode.SVR_NOT_IMPLEMENTED), {{package_name}}.srv.{{service_name}}.Response())

class {{service_name}}ServiceProxy(aimrt_py.ProxyBase):
    def __init__(self, rpc_handle_ref=aimrt_py.RpcHandleRef()):
        super().__init__(rpc_handle_ref, "ros2", "{{package_name}}")
        self.rpc_handle_ref = rpc_handle_ref

    @overload
    def {{service_name}}(
        self, req: {{package_name}}.srv.{{service_name}}.Request
    ) -> tuple[aimrt_py.RpcStatus, {{package_name}}.srv.{{service_name}}.Response]: ...

    @overload
    def {{service_name}}(
        self, ctx: aimrt_py.RpcContext, req: {{package_name}}.srv.{{service_name}}.Request
    ) -> tuple[aimrt_py.RpcStatus, {{package_name}}.srv.{{service_name}}.Response]: ...

    @overload
    def {{service_name}}(
        self, ctx_ref: aimrt_py.RpcContextRef, req: {{package_name}}.srv.{{service_name}}.Request
    ) -> tuple[aimrt_py.RpcStatus, {{package_name}}.srv.{{service_name}}.Response]: ...

    def {{service_name}}(self, *args) -> tuple[aimrt_py.RpcStatus, {{package_name}}.srv.{{service_name}}.Response]:
        if len(args) == 1:
            ctx = super().NewContextSharedPtr()
            req = args[0]
        elif len(args) == 2:
            ctx = args[0]
            req = args[1]
        else:
            raise TypeError(f"{{service_name}} expects 1 or 2 arguments, got {len(args)}")

        if isinstance(ctx, aimrt_py.RpcContext):
            ctx_ref = aimrt_py.RpcContextRef(ctx)
        elif isinstance(ctx, aimrt_py.RpcContextRef):
            ctx_ref = ctx
        else:
            raise TypeError(f"ctx must be 'aimrt_py.RpcContext' or 'aimrt_py.RpcContextRef', got {type(ctx)}")

        if ctx_ref:
            if ctx_ref.GetSerializationType() == "":
                ctx_ref.SetSerializationType("ros2")
        else:
            real_ctx = aimrt_py.RpcContext()
            ctx_ref = aimrt_py.RpcContextRef(real_ctx)
            ctx_ref.SetSerializationType("ros2")

        status, rsp = aimrt_py_ros2.Ros2Invoke(
            self.rpc_handle_ref,
            "ros2:/{{package_name}}/srv/{{service_name}}",
            ctx_ref,
            req,
            {{package_name}}.srv.{{service_name}}.Response)

        return status, rsp


    @staticmethod
    def RegisterClientFunc(rpc_handle):
        # {{service_name}}
        {{service_name}}_req_aimrt_ts = aimrt_py_ros2.PyRos2TypeSupport({{package_name}}.srv.{{service_name}}.Request)
        {{service_name}}_req_aimrt_ts.SetTypeName(aimrt_py.GetRos2MessageTypeName({{package_name}}.srv.{{service_name}}.Request))
        {{service_name}}_req_aimrt_ts.SetSerializationTypesSupportedList(["ros2"])

        {{service_name}}_rsp_aimrt_ts = aimrt_py_ros2.PyRos2TypeSupport({{package_name}}.srv.{{service_name}}.Response)
        {{service_name}}_rsp_aimrt_ts.SetTypeName(aimrt_py.GetRos2MessageTypeName({{package_name}}.srv.{{service_name}}.Response))
        {{service_name}}_rsp_aimrt_ts.SetSerializationTypesSupportedList(["ros2"])

        return aimrt_py_ros2.Ros2RegisterClientFunc(
            rpc_handle,
            "ros2:/{{package_name}}/srv/{{service_name}}",
            {{package_name}}.srv.{{service_name}},
            {{service_name}}_req_aimrt_ts,
            {{service_name}}_rsp_aimrt_ts)

"""


def generate():
    package_name = ""
    service_file = ""
    output_path = ""
    for arg in sys.argv:
        kv = arg.split('=')
        if (kv[0] == '--pkg_name'):
            package_name = kv[1]
        elif (kv[0] == '--srv_file'):
            service_file = kv[1]
        elif (kv[0] == '--output_path'):
            output_path = kv[1]

    if package_name == "" or service_file == "" or output_path == "":
        print("Usage: python3 ros2_py_gen_aimrt_py_rpc.py "
              "--pkg_name=<package_name> "
              "--srv_file=<service_file> "
              "--output_path=<output_path>")
        exit(1)

    path, file = os.path.split(service_file)
    filename, ext = os.path.splitext(file)
    py_file_path = os.path.join(output_path, filename + "_aimrt_rpc_ros2.py")

    # py file
    f_py_file = open(py_file_path, 'w')
    f_py_file.write(t_pyfile
                    .replace("{{package_name}}", package_name)
                    .replace("{{service_name}}", filename)
                    .replace("{{cur_gencode_version}}", "1000000")
                    .replace("{{cur_gencode_version_str}}", "1.0.0"))
    f_py_file.close()


if __name__ == '__main__':
    generate()
