import sys
import os
import time

# Ensure pytrace is in python path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from pytrace.api.node import TraceNode
from pytrace.api.decorators import register_class, register_node, input_port, output_port
from pytrace.runtime.context import TraceContext
from pytrace.remote.worker import TraceWorker, expose_node

# =============================================================================
# 1. Define Your Custom Nodes
# =============================================================================
# Use @expose_node to make the class available to the Remote Server.
# Use @register_class to define the category and display name in the UI.

@expose_node
@register_class(category="Tutorial", display_name="Client Demo")
class ClientDemoNode(TraceNode):
    """
    A collection of demo nodes running on the client side.
    Developers can add their own methods here.
    """
    
    @register_node(display_name="Remote Add", description="Performs addition on the client.")
    @input_port(name="a", type="float")
    @input_port(name="b", type="float")
    @output_port(name="result", type="float")
    def add(self, ctx: TraceContext):
        # 1. Get Inputs
        a = ctx.get_input("a", 0.0)
        b = ctx.get_input("b", 0.0)
        
        # 2. Execute Logic
        result = a + b
        
        # 3. Log execution (This log will appear on the Server console!)
        ctx.log(f"[Client] Executing Add: {a} + {b} = {result}")
        
        # 4. Set Output
        ctx.set_output("result", result)

    @register_node(display_name="Remote Echo", description="Echoes a message.")
    @input_port(name="message", type="string")
    @output_port(name="echo", type="string")
    def echo(self, ctx: TraceContext):
        msg = ctx.get_input("message", "")
        ctx.log(f"[Client] Echoing: {msg}")
        ctx.set_output("echo", f"Client says: {msg}")

# =============================================================================
# 2. Start the Worker
# =============================================================================

if __name__ == "__main__":
    SERVER_HOST = 'localhost'
    SERVER_PORT = 9999
    
    print(f"--- TraceStudio Client Demo ---")
    print(f"Connecting to Server at {SERVER_HOST}:{SERVER_PORT}...")
    
    # Initialize the Worker
    worker = TraceWorker(host=SERVER_HOST, port=SERVER_PORT)
    
    try:
        # Start the worker (This blocks until stopped)
        worker.start()
    except KeyboardInterrupt:
        print("\n[Client] Worker stopped.")
    except Exception as e:
        print(f"\n[Client] Error: {e}")