# We still need to import the message definition (ModelProto) from the onnx library,
# as this IS the class generated by protobuf for ONNX models.
import os
import sys

from onnx import ModelProto

# Replace with the actual path to your ONNX model
onnx_model_path = "docs/onnx/examples/nnx/classification_head.onnx"
# Example using one of the previous models:
# onnx_model_path = "function_with_intermediate_shape.onnx"

print(f"Attempting to read ONNX file as raw protobuf: {onnx_model_path}")

# --- Reading Protobuf Data ---

# Check if file exists before reading
if not os.path.exists(onnx_model_path):
    print(f"Error: File '{onnx_model_path}' not found.")
    sys.exit(1)

# Read the entire binary content of the .onnx file
print("Reading raw bytes from file...")
try:
    with open(onnx_model_path, "rb") as f:  # Open in binary read mode ('rb')
        onnx_bytes = f.read()
    print(f"Read {len(onnx_bytes)} bytes.")
except OSError as e:
    print(f"Error reading file: {e}")
    sys.exit(1)

# Create an empty ModelProto object
# ModelProto is the Python class generated by protoc from onnx.proto
model_proto = ModelProto()
print("Created empty onnx.ModelProto object.")

# Parse the raw bytes into the ModelProto object
# This uses the protobuf library's deserialization logic directly.
print("Parsing bytes into ModelProto object using ParseFromString...")
try:
    model_proto.ParseFromString(onnx_bytes)
    print("Parsing successful.")
except Exception as e:
    # Errors here usually mean the file is corrupted, not a valid ONNX file,
    # or potentially incompatible protobuf versions (less common).
    print(f"\nError during protobuf parsing: {e}")
    print("The file might be corrupted or not a valid ONNX protobuf file.")
    sys.exit(1)

# --- Inspecting the Raw Protobuf Object ---
print("\n--- Raw ModelProto Object Content (Protobuf Structure) ---")

# Printing the object itself gives the default protobuf text representation.
# This shows the raw fields and values as defined in onnx.proto,
# without the helper methods or higher-level structure provided by onnx.load().
print(model_proto)

# You can also access fields directly using protobuf object syntax:
# print("\n--- Accessing Specific Fields ---")
# print(f"IR Version: {model_proto.ir_version}")
# print(f"Producer Name: {model_proto.producer_name}")
# print(f"Graph Name: {model_proto.graph.name}")
# print(f"Number of nodes in graph: {len(model_proto.graph.node)}")
# # Accessing deeper structures requires navigating the object hierarchy
# if len(model_proto.graph.node) > 0:
#      print(f"First node's op_type: {model_proto.graph.node[0].op_type}")
#      print(f"First node's name: {model_proto.graph.node[0].name}")
#      print(f"First node's inputs: {model_proto.graph.node[0].input}")
#      print(f"First node's outputs: {model_proto.graph.node[0].output}")
