WebUITest-RealProjecT/Handler/Project/WebUI/WebSocketServer.cs

using System;
using System.Collections.Generic;
using System.Linq;
using System.Net;
using System.Net.WebSockets;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using System.Windows.Forms;
using Project.Dialog;

namespace Project.WebUI
{
    public class WebSocketServer
    {
        private HttpListener _httpListener;
        private List<WebSocket> _clients = new List<WebSocket>();
        private fWebView _mainForm;

        public WebSocketServer(string url, fWebView form)
        {
            _mainForm = form;
            _httpListener = new HttpListener();
            _httpListener.Prefixes.Add(url);
            _httpListener.Start();
            Console.WriteLine($"[WS] Listening on {url}");
            Task.Run(AcceptConnections);
        }

        private async Task AcceptConnections()
        {
            while (_httpListener.IsListening)
            {
                try
                {
                    var context = await _httpListener.GetContextAsync();
                    if (context.Request.IsWebSocketRequest)
                    {
                        ProcessRequest(context);
                    }
                    else
                    {
                        context.Response.StatusCode = 400;
                        context.Response.Close();
                    }
                }
                catch (Exception ex)
                {
                    Console.WriteLine($"[WS] Error: {ex.Message}");
                }
            }
        }

        private System.Collections.Concurrent.ConcurrentDictionary<WebSocket, SemaphoreSlim> _socketLocks = new System.Collections.Concurrent.ConcurrentDictionary<WebSocket, SemaphoreSlim>();

        private async void ProcessRequest(HttpListenerContext context)
        {
            WebSocketContext wsContext = null;
            try
            {
                wsContext = await context.AcceptWebSocketAsync(subProtocol: null);
                WebSocket socket = wsContext.WebSocket;
                _socketLocks.TryAdd(socket, new SemaphoreSlim(1, 1));

                lock (_clients) { _clients.Add(socket); }
                Console.WriteLine("[WS] Client Connected");

                await ReceiveLoop(socket);
            }
            catch (Exception ex)
            {
                Console.WriteLine($"[WS] Accept Error: {ex.Message}");
            }
            finally
            {
                if (wsContext != null)
                {
                    WebSocket socket = wsContext.WebSocket;
                    lock (_clients) { _clients.Remove(socket); }

                    if (_socketLocks.TryRemove(socket, out var semaphore))
                    {
                        semaphore.Dispose();
                    }
                    socket.Dispose();
                }
            }
        }

        private async Task ReceiveLoop(WebSocket socket)
        {
            var buffer = new byte[1024 * 4];
            while (socket.State == WebSocketState.Open)
            {
                try
                {
                    var result = await socket.ReceiveAsync(new ArraySegment<byte>(buffer), CancellationToken.None);
                    if (result.MessageType == WebSocketMessageType.Close)
                    {
                        await socket.CloseAsync(WebSocketCloseStatus.NormalClosure, "Closing", CancellationToken.None);
                    }
                    else if (result.MessageType == WebSocketMessageType.Text)
                    {
                        string msg = Encoding.UTF8.GetString(buffer, 0, result.Count);
                        HandleMessage(msg, socket);
                    }
                }
                catch
                {
                    break;
                }
            }
        }

        private async void HandleMessage(string msg, WebSocket socket)
        {
            // Simple JSON parsing (manual or Newtonsoft)
            // Expected format: { "type": "...", "data": ... }
            try
            {
                dynamic json = Newtonsoft.Json.JsonConvert.DeserializeObject(msg);
                string type = json.type;

                Console.WriteLine($"HandleMessage:{type}");
                if (type == "GET_CONFIG")
                {
                    // Simulate Delay for Loading Screen Test
                    await Task.Delay(1000);

                    // Send Config back
                    var bridge = new MachineBridge(_mainForm); // Re-use logic
                    string configJson = bridge.GetConfig();
                    var response = new { type = "CONFIG_DATA", data = Newtonsoft.Json.JsonConvert.DeserializeObject(configJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "GET_IO_LIST")
                {
                    var bridge = new MachineBridge(_mainForm);
                    string ioJson = bridge.GetIOList();
                    var response = new { type = "IO_LIST_DATA", data = Newtonsoft.Json.JsonConvert.DeserializeObject(ioJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "GET_RECIPE_LIST")
                {
                    var bridge = new MachineBridge(_mainForm);
                    string recipeJson = bridge.GetRecipeList();
                    var response = new { type = "RECIPE_LIST_DATA", data = Newtonsoft.Json.JsonConvert.DeserializeObject(recipeJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "SAVE_CONFIG")
                {
                    string configJson = Newtonsoft.Json.JsonConvert.SerializeObject(json.data);
                    var bridge = new MachineBridge(_mainForm);
                    bridge.SaveConfig(configJson);
                }
                else if (type == "CONTROL")
                {
                    string cmd = json.command;
                    _mainForm.Invoke(new Action(() => _mainForm.HandleCommand(cmd)));
                }
                else if (type == "MOVE")
                {
                    string axis = json.axis;
                    double val = json.value;
                    _mainForm.Invoke(new Action(() => _mainForm.SetTargetPosition(axis, val)));
                }
                else if (type == "SET_IO")
                {
                    int id = json.id;
                    bool state = json.state;
                    _mainForm.Invoke(new Action(() => _mainForm.SetOutput(id, state)));
                }
                else if (type == "SELECT_RECIPE")
                {
                    string recipeId = json.recipeId;
                    var bridge = new MachineBridge(_mainForm);
                    string resultJson = bridge.SelectRecipe(recipeId);
                    var response = new { type = "RECIPE_SELECTED", data = Newtonsoft.Json.JsonConvert.DeserializeObject(resultJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "COPY_RECIPE")
                {
                    string recipeId = json.recipeId;
                    string newName = json.newName;
                    var bridge = new MachineBridge(_mainForm);
                    string resultJson = bridge.CopyRecipe(recipeId, newName);
                    var response = new { type = "RECIPE_COPIED", data = Newtonsoft.Json.JsonConvert.DeserializeObject(resultJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "DELETE_RECIPE")
                {
                    string recipeId = json.recipeId;
                    var bridge = new MachineBridge(_mainForm);
                    string resultJson = bridge.DeleteRecipe(recipeId);
                    var response = new { type = "RECIPE_DELETED", data = Newtonsoft.Json.JsonConvert.DeserializeObject(resultJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "GET_RECIPE")
                {
                    string recipeTitle = json.recipeTitle;
                    var bridge = new MachineBridge(_mainForm);
                    string recipeJson = bridge.GetRecipe(recipeTitle);
                    var response = new { type = "RECIPE_DATA", data = Newtonsoft.Json.JsonConvert.DeserializeObject(recipeJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "SAVE_RECIPE")
                {
                    string recipeTitle = json.recipeTitle;
                    string recipeData = Newtonsoft.Json.JsonConvert.SerializeObject(json.recipeData);
                    var bridge = new MachineBridge(_mainForm);
                    string resultJson = bridge.SaveRecipe(recipeTitle, recipeData);
                    var response = new { type = "RECIPE_SAVED", data = Newtonsoft.Json.JsonConvert.DeserializeObject(resultJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "INITIALIZE_DEVICE")
                {
                    var bridge = new MachineBridge(_mainForm);
                    string resultJson = bridge.InitializeDevice();
                    var response = new { type = "DEVICE_INITIALIZED", data = Newtonsoft.Json.JsonConvert.DeserializeObject(resultJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "GET_INITIALIZE_STATUS")
                {
                    var bridge = new MachineBridge(_mainForm);
                    string statusJson = bridge.GetInitializeStatus();
                    var response = new { type = "INITIALIZE_STATUS_DATA", data = Newtonsoft.Json.JsonConvert.DeserializeObject(statusJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "GET_PROCESSED_DATA")
                {
                    var bridge = new MachineBridge(_mainForm);
                    string dataJson = bridge.GetProcessedData();
                    var response = new { type = "PROCESSED_DATA", data = Newtonsoft.Json.JsonConvert.DeserializeObject(dataJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "CAMERA_CONNECT")
                {
                    var bridge = new MachineBridge(_mainForm);
                    string resultJson = bridge.CameraConnect();
                    var response = new { type = "CAMERA_RESULT", data = Newtonsoft.Json.JsonConvert.DeserializeObject(resultJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "CAMERA_DISCONNECT")
                {
                    var bridge = new MachineBridge(_mainForm);
                    string resultJson = bridge.CameraDisconnect();
                    var response = new { type = "CAMERA_RESULT", data = Newtonsoft.Json.JsonConvert.DeserializeObject(resultJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "CAMERA_GET_IMAGE")
                {
                    var bridge = new MachineBridge(_mainForm);
                    string resultJson = bridge.CameraGetImage();
                    var response = new { type = "CAMERA_RESULT", data = Newtonsoft.Json.JsonConvert.DeserializeObject(resultJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "CAMERA_LIVE_VIEW")
                {
                    var bridge = new MachineBridge(_mainForm);
                    string resultJson = bridge.CameraLiveView();
                    var response = new { type = "CAMERA_RESULT", data = Newtonsoft.Json.JsonConvert.DeserializeObject(resultJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "CAMERA_READ_TEST")
                {
                    var bridge = new MachineBridge(_mainForm);
                    string resultJson = bridge.CameraReadTest();
                    var response = new { type = "CAMERA_RESULT", data = Newtonsoft.Json.JsonConvert.DeserializeObject(resultJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "KEYENCE_TRIGGER_ON")
                {
                    var bridge = new MachineBridge(_mainForm);
                    string resultJson = bridge.KeyenceTriggerOn();
                    var response = new { type = "KEYENCE_RESULT", data = Newtonsoft.Json.JsonConvert.DeserializeObject(resultJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "KEYENCE_TRIGGER_OFF")
                {
                    var bridge = new MachineBridge(_mainForm);
                    string resultJson = bridge.KeyenceTriggerOff();
                    var response = new { type = "KEYENCE_RESULT", data = Newtonsoft.Json.JsonConvert.DeserializeObject(resultJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "KEYENCE_GET_IMAGE")
                {
                    var bridge = new MachineBridge(_mainForm);
                    string resultJson = bridge.KeyenceGetImage();
                    var response = new { type = "KEYENCE_RESULT", data = Newtonsoft.Json.JsonConvert.DeserializeObject(resultJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "KEYENCE_SAVE_IMAGE")
                {
                    var bridge = new MachineBridge(_mainForm);
                    string resultJson = bridge.KeyenceSaveImage();
                    var response = new { type = "KEYENCE_RESULT", data = Newtonsoft.Json.JsonConvert.DeserializeObject(resultJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "TOGGLE_LIGHT")
                {
                    var bridge = new MachineBridge(_mainForm);
                    string resultJson = bridge.ToggleLight();
                    var response = new { type = "LIGHT_RESULT", data = Newtonsoft.Json.JsonConvert.DeserializeObject(resultJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "OPEN_MANUAL_PRINT")
                {
                    var bridge = new MachineBridge(_mainForm);
                    string resultJson = bridge.OpenManualPrint();
                    var response = new { type = "MANUAL_PRINT_RESULT", data = Newtonsoft.Json.JsonConvert.DeserializeObject(resultJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "CANCEL_JOB")
                {
                    var bridge = new MachineBridge(_mainForm);
                    string resultJson = bridge.CancelJob();
                    var response = new { type = "CANCEL_JOB_RESULT", data = Newtonsoft.Json.JsonConvert.DeserializeObject(resultJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "OPEN_MANAGE")
                {
                    var bridge = new MachineBridge(_mainForm);
                    string resultJson = bridge.OpenManage();
                    var response = new { type = "MANAGE_RESULT", data = Newtonsoft.Json.JsonConvert.DeserializeObject(resultJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "OPEN_MANUAL")
                {
                    var bridge = new MachineBridge(_mainForm);
                    string resultJson = bridge.OpenManual();
                    var response = new { type = "MANUAL_RESULT", data = Newtonsoft.Json.JsonConvert.DeserializeObject(resultJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "OPEN_LOG_VIEWER")
                {
                    var bridge = new MachineBridge(_mainForm);
                    string resultJson = bridge.OpenLogViewer();
                    var response = new { type = "LOG_VIEWER_RESULT", data = Newtonsoft.Json.JsonConvert.DeserializeObject(resultJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "OPEN_PROGRAM_FOLDER")
                {
                    var bridge = new MachineBridge(_mainForm);
                    string resultJson = bridge.OpenProgramFolder();
                    var response = new { type = "FOLDER_RESULT", data = Newtonsoft.Json.JsonConvert.DeserializeObject(resultJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "OPEN_LOG_FOLDER")
                {
                    var bridge = new MachineBridge(_mainForm);
                    string resultJson = bridge.OpenLogFolder();
                    var response = new { type = "FOLDER_RESULT", data = Newtonsoft.Json.JsonConvert.DeserializeObject(resultJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "OPEN_SCREENSHOT_FOLDER")
                {
                    var bridge = new MachineBridge(_mainForm);
                    string resultJson = bridge.OpenScreenshotFolder();
                    var response = new { type = "FOLDER_RESULT", data = Newtonsoft.Json.JsonConvert.DeserializeObject(resultJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
                else if (type == "OPEN_SAVED_DATA_FOLDER")
                {
                    var bridge = new MachineBridge(_mainForm);
                    string resultJson = bridge.OpenSavedDataFolder();
                    var response = new { type = "FOLDER_RESULT", data = Newtonsoft.Json.JsonConvert.DeserializeObject(resultJson) };
                    await Send(socket, Newtonsoft.Json.JsonConvert.SerializeObject(response));
                }
            }
            catch (Exception ex)
            {
                Console.WriteLine($"[WS] Msg Error: {ex.Message}");
            }
        }

        private async Task Send(WebSocket socket, string message)
        {
            if (_socketLocks.TryGetValue(socket, out var semaphore))
            {
                await semaphore.WaitAsync();
                try
                {
                    if (socket.State == WebSocketState.Open)
                    {
                        byte[] buffer = Encoding.UTF8.GetBytes(message);
                        await socket.SendAsync(new ArraySegment<byte>(buffer), WebSocketMessageType.Text, true, CancellationToken.None);
                    }
                }
                finally
                {
                    semaphore.Release();
                }
            }
        }

        public async void Broadcast(string message)
        {
            byte[] buffer = Encoding.UTF8.GetBytes(message);
            WebSocket[] clientsCopy;

            lock (_clients)
            {
                clientsCopy = _clients.ToArray();
            }

            foreach (var client in clientsCopy)
            {
                if (client.State == WebSocketState.Open && _socketLocks.TryGetValue(client, out var semaphore))
                {
                    // Fire and forget, but safely
                    _ = Task.Run(async () =>
                    {
                        // Try to get lock immediately. If busy (sending previous frame), skip this frame to prevent lag.
                        if (await semaphore.WaitAsync(0))
                        {
                            try
                            {
                                if (client.State == WebSocketState.Open)
                                {
                                    await client.SendAsync(new ArraySegment<byte>(buffer), WebSocketMessageType.Text, true, CancellationToken.None);
                                }
                            }
                            catch { /* Ignore send errors */ }
                            finally
                            {
                                semaphore.Release();
                            }
                        }
                    });
                }
            }
        }
    }
}