#include "onnx.h"
#include "onnx_core/expressions/expressions.h"
#include "onnx_lib/checker.h"
#include "onnx_lib/common/ir_pb_converter.h"
#include "onnx_lib/defs/schema.h"
#include "onnx_lib/defs/shape_inference.h"
#include "onnx_lib/inliner/inliner.h"
#include "onnx_lib/shape_inference/implementation.h"
#include "onnx_lib/version_converter/convert.h"
#include "onnx_lib/version_converter/errors.h"
#include "onnx_manipulations/compose.h"
#include "onnx_manipulations/parser.h"
#include "onnx_manipulations/printer.h"
#include <algorithm>
#include <limits>
#include <nanobind/stl/map.h>
#include <nanobind/stl/pair.h>
#include <nanobind/stl/set.h>
#include <nanobind/stl/string.h>
#include <nanobind/stl/tuple.h>
#include <nanobind/stl/unordered_map.h>
#include <nanobind/stl/unordered_set.h>
#include <nanobind/stl/vector.h>
#include <unordered_map>
#include <unordered_set>

namespace nb = nanobind;
using namespace ONNX_LIGHT_NAMESPACE;

void AddOnnxPyLib(nb::module_ &m);

NB_MODULE(_onnxpyprotolib, m) {
  m.doc() = "onnx lib bindings (defs/parser/checker/inliner/shape/version_converter) from "
            "python without protobuf but using the same format";
  AddOnnxPyLib(m);
}

void AddOnnxPyLib(nb::module_ &m) {
  // -----------------------------------------------------------------------
  // Submodule `version_converter`
  // -----------------------------------------------------------------------
  auto version_converter = m.def_submodule("version_converter");
  version_converter.doc() = "Version converter submodule";

  nb::exception<ConvertError>(
      version_converter,
      "ConvertError"); // NOLINT(bugprone-unused-raii,bugprone-throw-keyword-missing)

  version_converter.def(
      "convert_version",
      [](const ModelProto &mp_in, int target_version) {
        return version_conversion::ConvertVersion(mp_in, target_version);
      },
      nb::arg("model"), nb::arg("target_version"),
      "Convert a model to the specified target opset version.");

  // -----------------------------------------------------------------------
  // Submodule `parser`
  // -----------------------------------------------------------------------
  auto parser_mod = m.def_submodule("parser");
  parser_mod.doc() = "Parser submodule";

  // Returns (ok: bool, error_message: str, proto_object)
  parser_mod.def("parse_model", [](const std::string &text) -> nb::tuple {
    ModelProto proto{};
    OnnxParser parser(text.c_str());
    auto status = parser.Parse(proto);
    const std::string &error_msg = status.ErrorMessage();
    return nb::make_tuple(status.IsOK(), error_msg, std::move(proto));
  });

  parser_mod.def("parse_graph", [](const std::string &text) -> nb::tuple {
    GraphProto proto{};
    OnnxParser parser(text.c_str());
    auto status = parser.Parse(proto);
    const std::string &error_msg = status.ErrorMessage();
    return nb::make_tuple(status.IsOK(), error_msg, std::move(proto));
  });

  parser_mod.def("parse_function", [](const std::string &text) -> nb::tuple {
    FunctionProto proto{};
    OnnxParser parser(text.c_str());
    auto status = parser.Parse(proto);
    const std::string &error_msg = status.ErrorMessage();
    return nb::make_tuple(status.IsOK(), error_msg, std::move(proto));
  });

  parser_mod.def("parse_node", [](const std::string &text) -> nb::tuple {
    NodeProto proto{};
    OnnxParser parser(text.c_str());
    auto status = parser.Parse(proto);
    const std::string &error_msg = status.ErrorMessage();
    return nb::make_tuple(status.IsOK(), error_msg, std::move(proto));
  });

  // -----------------------------------------------------------------------
  // Submodule `printer`
  // -----------------------------------------------------------------------
  auto printer_mod = m.def_submodule("printer");
  printer_mod.doc() = "Printer submodule – converts proto objects to ONNX text format";

  printer_mod.def(
      "model_to_text", [](const ModelProto &model) -> std::string { return ProtoToString(model); },
      nb::arg("model"), "Converts a ModelProto to ONNX textual representation.");

  printer_mod.def(
      "graph_to_text", [](const GraphProto &graph) -> std::string { return ProtoToString(graph); },
      nb::arg("graph"), "Converts a GraphProto to ONNX textual representation.");

  printer_mod.def(
      "function_to_text", [](const FunctionProto &fn) -> std::string { return ProtoToString(fn); },
      nb::arg("fn"), "Converts a FunctionProto to ONNX textual representation.");

  // -----------------------------------------------------------------------
  // Submodule `shape_inference`
  // -----------------------------------------------------------------------
  auto shape_inference_mod = m.def_submodule("shape_inference");
  shape_inference_mod.doc() = "Shape inference submodule";

  nb::exception<InferenceError>(
      shape_inference_mod,
      "InferenceError"); // NOLINT(bugprone-unused-raii,bugprone-throw-keyword-missing)

  shape_inference_mod.def(
      "infer_function_output_types",
      [](const FunctionProto &function, const std::vector<TypeProto> &input_types,
         const std::vector<AttributeProto> &attributes) -> nb::list {
        auto output_types =
            shape_inference::InferFunctionOutputTypes(function, input_types, attributes);
        nb::list result;
        for (auto &type_proto : output_types) {
          result.append(nb::cast(type_proto));
        }
        return result;
      },
      nb::arg("function"), nb::arg("input_types"), nb::arg("attributes"),
      "Infers output types of a FunctionProto given serialized input TypeProtos and "
      "AttributeProtos.");

  shape_inference_mod.def(
      "infer_shapes",
      [](ModelProto &model, bool check_type, bool strict_mode, bool data_prop) -> void {
        ShapeInferenceOptions options(/*check_type_val=*/check_type,
                                      /*strict_mode_val=*/strict_mode ? 1 : 0,
                                      /*data_prop_val=*/data_prop);
        shape_inference::InferShapes(model, OpSchemaRegistry::Instance(), options);
      },
      nb::arg("model"), nb::arg("check_type") = false, nb::arg("strict_mode") = false,
      nb::arg("data_prop") = false,
      "Runs whole-model shape inference in place on the given ModelProto, populating "
      "inferred shapes/types on every intermediate value and graph output. ``check_type`` "
      "checks the type-equality for input and output, ``strict_mode`` raises any node-level "
      "shape inference errors, and ``data_prop`` enables data propagation for limited "
      "operators to perform shape computation.");

  // -----------------------------------------------------------------------
  // Submodule `defs`
  // -----------------------------------------------------------------------
  auto defs = m.def_submodule("defs");
  defs.doc() = "Schema submodule";

  nb::exception<SchemaError>(
      defs, "SchemaError"); // NOLINT(bugprone-unused-raii,bugprone-throw-keyword-missing)

  nb::class_<OpSchema> op_schema(defs, "OpSchema", "Schema of an operator.");

  nb::enum_<OpSchema::FormalParameterOption>(op_schema, "FormalParameterOption",
                                             nb::is_arithmetic())
      .value("Single", OpSchema::Single)
      .value("Optional", OpSchema::Optional)
      .value("Variadic", OpSchema::Variadic);

  nb::enum_<OpSchema::DifferentiationCategory>(op_schema, "DifferentiationCategory",
                                               nb::is_arithmetic())
      .value("Unknown", OpSchema::Unknown)
      .value("Differentiable", OpSchema::Differentiable)
      .value("NonDifferentiable", OpSchema::NonDifferentiable);

  nb::enum_<OpSchema::NodeDeterminism>(op_schema, "NodeDeterminism")
      .value("Deterministic", OpSchema::NodeDeterminism::Deterministic)
      .value("NonDeterministic", OpSchema::NodeDeterminism::NonDeterministic)
      .value("Unknown", OpSchema::NodeDeterminism::Unknown);

  nb::enum_<OpSchema::SupportType>(op_schema, "SupportType", nb::is_arithmetic())
      .value("COMMON", OpSchema::SupportType::COMMON)
      .value("EXPERIMENTAL", OpSchema::SupportType::EXPERIMENTAL);

  nb::class_<OpSchema::Attribute>(op_schema, "Attribute")
      .def(
          "__init__",
          [](OpSchema::Attribute *self, std::string name, AttributeProto::AttributeType type,
             std::string description, bool required) {
            new (self) OpSchema::Attribute(std::move(name), std::move(description), type, required);
          },
          nb::arg("name"), nb::arg("type"), nb::arg("description") = "", nb::kw_only(),
          nb::arg("required") = true)
      .def(
          "__init__",
          [](OpSchema::Attribute *self, std::string name, const AttributeProto &default_value,
             std::string description) {
            new (self) OpSchema::Attribute(std::move(name), std::move(description),
                                           AttributeProto(default_value));
          },
          nb::arg("name"), nb::arg("default_value"), nb::arg("description") = "")
      .def_ro("name", &OpSchema::Attribute::name)
      .def_ro("description", &OpSchema::Attribute::description)
      .def_ro("type", &OpSchema::Attribute::type)
      .def_prop_ro(
          "default_value",
          [](const OpSchema::Attribute *attr) -> AttributeProto { return attr->default_value; })
      .def_ro("required", &OpSchema::Attribute::required);

  nb::class_<OpSchema::TypeConstraintParam>(op_schema, "TypeConstraintParam")
      .def(nb::init<std::string, std::vector<std::string>, std::string>(),
           nb::arg("type_param_str"), nb::arg("allowed_type_strs"), nb::arg("description") = "")
      .def_ro("type_param_str", &OpSchema::TypeConstraintParam::type_param_str)
      .def_ro("allowed_type_strs", &OpSchema::TypeConstraintParam::allowed_type_strs)
      .def_ro("description", &OpSchema::TypeConstraintParam::description);

  nb::class_<OpSchema::FormalParameter>(op_schema, "FormalParameter")
      .def(
          "__init__",
          [](OpSchema::FormalParameter *self, std::string name, std::string type_str,
             const std::string &description, OpSchema::FormalParameterOption param_option,
             bool is_homogeneous, int min_arity,
             OpSchema::DifferentiationCategory differentiation_category) {
            new (self) OpSchema::FormalParameter(std::move(name), description, std::move(type_str),
                                                 param_option, is_homogeneous, min_arity,
                                                 differentiation_category);
          },
          nb::arg("name"), nb::arg("type_str"), nb::arg("description") = "", nb::kw_only(),
          nb::arg("param_option") = OpSchema::Single, nb::arg("is_homogeneous") = true,
          nb::arg("min_arity") = 1,
          nb::arg("differentiation_category") = OpSchema::DifferentiationCategory::Unknown)
      .def_prop_ro("name", &OpSchema::FormalParameter::GetName)
      .def_prop_ro("types", &OpSchema::FormalParameter::GetTypes)
      .def_prop_ro("type_str", &OpSchema::FormalParameter::GetTypeStr)
      .def_prop_ro("description", &OpSchema::FormalParameter::GetDescription)
      .def_prop_ro("option", &OpSchema::FormalParameter::GetOption)
      .def_prop_ro("is_homogeneous", &OpSchema::FormalParameter::GetIsHomogeneous)
      .def_prop_ro("min_arity", &OpSchema::FormalParameter::GetMinArity)
      .def_prop_ro("differentiation_category",
                   &OpSchema::FormalParameter::GetDifferentiationCategory);

  op_schema
      .def(
          "__init__",
          [](OpSchema *self, std::string name, std::string domain, int since_version,
             const std::string &doc, std::vector<OpSchema::FormalParameter> inputs,
             std::vector<OpSchema::FormalParameter> outputs,
             std::vector<std::tuple<std::string, std::vector<std::string>, std::string>>
                 type_constraints,
             std::vector<OpSchema::Attribute> attributes,
             OpSchema::NodeDeterminism node_determinism) {
            new (self) OpSchema();
            self->SetName(std::move(name))
                .SetDomain(std::move(domain))
                .SinceVersion(since_version)
                .SetDoc(doc);
            self->SetNodeDeterminism(node_determinism);
            for (size_t i = 0; i < inputs.size(); ++i) {
              self->Input(static_cast<int>(i), std::move(inputs[i]));
            }
            for (size_t i = 0; i < outputs.size(); ++i) {
              self->Output(static_cast<int>(i), std::move(outputs[i]));
            }
            for (auto &tc : type_constraints) {
              std::string type_str;
              std::vector<std::string> constraints;
              std::string description;
              tie(type_str, constraints, description) = std::move(tc);
              self->TypeConstraint(std::move(type_str), std::move(constraints),
                                   std::move(description));
            }
            for (auto &attribute : attributes) {
              self->Attr(std::move(attribute));
            }
            self->Finalize();
          },
          nb::arg("name"), nb::arg("domain"), nb::arg("since_version"), nb::arg("doc") = "",
          nb::kw_only(), nb::arg("inputs") = std::vector<OpSchema::FormalParameter>{},
          nb::arg("outputs") = std::vector<OpSchema::FormalParameter>{},
          nb::arg("type_constraints") =
              std::vector<std::tuple<std::string, std::vector<std::string>, std::string>>{},
          nb::arg("attributes") = std::vector<OpSchema::Attribute>{},
          nb::arg("node_determinism") = OpSchema::NodeDeterminism::Unknown)
      .def_prop_rw("name", &OpSchema::Name,
                   [](OpSchema &self, const std::string &name) { self.SetName(name); })
      .def_prop_rw("domain", &OpSchema::domain,
                   [](OpSchema &self, const std::string &domain) { self.SetDomain(domain); })
      .def_prop_rw(
          "doc",
          [](const OpSchema &self) -> std::string {
            const char *d = self.doc();
            return d ? d : "";
          },
          [](OpSchema &self, const std::string &doc) { self.SetDoc(doc); })
      .def_prop_ro("file", &OpSchema::file)
      .def_prop_ro("line", &OpSchema::line)
      .def_prop_ro("support_level", &OpSchema::support_level)
      .def_prop_ro("since_version", &OpSchema::since_version)
      .def_prop_ro("deprecated", &OpSchema::deprecated)
      .def_prop_ro("function_opset_versions", &OpSchema::function_opset_versions)
      .def_prop_ro("context_dependent_function_opset_versions",
                   &OpSchema::context_dependent_function_opset_versions)
      .def_prop_ro("all_function_opset_versions",
                   [](const OpSchema *op) -> std::vector<int> {
                     auto all = op->function_opset_versions();
                     auto ctx = op->context_dependent_function_opset_versions();
                     all.insert(all.end(), ctx.begin(), ctx.end());
                     std::sort(all.begin(), all.end());
                     all.erase(std::unique(all.begin(), all.end()), all.end());
                     return all;
                   })
      .def_prop_ro("min_input", &OpSchema::min_input)
      .def_prop_ro("max_input", &OpSchema::max_input)
      .def_prop_ro("min_output", &OpSchema::min_output)
      .def_prop_ro("max_output", &OpSchema::max_output)
      .def_prop_ro("attributes", &OpSchema::attributes)
      .def_prop_ro("inputs", &OpSchema::inputs)
      .def_prop_ro("outputs", &OpSchema::outputs)
      .def_prop_ro("has_type_and_shape_inference_function",
                   &OpSchema::has_type_and_shape_inference_function)
      .def_prop_ro("has_data_propagation_function", &OpSchema::has_data_propagation_function)
      .def_prop_ro("type_constraints", &OpSchema::typeConstraintParams)
      .def_static("is_infinite", [](int v) { return v == std::numeric_limits<int>::max(); })
      .def_prop_ro("non_deterministic",
                   [](const OpSchema *op) {
                     return op->GetNodeDeterminism() == OpSchema::NodeDeterminism::NonDeterministic;
                   })
      .def_prop_ro("has_function", &OpSchema::HasFunction)
      .def_prop_ro("_function_body",
                   [](const OpSchema *op) -> nb::object {
                     const FunctionProto *fp = op->GetFunction();
                     if (!fp)
                       return nb::none();
                     FunctionProto copy;
                     copy.CopyFrom(*fp);
                     return nb::cast(std::move(copy));
                   })
      .def("get_function_with_opset_version",
           [](const OpSchema *op, int opset_version) -> nb::object {
             const FunctionProto *fp = op->GetFunction(opset_version);
             if (!fp)
               return nb::none();
             FunctionProto copy;
             copy.CopyFrom(*fp);
             return nb::cast(std::move(copy));
           })
      .def_prop_ro("has_context_dependent_function", &OpSchema::HasContextDependentFunction)
      .def("get_context_dependent_function",
           [](const OpSchema *op, const NodeProto &node,
              const std::vector<TypeProto> &input_types) -> nb::object {
             if (!op->HasContextDependentFunction())
               return nb::none();
             FunctionBodyBuildContextImpl ctx(node, input_types);
             FunctionProto func_proto;
             op->BuildContextDependentFunction(ctx, func_proto);
             return nb::cast(std::move(func_proto));
           })
      .def("get_context_dependent_function_with_opset_version",
           [](const OpSchema *op, int opset_version, const NodeProto &node,
              const std::vector<TypeProto> &input_types) -> nb::object {
             if (!op->HasContextDependentFunctionWithOpsetVersion(opset_version))
               return nb::none();
             FunctionBodyBuildContextImpl ctx(node, input_types);
             FunctionProto func_proto;
             op->BuildContextDependentFunction(ctx, func_proto, opset_version);
             return nb::cast(std::move(func_proto));
           })
      .def(
          "_infer_node_outputs",
          [](const OpSchema *schema, NodeProto &node,
             std::unordered_map<std::string, TypeProto *> input_types,
             std::unordered_map<std::string, TensorProto *> input_data,
             std::unordered_map<std::string, SparseTensorProto *> input_sparse_data,
             std::unordered_map<std::string, int> opset_imports,
             int ir_version) -> std::unordered_map<std::string, TypeProto> {
            // ``node`` and the value maps are taken by (non-const) reference /
            // pointer because the underlying protos hold move-only members and
            // cannot be copied; nanobind hands us references to the Python-
            // owned instances which we use directly as backing storage for
            // the inference context.
            // Early fail if node is badly defined - may throw ValidationError.
            schema->Verify(node);
            if (opset_imports.empty()) {
              opset_imports[schema->domain()] = schema->SinceVersion();
            }

            // Adapt the bound maps to the API expected by the inference
            // context, which takes pointers to const protos for the data
            // inputs and pointers to mutable TypeProto for the value-types.
            std::unordered_map<std::string, const TensorProto *> input_data_by_name;
            input_data_by_name.reserve(input_data.size());
            for (const auto &kv : input_data) {
              input_data_by_name[kv.first] = kv.second;
            }
            std::unordered_map<std::string, const SparseTensorProto *> input_sparse_data_by_name;
            input_sparse_data_by_name.reserve(input_sparse_data.size());
            for (const auto &kv : input_sparse_data) {
              input_sparse_data_by_name[kv.first] = kv.second;
            }

            shape_inference::GraphInferenceContext graph_inference_context(
                input_types, opset_imports,
                /*symbol_table_in=*/nullptr,
                /*model_local_functions_in=*/{},
                /*schema_registry_in=*/OpSchemaRegistry::Instance(),
                /*generated_shape_data_by_name_in=*/nullptr,
                /*ir_version_in=*/ir_version);

            // Construct inference context and get results - may throw
            // InferenceError.  We always use the default options here; if it
            // is desirable for infer_node_outputs to provide check_type,
            // strict_mode, data_prop, we can add them to the Python API.
            ShapeInferenceOptions options(/*check_type_val=*/false,
                                          /*strict_mode_val=*/0,
                                          /*data_prop_val=*/false);
            shape_inference::InferenceContextImpl ctx(
                node, input_types, input_data_by_name, input_sparse_data_by_name, options,
                /*generatedShapeData=*/nullptr, &graph_inference_context);

            schema->GetTypeAndShapeInferenceFunction()(ctx);
            // Verify the inference succeeded - may also throw
            // ValidationError.  Note that input types were not validated
            // until now (except that their count was correct).
            schema->CheckInputOutputType(ctx);

            std::unordered_map<std::string, TypeProto> result;
            for (size_t i = 0; i < ctx.allOutputTypes_.size(); ++i) {
              auto &proto = ctx.allOutputTypes_[i];
              if (proto.value_case() != TypeProto::VALUE_NOT_SET) {
                result.emplace(std::string(node.output(static_cast<int>(i))), std::move(proto));
              }
            }
            return result;
          },
          nb::arg("node"), nb::arg("input_types"),
          nb::arg("input_data") = std::unordered_map<std::string, TensorProto *>{},
          nb::arg("input_sparse_data") = std::unordered_map<std::string, SparseTensorProto *>{},
          nb::arg("opset_imports") = std::unordered_map<std::string, int>{},
          nb::arg("ir_version") = static_cast<int>(IR_VERSION),
          "Runs type and shape inference for a single node and returns output TypeProto map.");

  defs.def("register_onnx_operator_set_schema", &RegisterAllOnnxOperatorSchemas,
           "Registers all built-in ONNX operator schemas with type-and-shape inference "
           "functions across all opset versions.  Duplicate registrations are silently "
           "ignored so the function is safe to call more than once.")
      .def(
          "has_schema",
          [](const std::string &op_type, const std::string &domain) -> bool {
            return OpSchemaRegistry::Schema(op_type, domain) != nullptr;
          },
          nb::arg("op_type"), nb::arg("domain") = ONNX_DOMAIN)
      .def(
          "has_schema",
          [](const std::string &op_type, int max_inclusive_version,
             const std::string &domain) -> bool {
            // should we register if nothing has been done?
            return OpSchemaRegistry::Schema(op_type, max_inclusive_version, domain) != nullptr;
          },
          nb::arg("op_type"), nb::arg("max_inclusive_version"), nb::arg("domain") = ONNX_DOMAIN)
      .def("schema_version_map",
           []() -> std::unordered_map<std::string, std::pair<int, int>> {
             return OpSchemaRegistry::DomainToVersionRange::Instance().Map();
           })
      .def(
          "get_schema",
          [](const std::string &op_type, const int max_inclusive_version,
             const std::string &domain) -> OpSchema {
            const auto *const schema =
                OpSchemaRegistry::Schema(op_type, max_inclusive_version, domain);
            if (!schema) {
              fail_schema("No schema registered for '", op_type, "' version '",
                          max_inclusive_version, "' and domain '", domain, "'!");
            }
            return *schema;
          },
          nb::arg("op_type"), nb::arg("max_inclusive_version"), nb::arg("domain") = ONNX_DOMAIN,
          "Returns the schema of *op_type* for a specific version.")
      .def(
          "get_schema",
          [](const std::string &op_type, const std::string &domain) -> OpSchema {
            const auto *const schema = OpSchemaRegistry::Schema(op_type, domain);
            // should we register if nothing has been done?
            if (!schema) {
              fail_schema("No schema registered for '", op_type, "' and domain '", domain, "'");
            }
            return *schema;
          },
          nb::arg("op_type"), nb::arg("domain") = ONNX_DOMAIN,
          "Returns the latest schema of *op_type*.")
      .def(
          "get_all_schemas",
          []() -> std::vector<OpSchema> { return OpSchemaRegistry::get_all_schemas(); },
          "Returns the schema of all registered operators at their latest version.")
      .def(
          "get_all_schemas_with_history",
          []() -> std::vector<OpSchema> {
            return OpSchemaRegistry::get_all_schemas_with_history();
          },
          "Returns the schema of all registered operators across all versions.")
      .def(
          "set_domain_to_version",
          [](const std::string &domain, int min_version, int max_version,
             int last_release_version) {
            auto &obj = OpSchemaRegistry::DomainToVersionRange::Instance();
            if (obj.Map().count(domain) == 0) {
              obj.AddDomainToVersion(domain, min_version, max_version, last_release_version);
            } else {
              obj.UpdateDomainToVersion(domain, min_version, max_version, last_release_version);
            }
          },
          nb::arg("domain"), nb::arg("min_version"), nb::arg("max_version"),
          nb::arg("last_release_version") = -1)
      .def(
          "register_schema",
          [](OpSchema schema) { RegisterSchema(std::move(schema), 0, true, true); },
          nb::arg("schema"), "Registers a user-provided OpSchema.")
      .def("deregister_schema", &DeregisterSchema, nb::arg("op_type"), nb::arg("version"),
           nb::arg("domain"), "Deregisters the specified OpSchema.");

  // -----------------------------------------------------------------------
  // Submodule `checker`
  // -----------------------------------------------------------------------
  auto checker_mod = m.def_submodule("checker");
  checker_mod.doc() = "Checker submodule";

  // nb::exception registers a new Python exception class and maps C++ throws of that type to it.
  // The RAII object is intentionally discarded after the registration side-effect completes.
  nb::exception<checker::ValidationError>(
      checker_mod,
      "ValidationError"); // NOLINT(bugprone-unused-raii,bugprone-throw-keyword-missing)

  checker_mod.def(
      "check_model",
      [](const ModelProto &model, bool full_check, bool skip_opset_compatibility_check,
         bool check_custom_domain, bool skip_external_data_location_check) {
        checker::check_model(model, full_check, skip_opset_compatibility_check, check_custom_domain,
                             skip_external_data_location_check);
      },
      nb::arg("model"), nb::arg("full_check") = false,
      nb::arg("skip_opset_compatibility_check") = false, nb::arg("check_custom_domain") = false,
      nb::arg("skip_external_data_location_check") = false,
      "Validates the model for structural correctness, topological ordering of nodes, "
      "and schema compliance, raising ValidationError on failures.");

  checker_mod.def(
      "check_function_call_cycles",
      [](const ModelProto &model) { inliner::CheckFunctionCallCycles(model); }, nb::arg("model"),
      "Checks for cycles in model-local function call graph. Raises ValidationError if a cycle is "
      "found.");

  // -----------------------------------------------------------------------
  // Submodule `inliner`
  // -----------------------------------------------------------------------
  auto inliner_mod = m.def_submodule("inliner");
  inliner_mod.doc() = "Inliner submodule";

  inliner_mod.def(
      "inline_local_functions",
      [](const ModelProto &model, bool convert_version) {
        inliner::CheckFunctionCallCycles(model);
        ModelProto copy;
        copy.CopyFrom(model);
        inliner::InlineLocalFunctions(copy, convert_version);
        return copy;
      },
      nb::arg("model"), nb::arg("convert_version") = false,
      "Inlines all model-local functions. Returns a new model with functions inlined. Raises "
      "checker.ValidationError if the model contains cyclic function references.");

  inliner_mod.def(
      "inline_selected_local_functions",
      [](const ModelProto &model,
         const std::vector<std::pair<std::string, std::string>> &function_ids, bool invert) {
        inliner::FunctionIdVector ids(function_ids.begin(), function_ids.end());
        auto id_set = inliner::FunctionIdSet::Create(std::move(ids), invert);
        ModelProto copy;
        copy.CopyFrom(model);
        inliner::InlineSelectedLocalFunctions(copy, *id_set);
        return copy;
      },
      nb::arg("model"), nb::arg("function_ids"), nb::arg("invert") = false,
      "Inlines the specified model-local functions. If invert is True, inlines all functions "
      "except those listed. Returns a new model with functions inlined.");

  inliner_mod.def(
      "inline_selected_functions",
      [](const ModelProto &model,
         const std::vector<std::pair<std::string, std::string>> &function_ids, bool invert) {
        inliner::FunctionIdVector ids(function_ids.begin(), function_ids.end());
        auto id_set = inliner::FunctionIdSet::Create(std::move(ids), invert);
        ModelProto copy;
        copy.CopyFrom(model);
        inliner::InlineSelectedFunctions(copy, *id_set, nullptr);
        return copy;
      },
      nb::arg("model"), nb::arg("function_ids"), nb::arg("invert") = false,
      "Inlines the specified functions including schema-defined functions. If invert is True, "
      "inlines all functions except those listed. Returns a new model with functions inlined.");

  // -----------------------------------------------------------------------
  // Submodule `compose`
  // -----------------------------------------------------------------------
  auto compose_mod = m.def_submodule("compose");
  compose_mod.doc() = "Compose submodule – graph/model composition utilities";

  compose_mod.def(
      "check_overlapping_names",
      [](const GraphProto &g1, const GraphProto &g2,
         const std::vector<std::pair<std::string, std::string>> &io_map)
          -> std::vector<std::pair<std::string, std::vector<std::string>>> {
        return CheckOverlappingNames(g1, g2, io_map);
      },
      nb::arg("g1"), nb::arg("g2"),
      nb::arg("io_map") = std::vector<std::pair<std::string, std::string>>{},
      "Checks whether two graphs have overlapping names. Returns a list of (category, names) pairs "
      "for each category where names appear in both graphs. Recognised categories are 'edge', "
      "'value_info', 'initializer' and 'sparse_initializer'. The optional io_map lists "
      "output/input pairs that are intentionally shared; those overlaps are excluded.");

  compose_mod.def(
      "add_prefix_graph",
      [](const GraphProto &graph, const std::string &prefix, bool rename_nodes, bool rename_edges,
         bool rename_inputs, bool rename_outputs, bool rename_initializers,
         bool rename_value_infos) -> GraphProto {
        return AddPrefixGraph(graph, prefix, rename_nodes, rename_edges, rename_inputs,
                              rename_outputs, rename_initializers, rename_value_infos);
      },
      nb::arg("graph"), nb::arg("prefix"), nb::arg("rename_nodes") = true,
      nb::arg("rename_edges") = true, nb::arg("rename_inputs") = true,
      nb::arg("rename_outputs") = true, nb::arg("rename_initializers") = true,
      nb::arg("rename_value_infos") = true,
      "Adds a prefix to names of elements in a graph. Returns a new GraphProto with prefixed "
      "names. "
      "Empty names are never prefixed.");

  compose_mod.def(
      "add_prefix",
      [](const ModelProto &model, const std::string &prefix, bool rename_nodes, bool rename_edges,
         bool rename_inputs, bool rename_outputs, bool rename_initializers, bool rename_value_infos,
         bool rename_functions) -> ModelProto {
        return AddPrefix(model, prefix, rename_nodes, rename_edges, rename_inputs, rename_outputs,
                         rename_initializers, rename_value_infos, rename_functions);
      },
      nb::arg("model"), nb::arg("prefix"), nb::arg("rename_nodes") = true,
      nb::arg("rename_edges") = true, nb::arg("rename_inputs") = true,
      nb::arg("rename_outputs") = true, nb::arg("rename_initializers") = true,
      nb::arg("rename_value_infos") = true, nb::arg("rename_functions") = true,
      "Adds a prefix to names of elements in a model. Returns a new ModelProto with prefixed "
      "names. "
      "Empty names are never prefixed.");

  compose_mod.def(
      "merge_graphs",
      [](const GraphProto &g1, const GraphProto &g2,
         const std::vector<std::pair<std::string, std::string>> &io_map,
         const std::vector<std::string> &inputs, const std::vector<std::string> &outputs,
         const std::string &prefix1, const std::string &prefix2, const std::string &name,
         const std::string &doc_string) -> GraphProto {
        return MergeGraphs(g1, g2, io_map, inputs, outputs, prefix1, prefix2, name, doc_string);
      },
      nb::arg("g1"), nb::arg("g2"), nb::arg("io_map"),
      nb::arg("inputs") = std::vector<std::string>{},
      nb::arg("outputs") = std::vector<std::string>{}, nb::arg("prefix1") = "",
      nb::arg("prefix2") = "", nb::arg("name") = "", nb::arg("doc_string") = "",
      "Combines two ONNX graphs into a single one. io_map is a list of (output_name, input_name) "
      "pairs connecting outputs of g1 to inputs of g2.");

  compose_mod.def(
      "merge_models",
      [](const ModelProto &m1, const ModelProto &m2,
         const std::vector<std::pair<std::string, std::string>> &io_map,
         const std::vector<std::string> &inputs, const std::vector<std::string> &outputs,
         const std::string &prefix1, const std::string &prefix2, const std::string &name,
         const std::string &doc_string, const std::string &producer_name,
         const std::string &producer_version, const std::string &domain,
         int64_t model_version) -> ModelProto {
        return MergeModels(m1, m2, io_map, inputs, outputs, prefix1, prefix2, name, doc_string,
                           producer_name, producer_version, domain, model_version);
      },
      nb::arg("m1"), nb::arg("m2"), nb::arg("io_map"),
      nb::arg("inputs") = std::vector<std::string>{},
      nb::arg("outputs") = std::vector<std::string>{}, nb::arg("prefix1") = "",
      nb::arg("prefix2") = "", nb::arg("name") = "", nb::arg("doc_string") = "",
      nb::arg("producer_name") = "onnx_light.onnx.compose.merge_models",
      nb::arg("producer_version") = "1.0", nb::arg("domain") = "",
      nb::arg("model_version") = int64_t{1},
      "Combines two ONNX models into a single one. Both models must share the same IR version and "
      "operator sets. io_map is a list of (output_name, input_name) pairs connecting outputs of m1 "
      "to inputs of m2.");

  compose_mod.def(
      "expand_out_dim_graph",
      [](const GraphProto &graph, int64_t dim_idx) -> GraphProto {
        return ExpandOutDimGraph(graph, dim_idx);
      },
      nb::arg("graph"), nb::arg("dim_idx"),
      "Inserts an extra dimension with extent 1 to each output in the graph. Appends an Unsqueeze "
      "node for each output.");

  compose_mod.def(
      "expand_out_dim",
      [](const ModelProto &model, int64_t dim_idx) -> ModelProto {
        return ExpandOutDim(model, dim_idx);
      },
      nb::arg("model"), nb::arg("dim_idx"),
      "Inserts an extra dimension with extent 1 to each output in the model.");
}
