using System.Buffers.Binary;
using System.Net.WebSockets;
using System.Threading.Channels;
using Lumi.Companion.Protocol;

namespace Lumi.Companion.Core;

public sealed class CompanionSocket : IAsyncDisposable
{
    private const int NetworkAudioPayloadBytes = 640; // 20 ms of 16 kHz mono PCM16.
    private readonly Channel<ReadOnlyMemory<byte>> _audio = Channel.CreateBounded<ReadOnlyMemory<byte>>(new BoundedChannelOptions(250) { FullMode = BoundedChannelFullMode.Wait, SingleReader = true });
    private readonly object _bridgeAudioGate = new();
    private readonly byte[] _bridgeAudioBuffer = new byte[ProtocolV1.MaxAudioPayloadBytes + NetworkAudioPayloadBytes];
    private readonly byte[] _bridgeAudioHeader = new byte[ProtocolV1.AudioHeaderBytes];
    private int _bridgeAudioBuffered;
    private uint _bridgeAudioSequence;
    private ClientWebSocket? _socket;
    private CancellationTokenSource? _lifetime;
    private readonly SemaphoreSlim _sendLock = new(1, 1);
    private readonly List<Task> _backgroundTasks = [];
    private long _droppedFrames;
    private int _disconnectNotified;
    public long DroppedFrames => Interlocked.Read(ref _droppedFrames);
    public Guid SessionId { get; private set; }
    public event Func<ServerEnvelope, Task>? MessageReceived;
    public event Action<Exception?>? Disconnected;

    public async Task ConnectAsync(DeviceCredential credential, string companionVersion, string pluginVersion, string? obsVersion, CancellationToken cancellationToken)
    {
        _disconnectNotified = 0;
        _socket = new ClientWebSocket();
        _socket.Options.SetRequestHeader("Authorization", $"LumiDevice {credential.DeviceId}.{credential.DeviceSecret}");
        var host = new Uri(credential.Host);
        if (host.Scheme != "https" && !(host.Scheme == "http" && host.IsLoopback))
            throw new InvalidDataException("Companion transport requires HTTPS unless this credential was issued for localhost.");
        var uri = new UriBuilder(host) { Scheme = host.Scheme == "https" ? "wss" : "ws", Path = "/plugins/lumi_transcription/live" }.Uri;
        await _socket.ConnectAsync(uri, cancellationToken);
        var hello = ProtocolV1.EncodeEnvelope("hello", new { companion_version = companionVersion, plugin_version = pluginVersion, obs_version = obsVersion, capabilities = credential.Capabilities, audio = new { codec = "pcm_s16le", sample_rate = 16000, channels = 1, bits = 16 } });
        await _socket.SendAsync(hello, WebSocketMessageType.Text, true, cancellationToken);
        var acknowledgement = await ReceiveMessageAsync(_socket, cancellationToken);
        if (acknowledgement.Type != "hello_ack" || acknowledgement.Version != 1 || acknowledgement.SessionId is null)
            throw new InvalidDataException("Lumi returned an incompatible companion handshake.");
        SessionId = acknowledgement.SessionId.Value;
        _lifetime = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
        _backgroundTasks.Add(Task.Run(() => RunGuardedAsync(SendAudioAsync, _lifetime.Token), _lifetime.Token));
        _backgroundTasks.Add(Task.Run(() => RunGuardedAsync(ReceiveAsync, _lifetime.Token), _lifetime.Token));
        _backgroundTasks.Add(Task.Run(() => RunGuardedAsync(SendHeartbeatsAsync, _lifetime.Token), _lifetime.Token));
    }

    public bool QueueAudio(AudioFrame frame)
    {
        return QueueEncodedAudio(ProtocolV1.EncodeAudio(frame));
    }
    public bool QueueEncodedAudio(ReadOnlyMemory<byte> encoded)
    {
        if (encoded.Length < ProtocolV1.AudioHeaderBytes || encoded.Length > ProtocolV1.AudioHeaderBytes + ProtocolV1.MaxAudioPayloadBytes ||
            !encoded.Span[..4].SequenceEqual("LACP"u8)) return false;
        if (_audio.Writer.TryWrite(encoded)) return true;
        _audio.Reader.TryRead(out _);
        Interlocked.Increment(ref _droppedFrames);
        return _audio.Writer.TryWrite(encoded);
    }
    public bool QueueBridgeAudio(ReadOnlyMemory<byte> encoded)
    {
        var input = encoded.Span;
        if (input.Length < ProtocolV1.AudioHeaderBytes || input.Length > ProtocolV1.AudioHeaderBytes + ProtocolV1.MaxAudioPayloadBytes ||
            !input[..4].SequenceEqual("LACP"u8) || BinaryPrimitives.ReadUInt16LittleEndian(input[6..8]) != ProtocolV1.AudioHeaderBytes) return false;
        var payloadBytes = (int)BinaryPrimitives.ReadUInt32LittleEndian(input[60..64]);
        if (payloadBytes < 0 || payloadBytes % 2 != 0 || input.Length != ProtocolV1.AudioHeaderBytes + payloadBytes) return false;
        lock (_bridgeAudioGate)
        {
            var sourceChanged = !_bridgeAudioHeader.AsSpan(36, 16).SequenceEqual(input[36..52]);
            var stateChanged = _bridgeAudioHeader[5] != input[5];
            if (_bridgeAudioBuffered > 0 && (sourceChanged || stateChanged)) _bridgeAudioBuffered = 0;
            input[..ProtocolV1.AudioHeaderBytes].CopyTo(_bridgeAudioHeader);
            input[ProtocolV1.AudioHeaderBytes..].CopyTo(_bridgeAudioBuffer.AsSpan(_bridgeAudioBuffered));
            _bridgeAudioBuffered += payloadBytes;
            var accepted = true;
            while (_bridgeAudioBuffered >= NetworkAudioPayloadBytes)
            {
                var normalized = new byte[ProtocolV1.AudioHeaderBytes + NetworkAudioPayloadBytes];
                _bridgeAudioHeader.CopyTo(normalized, 0);
                BinaryPrimitives.WriteUInt32LittleEndian(normalized.AsSpan(8, 4), _bridgeAudioSequence++);
                BinaryPrimitives.WriteUInt64LittleEndian(normalized.AsSpan(12, 8), (ulong)DateTimeOffset.UtcNow.ToUnixTimeMilliseconds() * 1000UL);
                SessionId.TryWriteBytes(normalized.AsSpan(20, 16), bigEndian: true, out _);
                BinaryPrimitives.WriteUInt32LittleEndian(normalized.AsSpan(60, 4), NetworkAudioPayloadBytes);
                _bridgeAudioBuffer.AsSpan(0, NetworkAudioPayloadBytes).CopyTo(normalized.AsSpan(ProtocolV1.AudioHeaderBytes));
                _bridgeAudioBuffered -= NetworkAudioPayloadBytes;
                if (_bridgeAudioBuffered > 0)
                    _bridgeAudioBuffer.AsSpan(NetworkAudioPayloadBytes, _bridgeAudioBuffered).CopyTo(_bridgeAudioBuffer);
                accepted &= QueueEncodedAudio(normalized);
            }
            return accepted;
        }
    }
    public async Task SendAsync(string type, object payload, Guid? sessionId, CancellationToken cancellationToken)
    {
        if (_socket?.State != WebSocketState.Open) throw new InvalidOperationException("Lumi is not connected.");
        await SendLockedAsync(ProtocolV1.EncodeEnvelope(type, payload, sessionId), WebSocketMessageType.Text, cancellationToken);
    }
    private async Task SendAudioAsync(CancellationToken cancellationToken)
    {
        await foreach (var frame in _audio.Reader.ReadAllAsync(cancellationToken))
        {
            if (_socket?.State != WebSocketState.Open) continue;
            await SendLockedAsync(frame, WebSocketMessageType.Binary, cancellationToken);
        }
    }
    private async Task SendLockedAsync(ReadOnlyMemory<byte> body, WebSocketMessageType type, CancellationToken cancellationToken)
    {
        await _sendLock.WaitAsync(cancellationToken);
        try { if (_socket?.State == WebSocketState.Open) await _socket.SendAsync(body, type, true, cancellationToken); }
        finally { _sendLock.Release(); }
    }
    private async Task ReceiveAsync(CancellationToken cancellationToken)
    {
        while (_socket?.State == WebSocketState.Open && !cancellationToken.IsCancellationRequested)
        {
            var message = await ReceiveMessageAsync(_socket, cancellationToken);
            if (MessageReceived is { } handler) await handler(message);
        }
    }
    private async Task SendHeartbeatsAsync(CancellationToken cancellationToken)
    {
        using var timer = new PeriodicTimer(TimeSpan.FromSeconds(8));
        while (await timer.WaitForNextTickAsync(cancellationToken))
            await SendAsync("ping", new { }, SessionId, cancellationToken);
    }
    private async Task RunGuardedAsync(Func<CancellationToken, Task> operation, CancellationToken cancellationToken)
    {
        try { await operation(cancellationToken); }
        catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) { }
        catch (Exception error)
        {
            _lifetime?.Cancel();
            if (Interlocked.Exchange(ref _disconnectNotified, 1) == 0) Disconnected?.Invoke(error);
        }
    }
    private static async Task<ServerEnvelope> ReceiveMessageAsync(ClientWebSocket socket, CancellationToken cancellationToken)
    {
        using var body = new MemoryStream();
        var buffer = new byte[8192];
        WebSocketReceiveResult result;
        do
        {
            result = await socket.ReceiveAsync(new ArraySegment<byte>(buffer), cancellationToken);
            if (result.MessageType == WebSocketMessageType.Close)
            {
                if (socket.State == WebSocketState.CloseReceived)
                    try { await socket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, "close_ack", CancellationToken.None); } catch (WebSocketException) { }
                throw new EndOfStreamException($"Lumi closed the companion connection ({result.CloseStatus?.ToString() ?? "no status"}: {result.CloseStatusDescription ?? "no reason"}).");
            }
            body.Write(buffer, 0, result.Count);
            if (body.Length > ProtocolV1.MaxJsonBytes) throw new InvalidDataException("Lumi message exceeds the protocol limit.");
        } while (!result.EndOfMessage);
        if (result.MessageType != WebSocketMessageType.Text) throw new InvalidDataException("Unexpected binary message from Lumi.");
        return System.Text.Json.JsonSerializer.Deserialize<ServerEnvelope>(body.ToArray(), ProtocolV1.JsonOptions)
            ?? throw new InvalidDataException("Lumi message is invalid.");
    }
    public async ValueTask DisposeAsync()
    {
        _lifetime?.Cancel();
        if (_socket?.State == WebSocketState.Open)
        {
            try { await _socket.CloseAsync(WebSocketCloseStatus.NormalClosure, "companion_exit", CancellationToken.None); }
            catch (WebSocketException) { }
        }
        try { await Task.WhenAll(_backgroundTasks); } catch (OperationCanceledException) { }
        _socket?.Dispose(); _lifetime?.Dispose(); _sendLock.Dispose();
    }
}
