using System.Buffers;
using System.Collections.Concurrent;
using System.Net.WebSockets;
using System.Security.Cryptography;
using System.Threading.Channels;
using System.Text.Json;
using Fleet.Core;
using Fleet.Identity;
using Fleet.Rpc;

namespace Fleet.Server;

/// <summary>
/// 一个节点的 WebSocket 会话：挑战 → 签名认证 → hello → 状态推送 / 收件箱 / 产品 RPC。
/// 收到的消息按顺序处理，所以产品的处理函数不能等待“向同一个节点发起的请求”，否则会死锁。
/// </summary>
internal sealed class NodeConnection : IDisposable
{
    private static readonly TimeSpan s_authTimeout = TimeSpan.FromSeconds(15);
    private static readonly TimeSpan s_presenceRefresh = TimeSpan.FromSeconds(60);
    private static readonly TimeSpan s_statePushRetry = TimeSpan.FromSeconds(5);
    private static readonly TimeSpan s_closeTimeout = TimeSpan.FromSeconds(2);

    private readonly FleetLinkHost _host;
    private readonly WebSocket _socket;
    private readonly SemaphoreSlim _sendLock = new(1, 1);
    private readonly ConcurrentDictionary<long, TaskCompletionSource<RpcMessage>> _pending = new();
    private readonly Channel<bool> _stateSignal = Channel.CreateBounded<bool>(new BoundedChannelOptions(1) { FullMode = BoundedChannelFullMode.DropWrite });
    private readonly Dictionary<string, long> _lastPushedVersion = new(StringComparer.Ordinal);
    private readonly CancellationTokenSource _lifetime = new();
    private long _nextRequestId;
    private long _lastPresenceTicks;
    private NodeRecord? _node;
    private NodeSessionInfo? _session;
    private string? _closeReason;
    private bool _helloDone;

    public NodeConnection(FleetLinkHost host, WebSocket socket)
    {
        _host = host;
        _socket = socket;
    }

    public NodeId NodeId => _node!.Id;

    public NodeSessionInfo Session => _session!;

    public async Task RunAsync(CancellationToken ct)
    {
        using CancellationTokenSource linked = CancellationTokenSource.CreateLinkedTokenSource(ct, _lifetime.Token);
        CancellationToken token = linked.Token;
        bool registered = false;
        Task? pusher = null;
        try
        {
            if (!await AuthenticateAsync(token))
            {
                return;
            }

            NodeConnection? old = _host.Register(this);
            registered = true;
            if (old is not null)
            {
                await old.CloseAsync("replaced");
            }

            string status = StatusName(_node!.Status);
            FleetLog.Connected(_host.Logger, _node.Id.Value, status);
            await UpdatePresenceAsync(online: true, token);
            await FireAsync(NodeEventKind.Online, token);

            pusher = Task.Run(() => StatePushLoopAsync(token), CancellationToken.None);
            await ReceiveLoopAsync(token);
        }
        catch (OperationCanceledException)
        {
            // 关闭原因由触发取消的一方设置
        }
        catch (WebSocketException)
        {
            _closeReason = "connection lost";
        }
        finally
        {
            await _lifetime.CancelAsync();
            if (pusher is not null)
            {
                try
                {
                    await pusher;
                }
                catch (OperationCanceledException)
                {
                    // 正常结束
                }
            }

            foreach (TaskCompletionSource<RpcMessage> waiter in _pending.Values)
            {
                waiter.TrySetCanceled(CancellationToken.None);
            }

            if (registered)
            {
                _host.Unregister(this);
                FleetLog.Disconnected(_host.Logger, _node!.Id.Value, _closeReason ?? "closed");

                // 被新连接替换时，在线状态属于新连接。
                if (!_host.IsOnline(_node.Id))
                {
                    await UpdatePresenceAsync(online: false, CancellationToken.None);
                    await FireAsync(NodeEventKind.Offline, CancellationToken.None);
                }
            }

            if (_socket.State is WebSocketState.Open or WebSocketState.CloseReceived)
            {
                // 完整的关闭握手：发送 Close 并等待对端确认（有上限），这样对端能收到关闭原因
                try
                {
                    using var closeTimeout = new CancellationTokenSource(s_closeTimeout, _host.Time);
                    await _socket.CloseAsync(WebSocketCloseStatus.NormalClosure, _closeReason ?? "closed", closeTimeout.Token);
                }
                catch (WebSocketException)
                {
                    // 对端已经断开
                }
                catch (OperationCanceledException)
                {
                    // 对端没有在时限内确认：不再等待
                }
            }

        }
    }

    public void Dispose()
    {
        _lifetime.Dispose();
        _sendLock.Dispose();
    }

    public void SignalStateChanged() => _stateSignal.Writer.TryWrite(true);

    /// <summary>
    /// 服务端发起的关闭：发送 Close 帧，等对端确认之后接收循环自然结束。
    /// 注意不能直接取消接收的令牌——取消 <c>ReceiveAsync</c> 会中止整个 WebSocket，对端拿不到关闭原因；取消只作为超时兜底。
    /// </summary>
    public async ValueTask CloseAsync(string reason)
    {
        _closeReason ??= reason;
        try
        {
            await _sendLock.WaitAsync(CancellationToken.None);
        }
        catch (ObjectDisposedException)
        {
            return;
        }

        try
        {
            if (_socket.State == WebSocketState.Open)
            {
                await _socket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, reason, CancellationToken.None);
            }
        }
        catch (WebSocketException)
        {
            // 对端已经断开
        }
        finally
        {
            _sendLock.Release();
        }

        _lifetime.CancelAfter(s_closeTimeout);
    }

    public async ValueTask RevokeAsync(string reason, CancellationToken ct)
    {
        await SendAsync(RpcMessage.Notification(FleetMethods.Revoked, new RevokedParams { Reason = reason }, FleetRpcContext.Default.RevokedParams), ct);
        await CloseAsync(reason);
    }

    public async ValueTask<bool> SendAsync(RpcMessage message, CancellationToken ct)
    {
        byte[] bytes = JsonSerializer.SerializeToUtf8Bytes(message, FleetRpcContext.Default.RpcMessage);
        try
        {
            await _sendLock.WaitAsync(ct);
        }
        catch (ObjectDisposedException)
        {
            return false;
        }

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

            await _socket.SendAsync(bytes, WebSocketMessageType.Text, endOfMessage: true, ct);
            return true;
        }
        catch (WebSocketException)
        {
            return false;
        }
        finally
        {
            _sendLock.Release();
        }
    }

    public async ValueTask<Result<JsonElement>> RequestAsync(string method, JsonElement parameters, CancellationToken ct)
    {
        long id = Interlocked.Increment(ref _nextRequestId);
        var waiter = new TaskCompletionSource<RpcMessage>(TaskCreationOptions.RunContinuationsAsynchronously);
        _pending[id] = waiter;
        try
        {
            var request = new RpcMessage { Id = id, Method = method, Params = parameters };
            if (!await SendAsync(request, ct))
            {
                return Result<JsonElement>.Fail(FleetErrors.NodeOffline, "node is not connected");
            }

            using var timeout = new CancellationTokenSource(_host.Options.RequestTimeout, _host.Time);
            using CancellationTokenSource linked = CancellationTokenSource.CreateLinkedTokenSource(ct, timeout.Token);
            RpcMessage response = await waiter.Task.WaitAsync(linked.Token);
            return response.Error is not null
                ? Result<JsonElement>.Fail(response.Error.ToFleetError())
                : Result<JsonElement>.Ok(response.Result ?? default);
        }
        catch (OperationCanceledException) when (!ct.IsCancellationRequested)
        {
            return Result<JsonElement>.Fail(FleetErrors.Timeout, $"{method} timed out");
        }
        finally
        {
            _pending.TryRemove(id, out _);
        }
    }

    // ───────────────────────── 认证 ─────────────────────────

    private async Task<bool> AuthenticateAsync(CancellationToken ct)
    {
        string nonce = Convert.ToBase64String(RandomNumberGenerator.GetBytes(32));
        long expiresAtMs = _host.Time.GetUtcNow().Add(_host.Options.ChallengeTtl).ToUnixTimeMilliseconds();
        await SendAsync(
            RpcMessage.Notification(FleetMethods.Challenge, new ChallengeParams { Nonce = nonce, ExpiresAtMs = expiresAtMs }, FleetRpcContext.Default.ChallengeParams),
            ct);

        using var timeout = new CancellationTokenSource(s_authTimeout, _host.Time);
        using CancellationTokenSource linked = CancellationTokenSource.CreateLinkedTokenSource(ct, timeout.Token);
        Frame? frame;
        try
        {
            frame = await ReceiveFrameAsync(linked.Token);
        }
        catch (OperationCanceledException) when (!ct.IsCancellationRequested)
        {
            await RejectAsync(null, "auth timeout", ct);
            return false;
        }

        if (frame is not { IsText: true } || !TryParse(frame.Data, out RpcMessage? message) || message.Method != FleetMethods.Auth || message.Id is null)
        {
            await RejectAsync(null, "expected fleet.auth", ct);
            return false;
        }

        AuthParams? auth = null;
        try
        {
            auth = message.Params?.Deserialize(FleetRpcContext.Default.AuthParams);
        }
        catch (JsonException)
        {
            // 按缺失处理
        }

        NodeRecord? node = auth is null ? null : await _host.Store.GetNodeAsync(new NodeId(auth.NodeId), ct);
        bool signatureOk = node is not null
            && _host.Time.GetUtcNow().ToUnixTimeMilliseconds() <= expiresAtMs
            && TryDecode(auth!.Signature, out byte[] signature)
            && NodeIdentity.Verify(node.PublicKey, signature, SignaturePurpose.Auth, auth.NodeId, nonce);
        if (!signatureOk)
        {
            // 节点不存在与签名错误返回同样的错误，不泄露节点是否存在。
            await RejectAsync(message.Id, "authentication failed", ct);
            return false;
        }

        _node = node;
        if (node!.Status == NodeStatus.Disabled)
        {
            await SendAsync(RpcMessage.Failure(message.Id, new FleetError(FleetErrors.NodeDisabled, "node is disabled")), ct);
            await RevokeAsync("disabled", ct);
            return false;
        }

        _session = new NodeSessionInfo(node.Id, _host.Time.GetUtcNow().ToUnixTimeMilliseconds(), node.NodeVersion, [], null);
        await SendAsync(
            RpcMessage.Success(message.Id.Value, new AuthResult { Status = StatusName(node.Status) }, FleetRpcContext.Default.AuthResult),
            ct);
        return true;
    }

    private async Task RejectAsync(long? id, string reason, CancellationToken ct)
    {
        FleetLog.AuthRejected(_host.Logger, reason);
        _closeReason ??= reason;
        if (id is not null)
        {
            await SendAsync(RpcMessage.Failure(id, new FleetError(FleetErrors.Unauthorized, "authentication failed")), ct);
        }

        if (_socket.State == WebSocketState.Open)
        {
            try
            {
                await _socket.CloseOutputAsync(WebSocketCloseStatus.PolicyViolation, "unauthorized", ct);
            }
            catch (WebSocketException)
            {
                // 对端已经断开
            }
        }
    }

    // ───────────────────────── 接收循环 ─────────────────────────

    private async Task ReceiveLoopAsync(CancellationToken ct)
    {
        while (!ct.IsCancellationRequested)
        {
            using var idle = new CancellationTokenSource(_host.Options.HeartbeatTimeout, _host.Time);
            using CancellationTokenSource linked = CancellationTokenSource.CreateLinkedTokenSource(ct, idle.Token);
            Frame? frame;
            try
            {
                frame = await ReceiveFrameAsync(linked.Token);
            }
            catch (OperationCanceledException) when (idle.IsCancellationRequested && !ct.IsCancellationRequested)
            {
                _closeReason ??= "heartbeat timeout";
                return;
            }

            if (frame is null)
            {
                _closeReason ??= "closed by peer";
                return;
            }

            if (frame.IsText)
            {
                if (TryParse(frame.Data, out RpcMessage? message))
                {
                    await HandleMessageAsync(message, ct);
                }
            }
            else
            {
                await HandleBinaryAsync(frame.Data, ct);
            }
        }
    }

    private async Task HandleMessageAsync(RpcMessage message, CancellationToken ct)
    {
        if (message.IsResponse)
        {
            if (_pending.TryGetValue(message.Id!.Value, out TaskCompletionSource<RpcMessage>? waiter))
            {
                waiter.TrySetResult(message);
            }

            return;
        }

        switch (message.Method)
        {
            case FleetMethods.Hello:
                await HandleHelloAsync(message, ct);
                return;
            case FleetMethods.Heartbeat:
                await HandleHeartbeatAsync(message, ct);
                return;
        }

        if (_node!.Status != NodeStatus.Active)
        {
            if (message.Id is not null)
            {
                await SendAsync(RpcMessage.Failure(message.Id, new FleetError(FleetErrors.NodePending, "node is not approved yet")), ct);
            }

            return;
        }

        var context = new NodeCallContext(_node, _session!, _host);
        RpcMessage? response;
        try
        {
            response = await _host.Router.DispatchAsync(context, message, ct);
        }
        catch (OperationCanceledException) when (ct.IsCancellationRequested)
        {
            throw;
        }

        if (response is not null)
        {
            await SendAsync(response, ct);
        }
    }

    private async Task HandleHelloAsync(RpcMessage message, CancellationToken ct)
    {
        HelloParams? hello = null;
        try
        {
            hello = message.Params?.Deserialize(FleetRpcContext.Default.HelloParams);
        }
        catch (JsonException)
        {
            // 按缺失处理
        }

        if (hello is null)
        {
            if (message.Id is not null)
            {
                await SendAsync(RpcMessage.Failure(message.Id, new FleetError(FleetErrors.InvalidParams, "invalid hello")), ct);
            }

            return;
        }

        long nowMs = _host.Time.GetUtcNow().ToUnixTimeMilliseconds();
        _session = _session! with
        {
            NodeVersion = hello.NodeVersion,
            Capabilities = hello.Capabilities,
            HelloMetaJson = hello.Meta?.GetRawText(),
        };
        _node = _node! with { NodeVersion = hello.NodeVersion };
        _helloDone = true;

        await UpdatePresenceAsync(online: true, ct);
        await FireAsync(NodeEventKind.Hello, ct, hello);

        if (message.Id is not null)
        {
            await SendAsync(
                RpcMessage.Success(
                    message.Id.Value,
                    new HelloResult { ServerTimeMs = nowMs, ClockOffsetMs = hello.NowUtcMs - nowMs },
                    FleetRpcContext.Default.HelloResult),
                ct);
        }

        SignalStateChanged();
    }

    private async Task HandleHeartbeatAsync(RpcMessage message, CancellationToken ct)
    {
        JsonElement? extra = null;
        try
        {
            extra = message.Params?.Deserialize(FleetRpcContext.Default.HeartbeatParams)?.Extra;
        }
        catch (JsonException)
        {
            // 忽略格式不正确的附加字段
        }

        long now = _host.Time.GetTimestamp();
        if (_host.Time.GetElapsedTime(Interlocked.Read(ref _lastPresenceTicks), now) >= s_presenceRefresh)
        {
            await UpdatePresenceAsync(online: true, ct);
        }

        await FireAsync(NodeEventKind.Heartbeat, ct, hello: null, extra);
        if (message.Id is not null)
        {
            await SendAsync(new RpcMessage { Id = message.Id, Result = JsonSerializer.SerializeToElement(new EmptyObject(), FleetEmptyObject.Info) }, ct);
        }
    }

    private async Task HandleBinaryAsync(byte[] data, CancellationToken ct)
    {
        if (_node!.Status != NodeStatus.Active
            || !OutboxFrame.TryParse(data, out byte channel, out Guid msgId, out ReadOnlyMemory<byte> payload)
            || !_host.Inboxes.TryGetValue(channel, out IInbox? inbox))
        {
            return;
        }

        InboxOutcome outcome;
        try
        {
            outcome = await inbox.HandleAsync(_node, msgId, payload, ct);
        }
        catch (OperationCanceledException) when (ct.IsCancellationRequested)
        {
            throw;
        }
        catch (Exception ex)
        {
            // 没有持久化就不确认，节点会重发。
            FleetLog.InboxFailed(_host.Logger, ex, _node.Id.Value, channel);
            return;
        }

        if (outcome == InboxOutcome.Rejected)
        {
            return;
        }

        var ack = new OutboxAckParams();
        string id = msgId.ToString();
        if (outcome == InboxOutcome.Duplicate)
        {
            ack.Duplicates.Add(id);
        }
        else
        {
            ack.Ids.Add(id);
        }

        await SendAsync(RpcMessage.Notification(FleetMethods.OutboxAck, ack, FleetRpcContext.Default.OutboxAckParams), ct);
    }

    // ───────────────────────── 状态推送 ─────────────────────────

    private async Task StatePushLoopAsync(CancellationToken ct)
    {
        while (await _stateSignal.Reader.WaitToReadAsync(ct))
        {
            _stateSignal.Reader.TryRead(out _);
            bool complete = await PushStateAsync(ct);
            if (!complete)
            {
                await Task.Delay(s_statePushRetry, _host.Time, ct);
                SignalStateChanged();
            }
        }
    }

    private async Task<bool> PushStateAsync(CancellationToken ct)
    {
        NodeRecord? fresh = await _host.Store.GetNodeAsync(_node!.Id, ct);
        if (fresh is null)
        {
            return true;
        }

        _node = fresh with { NodeVersion = _node.NodeVersion ?? fresh.NodeVersion };
        if (fresh.Status == NodeStatus.Disabled)
        {
            await RevokeAsync("disabled", ct);
            return true;
        }

        if (fresh.Status != NodeStatus.Active || !_helloDone)
        {
            // pending 不下发状态；还没有 hello 的连接等 hello 之后再推送。
            return true;
        }

        bool complete = true;
        foreach (IStateProvider provider in _host.StateProviders)
        {
            StateDocument? document = await provider.GetAsync(_node, ct);
            if (document is null || (_lastPushedVersion.TryGetValue(provider.Type, out long last) && last == document.Version))
            {
                continue;
            }

            Result<JsonElement> raw = await RequestAsync(
                FleetMethods.StateSync,
                JsonSerializer.SerializeToElement(
                    new StateSyncParams { Type = provider.Type, Version = document.Version, Doc = document.Doc },
                    FleetRpcContext.Default.StateSyncParams),
                ct);
            if (!raw.TryGet(out JsonElement value, out FleetError? error))
            {
                FleetLog.StatePushFailed(_host.Logger, _node.Id.Value, provider.Type, error.Code.Value);
                complete = false;
                continue;
            }

            StateSyncResult? result = null;
            try
            {
                result = value.Deserialize(FleetRpcContext.Default.StateSyncResult);
            }
            catch (JsonException)
            {
                // 按缺失处理
            }

            if (result is null)
            {
                complete = false;
                continue;
            }

            _lastPushedVersion[provider.Type] = document.Version;
            await provider.OnAppliedAsync(_node, document.Version, result, ct);
        }

        return complete;
    }

    // ───────────────────────── 帮助方法 ─────────────────────────

    private async Task UpdatePresenceAsync(bool online, CancellationToken ct)
    {
        Interlocked.Exchange(ref _lastPresenceTicks, _host.Time.GetTimestamp());
        await _host.Store.UpdatePresenceAsync(
            _node!.Id,
            new Presence(online, _host.Time.GetUtcNow().ToUnixTimeMilliseconds(), _session?.NodeVersion, _session?.Capabilities ?? [], _session?.HelloMetaJson),
            ct);
    }

    private async Task FireAsync(NodeEventKind kind, CancellationToken ct, HelloParams? hello = null, JsonElement? extra = null)
    {
        var nodeEvent = new NodeEvent(kind, _node!, _host.Time.GetUtcNow().ToUnixTimeMilliseconds(), hello, extra);
        foreach (INodeEventSink sink in _host.Sinks)
        {
            try
            {
                await sink.OnEventAsync(nodeEvent, ct);
            }
            catch (OperationCanceledException) when (ct.IsCancellationRequested)
            {
                throw;
            }
            catch (Exception ex)
            {
                // 产品的事件处理失败不能影响连接。
                FleetLog.SinkFailed(_host.Logger, ex, kind.ToString(), _node!.Id.Value);
            }
        }
    }

    private async ValueTask<Frame?> ReceiveFrameAsync(CancellationToken ct)
    {
        var writer = new ArrayBufferWriter<byte>(1024);
        while (true)
        {
            Memory<byte> buffer = writer.GetMemory(4096);
            ValueWebSocketReceiveResult result = await _socket.ReceiveAsync(buffer, ct);
            if (result.MessageType == WebSocketMessageType.Close)
            {
                return null;
            }

            writer.Advance(result.Count);
            if (writer.WrittenCount > _host.Options.MaxFrameBytes)
            {
                _closeReason ??= "frame too large";
                await _socket.CloseOutputAsync(WebSocketCloseStatus.MessageTooBig, "frame too large", CancellationToken.None);
                return null;
            }

            if (result.EndOfMessage)
            {
                return new Frame(result.MessageType == WebSocketMessageType.Text, writer.WrittenSpan.ToArray());
            }
        }
    }

    private static bool TryParse(byte[] data, [System.Diagnostics.CodeAnalysis.NotNullWhen(true)] out RpcMessage? message)
    {
        try
        {
            message = JsonSerializer.Deserialize(data, FleetRpcContext.Default.RpcMessage);
            return message is not null;
        }
        catch (JsonException)
        {
            message = null;
            return false;
        }
    }

    private static bool TryDecode(string text, out byte[] data)
    {
        try
        {
            data = Convert.FromBase64String(text);
            return true;
        }
        catch (FormatException)
        {
            data = [];
            return false;
        }
    }

    private static string StatusName(NodeStatus status) => status switch
    {
        NodeStatus.Active => "active",
        NodeStatus.Pending => "pending",
        _ => "disabled",
    };

    private sealed record Frame(bool IsText, byte[] Data);
}

[System.Text.Json.Serialization.JsonSerializable(typeof(EmptyObject))]
internal sealed partial class FleetEmptyObjectContext : System.Text.Json.Serialization.JsonSerializerContext;

internal sealed class EmptyObject;

internal static class FleetEmptyObject
{
    public static System.Text.Json.Serialization.Metadata.JsonTypeInfo<EmptyObject> Info => FleetEmptyObjectContext.Default.EmptyObject;
}