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;

namespace HMIWeb
{
    public class WebSocketServer
    {
        private HttpListener _httpListener;
        private List<WebSocket> _clients = new List<WebSocket>();
        private MainForm _mainForm;

        public WebSocketServer(string url, MainForm 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;

                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 == "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)));
                }
            }
            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();
                            }
                        }
                    });
                }
            }
        }
    }
}