using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Net.WebSockets;
using System.Text;
using System.Threading;
using System.Threading.Tasks;

namespace JtkLibrary {
    /// <summary>
    /// JTK PC방 시스템에 최대한 종속적이지 않게 설계된 웹 소켓 관리용 정적 클래스입니다.<para/>
    /// 등록된 각 웹 소켓은 string형 키를 가지는 Dictionary 콜렉션으로 관리됩니다.
    /// </summary>
    public class WebSocketManager {
        /// <summary>
        /// 웹 소켓이 메시지를 수신할 때 호출되는 이벤트 처리기입니다.
        /// </summary>
        public static event EventHandler<WebSocketManagerEventArgs> OnReceive;

        /// <summary>
        /// 웹 소켓 연결이 닫힌 뒤에 호출되는 이벤트 처리기입니다.
        /// </summary>
        public static event EventHandler<WebSocketManagerEventArgs> OnClosed;

        private static Dictionary<string, WebSocket> webSockets = new Dictionary<string, WebSocket>();
        private const int SEND_BUFFER_LENGTH = 1024;
        private const int RECEIVE_BUFFER_LENGTH = 1024;

        /// <summary>
        /// 등록된 모든 웹 소켓 키의 배열입니다.
        /// </summary>
        public static string[] Keys {
            get {
                string[] result;
                lock (webSockets) {
                    result = webSockets.Keys.ToArray();
                }
                return result;
            }
        }

        public static Task CloseWebSocket(string key) {
            WebSocket ws;
            lock (webSockets) {
                if (webSockets.TryGetValue(key, out ws) == false) {
                    return Task.CompletedTask;
                }
            }
            return ws.CloseAsync(WebSocketCloseStatus.Empty, "", CancellationToken.None);
        }

        private static Task<bool> SendMessageAsync(string key, WebSocketMessageType messageType, ArraySegment<byte> message) =>
            Task.Factory.StartNew(obj => {
                (var _key, var _messageType, var _message) = ((string, WebSocketMessageType, ArraySegment<byte>))obj;
                var destWebSockets = new List<WebSocket>();
                var tasks = new List<Task>();
                var fails = 0;
                lock (webSockets) {
                    if (_key == null) {
                        foreach (var webSocket in webSockets.Values) {
                            destWebSockets.Add(webSocket);
                        }
                    } else if (webSockets.TryGetValue(_key, out var webSocket)) {
                        destWebSockets.Add(webSocket);
                    } else {
                        return false;
                    }
                }
                if (destWebSockets.Count == 0) {
                    return false;
                }
                foreach (var webSocket in destWebSockets) {
                    tasks.Add(Task.Factory.StartNew(_obj => {
                        var _webSocket = _obj as WebSocket;
                        lock (_webSocket) {
                            if (_webSocket.State == WebSocketState.Open) {
                                try {
                                    if (_message.Count <= SEND_BUFFER_LENGTH) {
                                        _webSocket.SendAsync(_message, _messageType, true, CancellationToken.None).Wait();
                                    } else {
                                        int i;
                                        for (i = 0; i + SEND_BUFFER_LENGTH < _message.Count; i += SEND_BUFFER_LENGTH) {
                                            _webSocket.SendAsync(new ArraySegment<byte>(_message.Array, _message.Offset + i, SEND_BUFFER_LENGTH), _messageType, false, CancellationToken.None).Wait();
                                        }
                                        _webSocket.SendAsync(new ArraySegment<byte>(_message.Array, _message.Offset + i, _message.Count % SEND_BUFFER_LENGTH), _messageType, true, CancellationToken.None).Wait();
                                    }
                                } catch {
                                    fails++;
                                }
                            } else {
                                fails++;
                            }
                        }
                    }, webSocket));
                }
                Task.WaitAll(tasks.ToArray());
                return fails == 0;
            }, (key, messageType, message));

        /// <summary>키에 해당하는 웹 소켓에게 비동기 작업으로 문자열 메시지를 보냅니다.</summary>
        public static Task<bool> SendMessageAsync(string key, string textMessage) =>
            SendMessageAsync(key, WebSocketMessageType.Text, new ArraySegment<byte>(Encoding.UTF8.GetBytes(textMessage)));

        /// <summary>키에 해당하는 웹 소켓에게 비동기 작업으로 이진 메시지를 보냅니다.</summary>
        public static Task<bool> SendMessageAsync(string key, ArraySegment<byte> binaryMessage) =>
            SendMessageAsync(key, WebSocketMessageType.Binary, binaryMessage);

        /// <summary>모든 웹 소켓에게 비동기 작업으로 문자열 메시지를 보냅니다.</summary>
        public static Task<bool> BroadcastMessageAsync(string textMessage) =>
            SendMessageAsync(null, WebSocketMessageType.Text, new ArraySegment<byte>(Encoding.UTF8.GetBytes(textMessage)));

        /// <summary>모든 웹 소켓에게 비동기 작업으로 이진 메시지를 보냅니다.</summary>
        public static Task<bool> BroadcastMessageAsync(ArraySegment<byte> binaryMessage) =>
            SendMessageAsync(null, WebSocketMessageType.Binary, binaryMessage);

        /// <summary>그룹키에 해당하는 그룹에 임의의 멤버키를 갖는 멤버 웹 소켓을 등록하고 비동기 작업으로 수신 루프를 반복 수행합니다.</summary>
        /// <param name="groupKey">대상 웹 소켓 그룹</param>
        /// <param name="host">클라이언트 호스트</param>
        /// <param name="webSocket">웹 소켓 인스턴스</param>
        /// <returns>등록 성공 여부</returns>
        public static Task<bool> AddAndLoop(string key, WebSocket webSocket, string firstSendingTextMessage = null) =>
            Task.Factory.StartNew((obj) => {
                MemoryStream bigBuffer = null;
                var buffer = new byte[RECEIVE_BUFFER_LENGTH];
                (var _key, var _webSocket) = ((string, WebSocket))obj;
                var closeStatus = WebSocketCloseStatus.Empty;
                var closeStatusDescription = string.Empty;

                if (_webSocket.State != WebSocketState.Open) {
                    return false;
                }

                // 그룹에 웹소켓 등록
                lock (webSockets) {
                    if (webSockets.ContainsKey(_key)) {
                        return false;
                    }
                    webSockets[_key] = _webSocket;
                }

                // 최초 메시지 보내기
                if (string.IsNullOrEmpty(firstSendingTextMessage) == false) {
                    SendMessageAsync(_key, firstSendingTextMessage).Wait();
                }

                // Receive 루프
                while (_webSocket.State == WebSocketState.Open) {
                    try {
                        // 수신
                        var t = _webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), CancellationToken.None);
                        t.Wait();
                        var receiveResult = t.Result;
                        // Close 체크
                        if (receiveResult.CloseStatus.HasValue) {
                            // 닫을 준비 후, 루프 아웃
                            closeStatus = receiveResult.CloseStatus.Value;
                            closeStatusDescription = receiveResult.CloseStatusDescription;
                            break;
                        }
                        // EoM 여부로 분기
                        if (receiveResult.EndOfMessage == false) {
                            // 수신 미 완료 시, 빅 버퍼에 저장
                            if (bigBuffer == null) {
                                bigBuffer = new MemoryStream();
                            }
                            bigBuffer.Write(buffer, 0, receiveResult.Count);
                        } else {
                            // 수신 완료 시
                            byte[] finalBuffer;
                            int finalBufferLength;
                            // 빅 버퍼 유무로 분기
                            if (bigBuffer != null) {
                                // 빅 버퍼 있을 시
                                bigBuffer.Write(buffer, 0, receiveResult.Count);  // ★
                                finalBuffer = bigBuffer.GetBuffer();
                                finalBufferLength = (int)bigBuffer.Length;
                            } else {
                                // 빅 버퍼 없을 시
                                finalBuffer = buffer;
                                finalBufferLength = receiveResult.Count;
                            }
                            // 메시지 타입으로 분기
                            if (receiveResult.MessageType == WebSocketMessageType.Text) {
                                // Text인 경우
                                var stringMessage = Encoding.UTF8.GetString(finalBuffer, 0, finalBufferLength);
                                // 이벤트 호출
                                OnReceive?.Invoke(_webSocket, new WebSocketManagerEventArgs(key, WebSocketMessageType.Text, stringMessage, null));
                            } else if (receiveResult.MessageType == WebSocketMessageType.Binary) {
                                // Binary인 경우
                                var binaryMessage = new byte[finalBufferLength];
                                Buffer.BlockCopy(finalBuffer, 0, binaryMessage, 0, finalBufferLength);
                                // 이벤트 호출
                                OnReceive?.Invoke(_webSocket, new WebSocketManagerEventArgs(key, WebSocketMessageType.Binary, null, binaryMessage));
                            } else {
                                break;
                            }
                            bigBuffer?.Close();  // ★
                            bigBuffer = null;  // ★
                        }
                    } catch (Exception e) {
                        // 닫을 준비 후, 루프 아웃
                        closeStatus = WebSocketCloseStatus.InternalServerError;
                        closeStatusDescription = e.Message;
                        break;
                    } finally {
                        // 빅 버퍼 닫기
                        //bigBuffer?.Close();  // ★
                        //bigBuffer = null;
                    }
                }

                // 웹소켓 닫기
                try {
                    if (_webSocket.State == WebSocketState.Open) {
                        _webSocket.CloseAsync(closeStatus, closeStatusDescription, CancellationToken.None).Wait();
                    }
                } catch { }
                // 이벤트 호출
                OnClosed?.Invoke(_webSocket, new WebSocketManagerEventArgs(key, WebSocketMessageType.Close, null, null));

                // 그룹에서 웹소켓 제거
                lock (webSockets) {
                    webSockets.Remove(_key);
                }

                return true;
            }, (key, webSocket));

    }

    public class WebSocketManagerEventArgs {
        public string Key { get; }
        public WebSocketMessageType MessageType { get; }
        public string TextMessage { get; }
        public byte[] BinaryMessage { get; }

        public WebSocketManagerEventArgs(string key, WebSocketMessageType messageType, string textMessage, byte[] binaryMessage) {
            Key = key;
            MessageType = messageType;
            TextMessage = textMessage;
            BinaryMessage = binaryMessage;
        }
    }
}
