From 56186e7eb67749a1679e2f27ea1cf112ef29e4c4 Mon Sep 17 00:00:00 2001
From: Cyber MacGeddon <cybermaggedon@gmail.com>
Date: Mon, 23 Jun 2025 23:43:08 +0100
Subject: [PATCH] Configurable rev_gateway

---
 .../trustgraph/rev_gateway/__main__.py        | 56 ++++++++++++++++++-
 .../trustgraph/rev_gateway/service.py         | 33 +++++++++--
 2 files changed, 83 insertions(+), 6 deletions(-)

diff --git a/trustgraph-flow/trustgraph/rev_gateway/__main__.py b/trustgraph-flow/trustgraph/rev_gateway/__main__.py
index 69fa0895..49f77d77 100644
--- a/trustgraph-flow/trustgraph/rev_gateway/__main__.py
+++ b/trustgraph-flow/trustgraph/rev_gateway/__main__.py
@@ -1,6 +1,8 @@
 import asyncio
+import argparse
 import logging
 import sys
+import os
 from .service import ReverseGateway
 
 logging.basicConfig(
@@ -8,8 +10,60 @@ logging.basicConfig(
     format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
 )
 
+def parse_args():
+    parser = argparse.ArgumentParser(
+        prog="reverse-gateway",
+        description="TrustGraph Reverse Gateway - WebSocket to Pulsar bridge"
+    )
+    
+    parser.add_argument(
+        '--websocket-uri',
+        default=None,
+        help='WebSocket URI to connect to (default: wss://api.trustgraph.ai/ws or WEBSOCKET_URI env var)'
+    )
+    
+    parser.add_argument(
+        '--max-workers',
+        type=int,
+        default=10,
+        help='Maximum concurrent message handlers (default: 10)'
+    )
+    
+    parser.add_argument(
+        '--pulsar-host',
+        default=None,
+        help='Pulsar host URL (default: pulsar://pulsar:6650 or PULSAR_HOST env var)'
+    )
+    
+    parser.add_argument(
+        '--pulsar-api-key',
+        default=None,
+        help='Pulsar API key for authentication (default: PULSAR_API_KEY env var)'
+    )
+    
+    parser.add_argument(
+        '--pulsar-listener',
+        default=None,
+        help='Pulsar listener name'
+    )
+    
+    return parser.parse_args()
+
 async def main():
-    gateway = ReverseGateway()
+    args = parse_args()
+    
+    gateway = ReverseGateway(
+        websocket_uri=args.websocket_uri,
+        max_workers=args.max_workers,
+        pulsar_host=args.pulsar_host,
+        pulsar_api_key=args.pulsar_api_key,
+        pulsar_listener=args.pulsar_listener
+    )
+    
+    print(f"Starting reverse gateway:")
+    print(f"  WebSocket URI: {gateway.url}")
+    print(f"  Max workers: {args.max_workers}")
+    print(f"  Pulsar host: {gateway.pulsar_host}")
     
     try:
         await gateway.run()
diff --git a/trustgraph-flow/trustgraph/rev_gateway/service.py b/trustgraph-flow/trustgraph/rev_gateway/service.py
index 66b6ed92..067a8e45 100644
--- a/trustgraph-flow/trustgraph/rev_gateway/service.py
+++ b/trustgraph-flow/trustgraph/rev_gateway/service.py
@@ -5,6 +5,7 @@ import sys
 import os
 from aiohttp import ClientSession, WSMsgType, ClientWebSocketResponse
 from typing import Optional
+from urllib.parse import urlparse, urlunparse
 import pulsar
 
 from .dispatcher import MessageDispatcher
@@ -15,11 +16,33 @@ logger.setLevel(logging.INFO)
 
 class ReverseGateway:
     
-    def __init__(self, host: str = "api.trustgraph.ai", max_workers: int = 10, 
+    def __init__(self, websocket_uri: str = None, max_workers: int = 10, 
                  pulsar_host: str = None, pulsar_api_key: str = None, 
                  pulsar_listener: str = None):
-        self.host = host
-        self.url = f"wss://{host}/ws"
+        # Set default WebSocket URI with environment variable support
+        if websocket_uri is None:
+            websocket_uri = os.getenv("WEBSOCKET_URI", "wss://api.trustgraph.ai/ws")
+        
+        # Parse and validate the WebSocket URI
+        parsed_uri = urlparse(websocket_uri)
+        if parsed_uri.scheme not in ('ws', 'wss'):
+            raise ValueError(f"WebSocket URI must use ws:// or wss:// scheme, got: {parsed_uri.scheme}")
+        if not parsed_uri.netloc:
+            raise ValueError(f"WebSocket URI must include hostname, got: {websocket_uri}")
+        
+        # Store parsed components for debugging/logging
+        self.websocket_uri = websocket_uri
+        self.host = parsed_uri.hostname
+        self.port = parsed_uri.port
+        self.scheme = parsed_uri.scheme
+        self.path = parsed_uri.path or "/ws"
+        
+        # Construct the full URL (in case path was missing)
+        if not parsed_uri.path:
+            self.url = f"{self.scheme}://{parsed_uri.netloc}/ws"
+        else:
+            self.url = websocket_uri
+            
         self.max_workers = max_workers
         self.ws: Optional[ClientWebSocketResponse] = None
         self.session: Optional[ClientSession] = None
@@ -57,11 +80,11 @@ class ReverseGateway:
                 
             logger.info(f"Connecting to {self.url}")
             self.ws = await self.session.ws_connect(self.url)
-            logger.info("WebSocket connection established")
+            logger.info(f"WebSocket connection established to {self.host}:{self.port or 'default'}")
             return True
             
         except Exception as e:
-            logger.error(f"Failed to connect: {e}")
+            logger.error(f"Failed to connect to {self.url}: {e}")
             return False
     
     async def disconnect(self):