﻿using System;
using System.Net.WebSockets;
using System.Threading;
using System.Threading.Tasks;

namespace SchrodingersForward.Models
{
    public class WebSocketWrapper : IDisposable
    {
        private WebSocket webSocket;
        private readonly ChatRoom room;

        public WebSocket WebSocket => webSocket;

        private Task receiveTask;
        private CancellationTokenSource cancellationTokenSource;
        private byte[] buffer;

        public WebSocketWrapper(WebSocket webSocket, ChatRoom room, int bufferSize)
        {
            this.webSocket = webSocket;
            this.room = room;
            buffer = new byte[bufferSize];
            cancellationTokenSource = new CancellationTokenSource();
            receiveTask = new Task(ReceiveMessageLoop, cancellationTokenSource.Token, cancellationTokenSource.Token);
            room.AddClient(this);
        }

        private void ReceiveMessageLoop(object state)
        {
            CancellationToken cancellationToken = (CancellationToken)state;
            WebSocketReceiveResult receiveResult = null;
            while (true)
            {
                if (cancellationToken.IsCancellationRequested)
                    break;
                receiveResult = webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), cancellationToken).GetAwaiter().GetResult();
                if (receiveResult == null)
                    break;
                if (receiveResult.CloseStatus.HasValue)
                    break;
                room.ForwardAsync(this, receiveResult, buffer, cancellationToken).Wait();
            }
            if (cancellationToken.IsCancellationRequested || receiveResult == null)
            {
                webSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, "a cancellation request ends the connection", CancellationToken.None).Wait();
            }
            else
            {
                webSocket.CloseAsync(receiveResult.CloseStatus.GetValueOrDefault(), receiveResult.CloseStatusDescription, CancellationToken.None).Wait();
            }
            room.RemoveClient(this);
            this.Dispose();
        }

        public Task StartTask()
        {
            receiveTask.Start();
            return receiveTask;
        }

        #region IDisposable Support
        private bool disposedValue = false;

        protected virtual void Dispose(bool disposing)
        {
            if (!disposedValue)
            {
                if (disposing)
                {
                }

                if (receiveTask.Status == TaskStatus.Running)
                {
                    cancellationTokenSource.Cancel();
                    receiveTask.Wait();
                }
                receiveTask.Dispose();
                receiveTask = null;
                cancellationTokenSource.Dispose();
                cancellationTokenSource = null;
                buffer = null;
                disposedValue = true;
            }
        }

        ~WebSocketWrapper()
        {
            Dispose(false);
        }

        public void Dispose() => Dispose(true);
        #endregion
    }
}
