import sys
import os

from pytrace.runtime.registry import NodeRegistry

# Add the project root to the Python path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from pytrace.nodes.test_math_suite import MathSuite

# Mock for NodeIO
class MockNodeIO:
    def __init__(self, inputs=None):
        self._inputs = inputs or {}
        self._outputs = {}

    def get_input(self, name, default=None):
        return self._inputs.get(name, default)

    def set_output(self, name, value):
        self._outputs[name] = value
        print(f"Output '{name}' set to: {value}")

    def get_output(self, name):
        return self._outputs.get(name)

# Mock for TraceContext
class MockTraceContext:
    def log(self, message):
        print(f"LOG: {message}")

def run_math_suite_tests():
    """
    Runs tests for the MathSuite node.
    """
    print("--- Testing MathSuite ---")

    # 1. Test the 'add' method
    print("\n--- Testing add method ---")
    add_id = MathSuite.__name__ + "." + "add"
    subtract_id = MathSuite.__name__ + "." + "subtract"
    
    node_add = NodeRegistry.create_node(add_id)
    node_subtract = NodeRegistry.create_node(subtract_id)
    
    io_add = MockNodeIO({"a": 5.0, "b": 3.0})
    io_subtract = MockNodeIO({"a": 10.0, "b": 7.0})
    
    context_add = MockTraceContext()
    context_subtract = MockTraceContext()
    
    node_add(io_add, context_add)
    node_subtract(io_subtract, context_subtract)
    
    assert io_add.get_output("result") == 8.0
    assert io_subtract.get_output("result") == 3.0

    print("\n--- All MathSuite tests passed ---")


if __name__ == "__main__":
    run_math_suite_tests()