// Copyright (C) 2026
// License: Boost Software License   See LICENSE.txt for the full license.
#ifndef DLIB_TEST_ONNX_TEST_HELPERS_H_
#define DLIB_TEST_ONNX_TEST_HELPERS_H_

#include "tester.h"

#include <cstddef>
#include <cstdint>
#include <map>
#include <string>
#include <vector>

namespace test
{
    namespace onnx
    {

    // ----------------------------------------------------------------------------------------

        /*
            Minimal ONNX reader for dependency-free structural tests.  It
            recognizes just enough ModelProto/GraphProto/NodeProto/
            AttributeProto/TensorProto wire format to inspect node op_types,
            node inputs, int/int-list/tensor attributes, and small int64
            initializers generated by dlib's ONNX exporter.
        */

        inline bool contains_substring (
            const std::string& str,
            const std::string& substr
        )
        {
            return str.find(substr) != std::string::npos;
        }

        struct onnx_tensor_info
        {
            std::string name;
            std::vector<int64_t> dims;
            std::vector<int64_t> int64_data;
        };

        struct onnx_node_info
        {
            std::string op_type;
            std::vector<std::string> inputs;
            std::map<std::string, int64_t> int_attributes;
            std::map<std::string, std::vector<int64_t> > ints_attributes;
            std::map<std::string, std::string> string_attributes;
            std::map<std::string, onnx_tensor_info> tensor_attributes;
        };

        inline bool read_onnx_varint (
            const std::string& data,
            size_t& pos,
            uint64_t& value
        )
        {
            value = 0;
            int shift = 0;
            while (pos < data.size() && shift < 64)
            {
                const unsigned char byte = static_cast<unsigned char>(data[pos++]);
                value |= static_cast<uint64_t>(byte & 0x7f) << shift;
                if ((byte & 0x80) == 0)
                    return true;
                shift += 7;
            }
            return false;
        }

        inline bool read_onnx_message (
            const std::string& data,
            size_t& pos,
            std::string& value
        )
        {
            uint64_t size = 0;
            if (!read_onnx_varint(data, pos, size))
                return false;
            if (size > data.size() - pos)
                return false;
            value.assign(data.data()+pos, static_cast<size_t>(size));
            pos += static_cast<size_t>(size);
            return true;
        }

        inline void skip_onnx_field (
            const std::string& data,
            size_t& pos,
            int wire_type
        )
        {
            uint64_t ignored = 0;
            std::string ignored_message;
            if (wire_type == 0)
                read_onnx_varint(data, pos, ignored);
            else if (wire_type == 2)
                read_onnx_message(data, pos, ignored_message);
            else
                pos = data.size();
        }

        inline onnx_tensor_info parse_onnx_tensor_metadata (
            const std::string& data
        )
        {
            onnx_tensor_info tensor;
            size_t pos = 0;
            while (pos < data.size())
            {
                uint64_t key = 0;
                if (!read_onnx_varint(data, pos, key))
                    break;
                const int field_number = static_cast<int>(key >> 3);
                const int wire_type = static_cast<int>(key & 0x07);
                if (field_number == 1 && wire_type == 0)
                {
                    uint64_t value = 0;
                    read_onnx_varint(data, pos, value);
                    tensor.dims.push_back(static_cast<int64_t>(value));
                }
                else if (field_number == 8 && wire_type == 2)
                {
                    read_onnx_message(data, pos, tensor.name);
                }
                else
                {
                    skip_onnx_field(data, pos, wire_type);
                }
            }
            return tensor;
        }

        inline void parse_onnx_attribute (
            const std::string& data,
            onnx_node_info& node
        )
        {
            std::string name;
            bool has_int = false;
            int64_t int_value = 0;
            std::vector<int64_t> ints_values;
            bool has_string = false;
            std::string string_value;
            bool has_tensor = false;
            onnx_tensor_info tensor_value;

            size_t pos = 0;
            while (pos < data.size())
            {
                uint64_t key = 0;
                if (!read_onnx_varint(data, pos, key))
                    break;
                const int field_number = static_cast<int>(key >> 3);
                const int wire_type = static_cast<int>(key & 0x07);
                if (field_number == 1 && wire_type == 2)
                {
                    read_onnx_message(data, pos, name);
                }
                else if (field_number == 3 && wire_type == 0)
                {
                    uint64_t value = 0;
                    read_onnx_varint(data, pos, value);
                    int_value = static_cast<int64_t>(value);
                    has_int = true;
                }
                else if (field_number == 4 && wire_type == 2)
                {
                    has_string = read_onnx_message(data, pos, string_value);
                }
                else if (field_number == 5 && wire_type == 2)
                {
                    std::string tensor_data;
                    if (read_onnx_message(data, pos, tensor_data))
                    {
                        tensor_value = parse_onnx_tensor_metadata(tensor_data);
                        has_tensor = true;
                    }
                }
                else if (field_number == 8 && wire_type == 0)
                {
                    uint64_t value = 0;
                    read_onnx_varint(data, pos, value);
                    ints_values.push_back(static_cast<int64_t>(value));
                }
                else
                {
                    skip_onnx_field(data, pos, wire_type);
                }
            }

            if (!name.empty() && has_int)
                node.int_attributes[name] = int_value;
            if (!name.empty() && !ints_values.empty())
                node.ints_attributes[name] = ints_values;
            if (!name.empty() && has_string)
                node.string_attributes[name] = string_value;
            if (!name.empty() && has_tensor)
                node.tensor_attributes[name] = tensor_value;
        }

        inline onnx_node_info parse_onnx_node (
            const std::string& data
        )
        {
            onnx_node_info node;
            size_t pos = 0;
            while (pos < data.size())
            {
                uint64_t key = 0;
                if (!read_onnx_varint(data, pos, key))
                    break;
                const int field_number = static_cast<int>(key >> 3);
                const int wire_type = static_cast<int>(key & 0x07);
                if (field_number == 1 && wire_type == 2)
                {
                    std::string input_name;
                    if (read_onnx_message(data, pos, input_name))
                        node.inputs.push_back(input_name);
                }
                else if (field_number == 4 && wire_type == 2)
                {
                    read_onnx_message(data, pos, node.op_type);
                }
                else if (field_number == 5 && wire_type == 2)
                {
                    std::string attribute_data;
                    if (read_onnx_message(data, pos, attribute_data))
                        parse_onnx_attribute(attribute_data, node);
                }
                else
                {
                    skip_onnx_field(data, pos, wire_type);
                }
            }
            return node;
        }

        inline int64_t decode_fixed_int64(const std::string& data, size_t offset)
        {
            uint64_t value = 0;
            for (size_t i = 0; i < 8; ++i)
                value |= static_cast<uint64_t>(static_cast<unsigned char>(data[offset+i])) << (8*i);
            return static_cast<int64_t>(value);
        }

        inline onnx_tensor_info parse_onnx_tensor (
            const std::string& data
        )
        {
            onnx_tensor_info tensor;
            size_t pos = 0;
            while (pos < data.size())
            {
                uint64_t key = 0;
                if (!read_onnx_varint(data, pos, key))
                    break;
                const int field_number = static_cast<int>(key >> 3);
                const int wire_type = static_cast<int>(key & 0x07);
                if (field_number == 1 && wire_type == 0)
                {
                    uint64_t value = 0;
                    read_onnx_varint(data, pos, value);
                    tensor.dims.push_back(static_cast<int64_t>(value));
                }
                else if (field_number == 8 && wire_type == 2)
                {
                    read_onnx_message(data, pos, tensor.name);
                }
                else if (field_number == 9 && wire_type == 2)
                {
                    std::string raw_data;
                    if (read_onnx_message(data, pos, raw_data))
                    {
                        for (size_t i = 0; i + 7 < raw_data.size(); i += 8)
                            tensor.int64_data.push_back(decode_fixed_int64(raw_data, i));
                    }
                }
                else
                {
                    skip_onnx_field(data, pos, wire_type);
                }
            }
            return tensor;
        }

        inline void parse_onnx_graph_nodes (
            const std::string& data,
            std::vector<onnx_node_info>& nodes
        )
        {
            size_t pos = 0;
            while (pos < data.size())
            {
                uint64_t key = 0;
                if (!read_onnx_varint(data, pos, key))
                    break;
                const int field_number = static_cast<int>(key >> 3);
                const int wire_type = static_cast<int>(key & 0x07);
                if (field_number == 1 && wire_type == 2)
                {
                    std::string node_data;
                    if (read_onnx_message(data, pos, node_data))
                        nodes.push_back(parse_onnx_node(node_data));
                }
                else
                {
                    skip_onnx_field(data, pos, wire_type);
                }
            }
        }

        inline void parse_onnx_graph_initializers (
            const std::string& data,
            std::vector<onnx_tensor_info>& tensors
        )
        {
            size_t pos = 0;
            while (pos < data.size())
            {
                uint64_t key = 0;
                if (!read_onnx_varint(data, pos, key))
                    break;
                const int field_number = static_cast<int>(key >> 3);
                const int wire_type = static_cast<int>(key & 0x07);
                if (field_number == 5 && wire_type == 2)
                {
                    std::string tensor_data;
                    if (read_onnx_message(data, pos, tensor_data))
                        tensors.push_back(parse_onnx_tensor(tensor_data));
                }
                else
                {
                    skip_onnx_field(data, pos, wire_type);
                }
            }
        }

        inline std::vector<onnx_node_info> parse_onnx_nodes (
            const std::string& data
        )
        {
            std::vector<onnx_node_info> nodes;
            size_t pos = 0;
            while (pos < data.size())
            {
                uint64_t key = 0;
                if (!read_onnx_varint(data, pos, key))
                    break;
                const int field_number = static_cast<int>(key >> 3);
                const int wire_type = static_cast<int>(key & 0x07);
                if (field_number == 7 && wire_type == 2)
                {
                    std::string graph_data;
                    if (read_onnx_message(data, pos, graph_data))
                        parse_onnx_graph_nodes(graph_data, nodes);
                }
                else
                {
                    skip_onnx_field(data, pos, wire_type);
                }
            }
            return nodes;
        }

        inline std::vector<onnx_tensor_info> parse_onnx_initializers (
            const std::string& data
        )
        {
            std::vector<onnx_tensor_info> tensors;
            size_t pos = 0;
            while (pos < data.size())
            {
                uint64_t key = 0;
                if (!read_onnx_varint(data, pos, key))
                    break;
                const int field_number = static_cast<int>(key >> 3);
                const int wire_type = static_cast<int>(key & 0x07);
                if (field_number == 7 && wire_type == 2)
                {
                    std::string graph_data;
                    if (read_onnx_message(data, pos, graph_data))
                        parse_onnx_graph_initializers(graph_data, tensors);
                }
                else
                {
                    skip_onnx_field(data, pos, wire_type);
                }
            }
            return tensors;
        }

        inline int count_onnx_nodes (
            const std::vector<onnx_node_info>& nodes,
            const std::string& op_type
        )
        {
            int count = 0;
            for (const auto& node : nodes)
            {
                if (node.op_type == op_type)
                    ++count;
            }
            return count;
        }

        inline int64_t get_onnx_int_attribute (
            const onnx_node_info& node,
            const std::string& name
        )
        {
            const auto found = node.int_attributes.find(name);
            DLIB_TEST_MSG(found != node.int_attributes.end(), "Missing ONNX int attribute: " << name);
            return found->second;
        }

        inline const std::vector<int64_t>& get_onnx_ints_attribute (
            const onnx_node_info& node,
            const std::string& name
        )
        {
            const auto found = node.ints_attributes.find(name);
            DLIB_TEST_MSG(found != node.ints_attributes.end(), "Missing ONNX ints attribute: " << name);
            static const std::vector<int64_t> dummy;
            if (found == node.ints_attributes.end())
                return dummy;
            return found->second;
        }

        inline const std::string& get_onnx_string_attribute (
            const onnx_node_info& node,
            const std::string& name
        )
        {
            const auto found = node.string_attributes.find(name);
            DLIB_TEST_MSG(found != node.string_attributes.end(), "Missing ONNX string attribute: " << name);
            static const std::string dummy;
            if (found == node.string_attributes.end())
                return dummy;
            return found->second;
        }

        inline const onnx_tensor_info& get_onnx_tensor_attribute (
            const onnx_node_info& node,
            const std::string& name
        )
        {
            const auto found = node.tensor_attributes.find(name);
            DLIB_TEST_MSG(found != node.tensor_attributes.end(), "Missing ONNX tensor attribute: " << name);
            static const onnx_tensor_info dummy;
            if (found == node.tensor_attributes.end())
                return dummy;
            return found->second;
        }

        inline const onnx_node_info& nth_onnx_node (
            const std::vector<onnx_node_info>& nodes,
            const std::string& op_type,
            int n
        )
        {
            for (const auto& node : nodes)
            {
                if (node.op_type == op_type && n-- == 0)
                    return node;
            }
            DLIB_TEST_MSG(false, "Missing ONNX node: " << op_type);
            static const onnx_node_info dummy;
            return dummy;
        }

        inline std::vector<int64_t> get_onnx_ints_initializer (
            const std::string& data,
            const std::string& name_substring
        )
        {
            const auto tensors = parse_onnx_initializers(data);
            for (const auto& tensor : tensors)
            {
                if (tensor.name.find(name_substring) != std::string::npos)
                    return tensor.int64_data;
            }
            DLIB_TEST_MSG(false, "Missing ONNX int64 initializer containing: " << name_substring);
            return std::vector<int64_t>();
        }

    // ----------------------------------------------------------------------------------------

    }
}

#endif // DLIB_TEST_ONNX_TEST_HELPERS_H_