1 files changed, 381 insertions, 0 deletions
diff --git a/wss/main.go b/wss/main.go
new file mode 100755
index 0000000..a3ea063
--- /dev/null
+++ b/wss/main.go
@@ -0,0 +1,381 @@
+package main
+
+import (
+	"encoding/json"
+	"flag"
+	"fmt"
+	"log"
+	"net/http"
+	"os"
+	"sync"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+// Configuration for the server
+type Config struct {
+	Debug      bool
+	Port       int
+	StaticDir  string
+}
+
+// Environment variables to pass to frontend
+type EnvVars struct {
+	Thing int `json:"thing"`
+}
+
+func getEnvVars() EnvVars {
+	return EnvVars{Thing: 1}
+}
+
+// WebSocket connection upgrader
+var upgrader = websocket.Upgrader{
+	ReadBufferSize:  1024,
+	WriteBufferSize: 1024,
+	CheckOrigin: func(r *http.Request) bool {
+		return true // Allow all connections (restrict in production)
+	},
+}
+
+// Connection represents a WebSocket connection
+type Connection struct {
+	ws   *websocket.Conn
+	send chan []byte
+}
+
+// Node represents a graph node
+type Node struct {
+	ID int     `json:"id"`
+	X  float64 `json:"x"`
+	Y  float64 `json:"y"`
+	Z  float64 `json:"z"`
+}
+
+// Edge represents a graph edge
+type Edge struct {
+	From int `json:"from"`
+	To   int `json:"to"`
+}
+
+// Graph holds the shared graph state
+type Graph struct {
+	Nodes []Node `json:"nodes"`
+	Edges []Edge `json:"edges"`
+}
+
+// Hub maintains the set of active connections and broadcasts messages
+type Hub struct {
+	connections map[*Connection]bool
+	broadcast   chan []byte
+	register    chan *Connection
+	unregister  chan *Connection
+	graph       Graph
+	graphMutex  sync.RWMutex
+}
+
+// Message types
+const (
+	MessageTypeAddNode    = "addNode"
+	MessageTypeAddEdge    = "addEdge"
+	MessageTypeRemoveNode = "removeNode"
+	MessageTypeFullSync   = "fullSync"
+)
+
+// WebSocket message format
+type WebSocketMessage struct {
+	Type  string  `json:"type"`
+	Node  *Node   `json:"node,omitempty"`
+	Edge  *Edge   `json:"edge,omitempty"`
+	Graph *Graph  `json:"graph,omitempty"`
+}
+
+var hub = Hub{
+	broadcast:   make(chan []byte),
+	register:    make(chan *Connection),
+	unregister:  make(chan *Connection),
+	connections: make(map[*Connection]bool),
+	graph:       createInitialGraph(),
+}
+
+func createInitialGraph() Graph {
+	nodes := []Node{
+		{ID: 0, X: 0, Y: 0, Z: 0},
+		{ID: 1, X: 3, Y: 2, Z: 1},
+		{ID: 2, X: -3, Y: 2, Z: -1},
+		{ID: 3, X: 2, Y: -2, Z: 2},
+		{ID: 4, X: -2, Y: -2, Z: -2},
+		{ID: 5, X: 0, Y: 3, Z: -2},
+	}
+	edges := []Edge{
+		{From: 0, To: 1},
+		{From: 0, To: 2},
+		{From: 0, To: 3},
+		{From: 0, To: 4},
+		{From: 1, To: 5},
+		{From: 2, To: 5},
+		{From: 3, To: 4},
+	}
+	return Graph{Nodes: nodes, Edges: edges}
+}
+
+func (h *Hub) run() {
+	for {
+		select {
+		case conn := <-h.register:
+			h.connections[conn] = true
+			log.Printf("Client connected. Total connections: %d", len(h.connections))
+
+			// Send current graph state to new connection
+			h.graphMutex.RLock()
+			fullSync := WebSocketMessage{
+				Type:  MessageTypeFullSync,
+				Graph: &h.graph,
+			}
+			h.graphMutex.RUnlock()
+
+			data, err := json.Marshal(fullSync)
+			if err == nil {
+				conn.send <- data
+			}
+
+		case conn := <-h.unregister:
+			if _, ok := h.connections[conn]; ok {
+				delete(h.connections, conn)
+				close(conn.send)
+				log.Printf("Client disconnected. Total connections: %d", len(h.connections))
+			}
+
+		case message := <-h.broadcast:
+			// Parse and apply the operation to server state
+			var msg WebSocketMessage
+			if err := json.Unmarshal(message, &msg); err != nil {
+				log.Printf("Error parsing message: %v", err)
+				continue
+			}
+
+			response := h.handleOperation(msg)
+			if response == nil {
+				continue
+			}
+
+			// Broadcast the response to all connections
+			data, err := json.Marshal(response)
+			if err != nil {
+				log.Printf("Error marshaling response: %v", err)
+				continue
+			}
+
+			for conn := range h.connections {
+				select {
+				case conn.send <- data:
+				default:
+					close(conn.send)
+					delete(h.connections, conn)
+				}
+			}
+		}
+	}
+}
+
+// handleOperation processes graph operations and returns the message to broadcast
+func (h *Hub) handleOperation(msg WebSocketMessage) *WebSocketMessage {
+	h.graphMutex.Lock()
+	defer h.graphMutex.Unlock()
+
+	switch msg.Type {
+	case MessageTypeAddNode:
+		if msg.Node == nil {
+			return nil
+		}
+		// Assign server-side ID to avoid conflicts
+		newID := 0
+		for _, n := range h.graph.Nodes {
+			if n.ID >= newID {
+				newID = n.ID + 1
+			}
+		}
+		node := Node{
+			ID: newID,
+			X:  msg.Node.X,
+			Y:  msg.Node.Y,
+			Z:  msg.Node.Z,
+		}
+		h.graph.Nodes = append(h.graph.Nodes, node)
+		log.Printf("Added node %d at (%.2f, %.2f, %.2f)", newID, node.X, node.Y, node.Z)
+		return &WebSocketMessage{Type: MessageTypeAddNode, Node: &node}
+
+	case MessageTypeAddEdge:
+		if msg.Edge == nil {
+			return nil
+		}
+		// Validate edge
+		if msg.Edge.From < 0 || msg.Edge.From >= len(h.graph.Nodes) ||
+			msg.Edge.To < 0 || msg.Edge.To >= len(h.graph.Nodes) ||
+			msg.Edge.From == msg.Edge.To {
+			return nil
+		}
+		// Check for duplicate
+		for _, e := range h.graph.Edges {
+			if (e.From == msg.Edge.From && e.To == msg.Edge.To) ||
+				(e.From == msg.Edge.To && e.To == msg.Edge.From) {
+				return nil
+			}
+		}
+		edge := Edge{From: msg.Edge.From, To: msg.Edge.To}
+		h.graph.Edges = append(h.graph.Edges, edge)
+		log.Printf("Added edge %d -> %d", edge.From, edge.To)
+		return &WebSocketMessage{Type: MessageTypeAddEdge, Edge: &edge}
+
+	case MessageTypeRemoveNode:
+		if msg.Node == nil || msg.Node.ID < 0 || msg.Node.ID >= len(h.graph.Nodes) {
+			return nil
+		}
+		// For simplicity, mark node as removed by setting special coordinates
+		// A full implementation would handle ID remapping
+		nodeID := msg.Node.ID
+		// Remove edges connected to this node
+		newEdges := []Edge{}
+		for _, e := range h.graph.Edges {
+			if e.From != nodeID && e.To != nodeID {
+				newEdges = append(newEdges, e)
+			}
+		}
+		h.graph.Edges = newEdges
+		log.Printf("Removed node %d", nodeID)
+		// Send full sync after removal for simplicity
+		return &WebSocketMessage{Type: MessageTypeFullSync, Graph: &h.graph}
+
+	case "reset":
+		h.graph = Graph{Nodes: []Node{}, Edges: []Edge{}}
+		log.Printf("Graph reset")
+		return &WebSocketMessage{Type: MessageTypeFullSync, Graph: &h.graph}
+	}
+
+	return nil
+}
+
+// WebSocket handler
+func serveWebSocket(w http.ResponseWriter, r *http.Request) {
+	// Upgrade to WebSocket
+	ws, err := upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		log.Printf("WebSocket upgrade error: %v", err)
+		return
+	}
+	
+	// Create a new connection
+	conn := &Connection{
+		ws:   ws,
+		send: make(chan []byte, 256),
+	}
+	
+	// Register the connection with the hub
+	hub.register <- conn
+	
+	// Start the connection handlers
+	go conn.writer()
+	conn.reader()
+}
+
+// Writer goroutine for connection
+func (c *Connection) writer() {
+	// Ensure clean close of connection
+	defer func() {
+		c.ws.Close()
+	}()
+	
+	for {
+		select {
+		case message, ok := <-c.send:
+			if !ok {
+				c.ws.WriteMessage(websocket.CloseMessage, []byte{})
+				return
+			}
+			
+			// Write the message
+			if err := c.ws.WriteMessage(websocket.TextMessage, message); err != nil {
+				return
+			}
+		}
+	}
+}
+
+// Reader goroutine for connection
+func (c *Connection) reader() {
+	// Ensure clean close of connection and unregister
+	defer func() {
+		hub.unregister <- c
+		c.ws.Close()
+	}()
+	
+	// Set read deadline
+	c.ws.SetReadDeadline(time.Now().Add(60 * time.Second))
+	c.ws.SetPongHandler(func(string) error {
+		c.ws.SetReadDeadline(time.Now().Add(60 * time.Second))
+		return nil
+	})
+	
+	// Read messages from the connection
+	for {
+		_, message, err := c.ws.ReadMessage()
+		if err != nil {
+			break
+		}
+		
+		// Broadcast the message to all connections
+		hub.broadcast <- message
+	}
+}
+
+
+func main() {
+	// Parse command line flags
+	config := Config{}
+	flag.IntVar(&config.Port, "port", 5000, "Port to run the server on")
+	flag.BoolVar(&config.Debug, "debug", true, "Enable debug mode")
+	flag.StringVar(&config.StaticDir, "static", "..", "Directory to serve static files from")
+	flag.Parse()
+
+	// Override port from environment variable if set
+	if portStr := os.Getenv("PORT"); portStr != "" {
+		fmt.Sscanf(portStr, "%d", &config.Port)
+	}
+
+	// Start the hub
+	go hub.run()
+
+	// Create a file server for static files (serve from parent directory)
+	fs := http.FileServer(http.Dir(config.StaticDir))
+
+	// Set up routes
+	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == "/" {
+			http.ServeFile(w, r, config.StaticDir+"/index.html")
+			return
+		}
+		fs.ServeHTTP(w, r)
+	})
+
+	// API routes
+	http.HandleFunc("/api/env", func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		json.NewEncoder(w).Encode(getEnvVars())
+	})
+
+	http.HandleFunc("/api/health", func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		json.NewEncoder(w).Encode(map[string]string{"status": "ok"})
+	})
+
+	// WebSocket endpoint
+	http.HandleFunc("/ws", serveWebSocket)
+
+	// Start the server
+	addr := fmt.Sprintf("0.0.0.0:%d", config.Port)
+	log.Printf("Server starting on http://localhost:%d", config.Port)
+	log.Printf("Serving static files from: %s", config.StaticDir)
+	if err := http.ListenAndServe(addr, nil); err != nil {
+		log.Fatalf("Server error: %v", err)
+	}
+}