namespace FinlyticAssets.Util;

using System;
using System.IO;
using System.Net.WebSockets;
using System.Text;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;

public abstract class ManagedWebSocket : IDisposable
{
    private ClientWebSocket? _webSocket;
    private CancellationTokenSource? _cts;
    private Task? _receiveTask;
    private Task? _keepAliveTask;

    /// <summary>
    /// Gets or sets the interval for keep-alive messages. 
    /// If set to null, the periodic timer is not started and no life messages are sent.
    /// </summary>
    public TimeSpan? KeepAliveInterval { get; set; }

    /// <summary>
    /// Gets a value indicating whether the WebSocket is currently open.
    /// </summary>
    public bool IsConnected => _webSocket?.State == WebSocketState.Open;

    /// <summary>
    /// Gets the current state of the WebSocket connection.
    /// </summary>
    public WebSocketState State => _webSocket?.State ?? WebSocketState.None;

    /// <summary>
    /// Connects to the specified WebSocket URI and starts the background tasks for receiving and keep-alive.
    /// </summary>
    /// <param name="uri">The URI of the WebSocket server to connect to.</param>
    /// <param name="keepAliveInterval">Optional. Sets or overrides the keep-alive interval for this connection.</param>
    public async Task ConnectAsync(string uri, TimeSpan? keepAliveInterval = null)
    {
        if (IsConnected) 
            throw new InvalidOperationException("WebSocket is already connected.");

        if (keepAliveInterval.HasValue)
        {
            KeepAliveInterval = keepAliveInterval;
        }

        _webSocket = new ClientWebSocket();
        _cts = new CancellationTokenSource();

        await _webSocket.ConnectAsync(new Uri(uri), _cts.Token);

        _receiveTask = ReceiveLoopAsync(_cts.Token);

        if (KeepAliveInterval.HasValue && KeepAliveInterval.Value > TimeSpan.Zero)
        {
            _keepAliveTask = KeepAliveLoopAsync(_cts.Token);
        }
    }

    /// <summary>
    /// Closes the WebSocket connection gracefully and cancels all running background tasks.
    /// </summary>
    public async Task DisconnectAsync()
    {
        if (_webSocket != null)
        {
            if (_cts != null)
            {
                await _cts.CancelAsync();
            }

            var tasksToWait = new List<Task>();
            if (_receiveTask != null) tasksToWait.Add(_receiveTask);
            if (_keepAliveTask != null) tasksToWait.Add(_keepAliveTask);

            if (tasksToWait.Count > 0)
            {
                try 
                { 
                    await Task.WhenAll(tasksToWait); 
                }
                catch (OperationCanceledException) { /* Ignore */ }
                catch { /* Ignore */ }
            }

            if (_webSocket.State == WebSocketState.Open)
            {
                try
                {
                    await _webSocket.CloseAsync(WebSocketCloseStatus.NormalClosure, "Connection closed by client", CancellationToken.None);
                }
                catch { /* Ignore exceptions during closure */ }
            }

            _webSocket.Dispose();
            _webSocket = null;
            
            _receiveTask = null;
            _keepAliveTask = null;
        }
    }
    /// <summary>
    /// Sends a raw string message over the WebSocket connection.
    /// </summary>
    /// <param name="message">The text message to send.</param>
    public async Task SendAsync(string message)
    {
        if (!IsConnected || _webSocket == null || _cts == null) 
            throw new InvalidOperationException("WebSocket is not connected.");

        var buffer = Encoding.UTF8.GetBytes(message);
        await _webSocket.SendAsync(new ArraySegment<byte>(buffer), WebSocketMessageType.Text, true, _cts.Token);
    }

    /// <summary>
    /// Serializes the given generic object to JSON and sends it over the WebSocket connection.
    /// </summary>
    /// <typeparam name="T">The type of the object to send.</typeparam>
    /// <param name="data">The object to serialize and send.</param>
    public Task SendAsync<T>(T data)
    {
        var json = JsonSerializer.Serialize(data);
        
        return SendAsync(json);
    }

    private async Task ReceiveLoopAsync(CancellationToken cancellationToken)
    {
        var buffer = new byte[8192];

        try
        {
            while (_webSocket?.State == WebSocketState.Open && !cancellationToken.IsCancellationRequested)
            {
                using var ms = new MemoryStream();
                WebSocketReceiveResult result;

                do
                {
                    result = await _webSocket.ReceiveAsync(new ArraySegment<byte>(buffer), cancellationToken);
                    ms.Write(buffer, 0, result.Count);
                }
                while (!result.EndOfMessage);

                if (result.MessageType == WebSocketMessageType.Close)
                {
                    await DisconnectAsync();
                    break;
                }

                if (result.MessageType == WebSocketMessageType.Text)
                {
                    var message = Encoding.UTF8.GetString(ms.ToArray());
                    
                    if (!IsKeepAliveMessage(message))
                    {
                        OnMessageReceived(message); 
                    }
                }
            }
        }
        catch (OperationCanceledException) { /* Expected behavior when cancellation is requested */ }
        catch (Exception ex)
        {
            OnError(ex);
        }
    }

    private async Task KeepAliveLoopAsync(CancellationToken cancellationToken)
    {
        if (!KeepAliveInterval.HasValue) return;

        using var timer = new PeriodicTimer(KeepAliveInterval.Value);
        try
        {
            while (await timer.WaitForNextTickAsync(cancellationToken))
            {
                if (IsConnected)
                {
                    await SendLifeMessageAsync();
                }
            }
        }
        catch (OperationCanceledException) { /* Expected behavior when cancellation is requested */ }
    }

    /// <summary>
    /// Triggered when a new text message is received that is not filtered as a keep-alive message.
    /// </summary>
    /// <param name="message">The received message content.</param>
    protected abstract void OnMessageReceived(string message);

    /// <summary>
    /// Triggered periodically based on the KeepAliveInterval to send a life message or ping to the server.
    /// Override this method in the derived class if you want to send actual keep-alive payloads.
    /// </summary>
    protected virtual Task SendLifeMessageAsync()
    {
        // Default implementation does nothing.
        return Task.CompletedTask;
    }

    /// <summary>
    /// Determines whether the incoming message is a keep-alive response (e.g., a "pong").
    /// If true is returned, the message is filtered out and OnMessageReceived is not triggered.
    /// </summary>
    /// <param name="message">The incoming text message.</param>
    /// <returns>True if the message should be filtered; otherwise, false.</returns>
    protected virtual bool IsKeepAliveMessage(string message)
    {
        // Default implementation does not filter any messages.
        return false; 
    }

    /// <summary>
    /// Optionally overridable method to handle exceptions that occur within the receive loop.
    /// </summary>
    /// <param name="ex">The exception that was caught.</param>
    protected virtual void OnError(Exception ex)
    {
        // Default implementation is empty. Can be overridden in the derived class.
    }

    /// <summary>
    /// Disposes the WebSocket connection and releases all associated resources.
    /// </summary>
    public void Dispose()
    {
        DisconnectAsync().Wait();
        _cts?.Dispose();
    }
}