using System.Collections.Concurrent;
using System.Diagnostics;
using System.Net;
using System.Net.Sockets;
using System.Security.Cryptography;
using System.Text.Json;
using Fleet.Core;
using Fleet.Identity;
using Fleet.Link;
using Fleet.Rpc;
using MoShiro.Protocol;

namespace MoShiro.TestSupport;

/// <summary>
/// 假 Probe：通过真实的点对点握手（TLS 固定证书 + 签名 + ECDH）连接真实的 Agent，接收 <c>line_task_set</c>，
/// 在自己的 UDP 套接字上发回显包、接收单向数据流、估计时钟偏移，并把 Probe 一侧的列块经发件箱通道 17 交给 Agent。
/// 它实现 Probe 的协议职责，是将来真实 Probe 的参照。
/// </summary>
public sealed class ProbeSimulator : IAsyncDisposable
{
    private readonly NodeIdentity _identity;
    private byte[]? _agentKey;
    private readonly RpcRouter<ProbeSimulator> _router;
    private readonly ConcurrentDictionary<Guid, TaskCompletionSource> _acks = new();
    private readonly ConcurrentDictionary<uint, EchoCollector> _echoes = new();
    private readonly ConcurrentDictionary<uint, TrainCollector> _trains = new();
    private readonly ConcurrentQueue<LineTaskSet> _taskSets = new();
    private readonly long _origin = TimeProvider.System.GetTimestamp();
    private readonly CancellationTokenSource _stop = new();
    private NodeId _agentId;
    private RpcSession<ProbeSimulator>? _rpc;
    private UdpClient? _udp;
    private IPEndPoint? _agentUdp;
    private PeerSessionInfo? _info;
    private Task? _receiver;
    private Task? _run;

    public ProbeSimulator(string probeId, NodeIdentity identity)
    {
        ProbeId = probeId;
        _identity = identity;
        _router = BuildRouter();
    }

    public string ProbeId { get; }

    /// <summary>同一个身份的另一个假 Probe（模拟 Probe 重启：新的进程、新的会话，序号从 0 重新开始）。</summary>
    public ProbeSimulator Restarted() => new(ProbeId, _identity);

    /// <summary>true 时拒绝按需测试的阶段（用来测试 Agent 对 Probe 拒绝的处理）。</summary>
    public bool RefusePhases { get; set; }

    /// <summary>收到 phase_end 之后不交出列块，只报告 0 个（用来测试 Agent 的等待逻辑）。</summary>
    public bool DropFragments { get; set; }

    public long ClockOffsetUs { get; private set; }

    public IReadOnlyCollection<LineTaskSet> TaskSets => [.. _taskSets];

    public ConcurrentQueue<TestPhaseStartParams> Phases { get; } = new();

    public ConcurrentQueue<TestId> Aborts { get; } = new();

    public bool IsConnected => _rpc is not null && !_stop.IsCancellationRequested;

    public long NowUs() => (TimeProvider.System.GetTimestamp() - _origin) * 1_000_000 / TimeProvider.System.TimestampFrequency;

    public async Task ConnectAsync(string host, int controlPort, byte[] agentPublicKey, int udpPort, NodeId agentId, TimeProvider? time = null, CancellationToken ct = default)
    {
        _agentId = agentId;
        (PeerSessionInfo info, TlsFrameTransport transport) = await PeerConnector.ConnectAsync(host, controlPort, agentPublicKey, _identity, ProbeId, "0.0.0-sim", 8 << 20, time ?? TimeProvider.System, ct);
        _info = info;
        _agentKey = agentPublicKey;
        _udp = new UdpClient(new IPEndPoint(IPAddress.Loopback, 0));
        _agentUdp = new IPEndPoint(IPAddress.Parse(host), udpPort);
        _rpc = new RpcSession<ProbeSimulator>(transport, _router, this, time ?? TimeProvider.System);
        _rpc.OnBinary = (_, _) => ValueTask.CompletedTask;
        _run = _rpc.RunAsync(_stop.Token);
        _receiver = Task.Run(ReceiveLoopAsync, _stop.Token);
        await SyncClockAsync();
    }

    public async Task<LineTaskSet> WaitForTasksAsync(Func<LineTaskSet, bool> predicate, TimeSpan? timeout = null)
    {
        long started = Stopwatch.GetTimestamp();
        TimeSpan limit = timeout ?? TimeSpan.FromSeconds(15);
        while (Stopwatch.GetElapsedTime(started) < limit)
        {
            LineTaskSet? found = _taskSets.ToArray().LastOrDefault(predicate);
            if (found is not null)
            {
                return found;
            }

            await Task.Delay(20);
        }

        throw new TimeoutException("the agent did not send the expected line task set");
    }

    /// <summary>NTP 式的时钟同步：偏移 = Agent 时钟 − (往返中点的) Probe 时钟。取最小往返的一次。</summary>
    public async Task SyncClockAsync(int rounds = 5)
    {
        long bestRoundTrip = long.MaxValue;
        for (int i = 0; i < rounds; i++)
        {
            long t0 = NowUs();
            Result<ClockResult> reply = await _rpc!.RequestAsync(MoShiroMethods.Clock, new ClockParams { ProbeMonoUs = t0 }, MoShiroJsonContext.Default.ClockParams, MoShiroJsonContext.Default.ClockResult, CancellationToken.None);
            long t1 = NowUs();
            long roundTrip = t1 - t0;
            if (reply.IsSuccess && roundTrip < bestRoundTrip)
            {
                bestRoundTrip = roundTrip;
                ClockOffsetUs = reply.Value.AgentMonoUs - ((t0 + t1) / 2);
            }
        }
    }

    /// <summary>对一个 udp_echo 线路任务发 <paramref name="count"/> 个回显包，等回复，然后交出一个列块。返回交给 Agent 的列块的序号范围。</summary>
    public async Task<(int Sent, int Received)> RunUdpEchoAsync(LineTask task, int count, TimeSpan interval, CancellationToken ct = default)
    {
        int size = (int)task.Params["size"].GetInt64();
        var collector = new EchoCollector();
        _echoes[task.StreamId] = collector;
        long startUtc = TimeProvider.System.GetUtcNow().ToUnixTimeMilliseconds();
        byte[] packet = new byte[size];
        for (uint seq = 0; seq < count; seq++)
        {
            long t1 = NowUs();
            collector.Sent[seq] = t1;
            UdpPacket.Write(packet, new UdpHeader(UdpPacketType.EchoRequest, _info!.Credentials.SessionIdValue, task.StreamId, seq, t1), _info.Credentials.Key);
            await _udp!.SendAsync(packet, packet.Length, _agentUdp!);
            await Task.Delay(interval, ct);
        }

        await Task.Delay(300, ct); // 等最后的回复
        _echoes.TryRemove(task.StreamId, out _);

        var msg = new StreamBuilder(new StreamSpec("msg", [Col.Seq("seq"), Col.U16("size"), Col.TsProbe("t1"), Col.TsProbe("t4", nullable: true)]));
        foreach ((uint seq, long t1) in collector.Sent.OrderBy(p => p.Key))
        {
            msg.Add(seq, size, t1, collector.Received.TryGetValue(seq, out long t4) ? t4 : null);
        }

        RawChunkHeader header = Header("udp_echo", task.Id, null, null, startUtc, (int)Math.Max(1, count * interval.TotalMilliseconds), task.Params);
        await ShipAsync(new RawChunk(header, [msg.Build()]), ct: ct);
        return (count, collector.Received.Count);
    }

    /// <summary>对一个 tcp_echo 线路任务建立一条 TCP（可选 TLS）连接，发 <paramref name="count"/> 条回显消息，然后交出一个列块。</summary>
    public async Task<(int Sent, int Received)> RunTcpEchoAsync(LineTask task, int tcpPort, int count, TimeSpan interval, bool tls, CancellationToken ct = default)
    {
        using var client = new System.Net.Sockets.TcpClient();
        await client.ConnectAsync(IPAddress.Loopback, tcpPort, ct);
        Stream stream = client.GetStream();
        await stream.WriteAsync(Knock.Build(Knock.DeriveKey(_identity, _agentKey!), TimeProvider.System.GetUtcNow().ToUnixTimeMilliseconds()), ct); // 预握手敲门
        if (tls)
        {
#pragma warning disable CA5359 // 测试：Agent 的证书是自签名的，这里只验证数据面
            var ssl = new System.Net.Security.SslStream(stream, false, (_, _, _, _) => true);
#pragma warning restore CA5359
            await ssl.AuthenticateAsClientAsync(new System.Net.Security.SslClientAuthenticationOptions { TargetHost = "agent" }, ct);
            stream = ssl;
        }

        byte[] auth = new byte[TcpWire.AuthBytes];
        TcpWire.WriteAuth(auth, _info!.Credentials.SessionIdValue, task.StreamId, RandomNumberGenerator.GetBytes(TcpWire.NonceBytes), _info.Credentials.Key);
        await stream.WriteAsync(auth, ct);

        long startUtc = TimeProvider.System.GetUtcNow().ToUnixTimeMilliseconds();
        var msg = new StreamBuilder(new StreamSpec("msg", [Col.Seq("seq"), Col.U16("size"), Col.TsProbe("t1"), Col.TsProbe("t4", nullable: true)]));
        int received = 0;
        for (uint seq = 0; seq < count; seq++)
        {
            byte[] body = new byte[48];
            long t1 = NowUs();
            TcpWire.WriteEchoBody(body, seq, t1);
            await TcpWire.WriteMessageAsync(stream, body, ct);
            await stream.FlushAsync(ct);
            byte[]? reply = await TcpWire.ReadMessageAsync(stream, ct);
            long? t4 = reply is not null && TcpWire.TryReadEchoBody(reply, out uint echoed, out _) && echoed == seq ? NowUs() : null;
            received += t4 is null ? 0 : 1;
            msg.Add(seq, body.Length, t1, t4);
            await Task.Delay(interval, ct);
        }

        RawChunkHeader header = Header("tcp_echo", task.Id, null, null, startUtc, (int)Math.Max(1, count * interval.TotalMilliseconds), task.Params);
        await ShipAsync(new RawChunk(header, [msg.Build()]), ct: ct);
        return (count, received);
    }

    /// <summary>发送任意字节到 Agent 的 UDP 端口（用来测试 Agent 对垃圾包的处理）。</summary>
    public Task SendRawUdpAsync(byte[] datagram) => _udp!.SendAsync(datagram, datagram.Length, _agentUdp!);

    /// <summary>用伪造的密钥发一个回显包。</summary>
    public Task SendForgedEchoAsync(LineTask task, uint seq, byte[]? key = null)
    {
        byte[] packet = new byte[64];
        UdpPacket.Write(packet, new UdpHeader(UdpPacketType.EchoRequest, _info!.Credentials.SessionIdValue, task.StreamId, seq, NowUs()), key ?? new byte[32]);
        return _udp!.SendAsync(packet, packet.Length, _agentUdp!);
    }

    /// <summary>重放：用同一个序号、同样的密钥再发一次。</summary>
    public Task SendEchoAsync(LineTask task, uint seq)
    {
        byte[] packet = new byte[64];
        UdpPacket.Write(packet, new UdpHeader(UdpPacketType.EchoRequest, _info!.Credentials.SessionIdValue, task.StreamId, seq, NowUs()), _info.Credentials.Key);
        return _udp!.SendAsync(packet, packet.Length, _agentUdp!);
    }

    public int EchoRepliesSeen(uint streamId) => _echoes.TryGetValue(streamId, out EchoCollector? c) ? c.Received.Count : 0;

    /// <summary>交出一个已经构造好的列块并等待 Agent 确认。</summary>
    public async Task ShipAsync(RawChunk fragment, bool waitForAck = true, CancellationToken ct = default)
    {
        Guid id = fragment.Header.ChunkId.Value;
        var ack = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
        _acks[id] = ack;
        await _rpc!.SendBinaryAsync(OutboxFrame.Build(MoShiroChannels.ColumnBlock, id, RawChunkCodec.Encode(fragment)), ct);
        if (waitForAck)
        {
            await ack.Task.WaitAsync(TimeSpan.FromSeconds(10), ct);
        }
    }

    public async ValueTask DisposeAsync()
    {
        await _stop.CancelAsync();
        _udp?.Dispose();
        if (_rpc is not null)
        {
            await _rpc.DisposeAsync();
        }

        foreach (Task? task in new[] { _run, _receiver })
        {
            if (task is not null)
            {
                try
                {
                    await task.WaitAsync(TimeSpan.FromSeconds(2));
                }
                catch (Exception ex) when (ex is OperationCanceledException or TimeoutException or ObjectDisposedException or SocketException or IOException)
                {
                    // 关闭中
                }
            }
        }

        _stop.Dispose();
    }

    private RawChunkHeader Header(string kind, TaskId? task, TestId? test, int? phase, long startUtcMs, int durationMs, IReadOnlyDictionary<string, JsonElement>? extra)
    {
        var header = new RawChunkHeader
        {
            ChunkId = new ChunkId(Guid.CreateVersion7()),
            Kind = kind,
            AgentId = _agentId,
            TaskId = task,
            TestId = test,
            Phase = phase,
            ProbeId = new ProbeId(ProbeId),
            StartUtcMs = startUtcMs,
            DurationMs = durationMs,
            Clock = new ChunkClock { ProbeOffsetUs = ClockOffsetUs },
        };
        if (extra is not null)
        {
            foreach ((string key, JsonElement value) in extra)
            {
                header.Params[key] = value;
            }
        }

        return header;
    }

    private async Task ReceiveLoopAsync()
    {
        while (!_stop.IsCancellationRequested)
        {
            UdpReceiveResult result;
            try
            {
                result = await _udp!.ReceiveAsync(_stop.Token);
            }
            catch (Exception ex) when (ex is OperationCanceledException or ObjectDisposedException or SocketException)
            {
                if (_stop.IsCancellationRequested || ex is ObjectDisposedException)
                {
                    return;
                }

                continue;
            }

            long now = NowUs();
            byte[] data = result.Buffer;
            if (!UdpPacket.TryParseHeader(data, out UdpHeader header) || !UdpPacket.VerifyMac(data, _info!.Credentials.Key))
            {
                continue;
            }

            if (header.Type == UdpPacketType.EchoReply && _echoes.TryGetValue(header.StreamId, out EchoCollector? echo))
            {
                echo.Received.TryAdd(header.Seq, now);
            }
            else if (header.Type == UdpPacketType.Data && _trains.TryGetValue(header.StreamId, out TrainCollector? train))
            {
                train.Received.TryAdd(header.Seq, now);
                train.Size = data.Length;
            }
        }
    }

    private RpcRouter<ProbeSimulator> BuildRouter() =>
        new RpcRouter<ProbeSimulator>(RpcSide.Dialer, MoShiroMethodTable.Table)
            .On(
                FleetMethods.StateSync,
                FleetRpcContext.Default.StateSyncParams,
                FleetRpcContext.Default.StateSyncResult,
                (probe, sync, _) =>
                {
                    if (sync.Type == LineTaskSet.StateType)
                    {
                        probe._taskSets.Enqueue(sync.Doc.Deserialize(MoShiroJsonContext.Default.LineTaskSet)!);
                    }

                    return ValueTask.FromResult(Result<StateSyncResult>.Ok(new StateSyncResult { AppliedVersion = sync.Version }));
                })
            .On(
                MoShiroMethods.TestPhase,
                MoShiroJsonContext.Default.TestPhaseStartParams,
                MoShiroJsonContext.Default.TestPhaseStartResult,
                async (probe, start, ct) => Result<TestPhaseStartResult>.Ok(new TestPhaseStartResult { Accepted = await probe.StartPhaseAsync(start, ct) }))
            .On(
                MoShiroMethods.TestPhaseEnd,
                MoShiroJsonContext.Default.TestPhaseEndParams,
                MoShiroJsonContext.Default.TestPhaseEndResult,
                async (probe, end, ct) => Result<TestPhaseEndResult>.Ok(new TestPhaseEndResult { Fragments = await probe.EndPhaseAsync(end, ct) }))
            .OnNotify(
                MoShiroMethods.TestAbort,
                MoShiroJsonContext.Default.TestAbortParams,
                (probe, abort, _) =>
                {
                    probe.Aborts.Enqueue(abort.TestId);
                    return ValueTask.CompletedTask;
                })
            .OnNotify(
                FleetMethods.OutboxAck,
                FleetRpcContext.Default.OutboxAckParams,
                (probe, ack, _) =>
                {
                    foreach (string id in ack.Ids.Concat(ack.Duplicates))
                    {
                        if (probe._acks.TryRemove(Guid.Parse(id), out TaskCompletionSource? waiter))
                        {
                            waiter.TrySetResult();
                        }
                    }

                    return ValueTask.CompletedTask;
                });

    private readonly ConcurrentDictionary<uint, TestPhaseStartParams> _activePhases = new();

    private async Task<bool> StartPhaseAsync(TestPhaseStartParams start, CancellationToken ct)
    {
        Phases.Enqueue(start);
        if (RefusePhases)
        {
            return false;
        }

        _activePhases[start.StreamId] = start;
        if (start.Kind == "udp_train")
        {
            _trains[start.StreamId] = new TrainCollector { StartUtcMs = TimeProvider.System.GetUtcNow().ToUnixTimeMilliseconds() };
            // 先发几个探活包，让 Agent 学到（最新的）源地址，再回复 accepted
            byte[] hello = new byte[64];
            for (uint seq = 0; seq < 3; seq++)
            {
                UdpPacket.Write(hello, new UdpHeader(UdpPacketType.Data, _info!.Credentials.SessionIdValue, start.StreamId, seq, NowUs()), _info.Credentials.Key);
                await _udp!.SendAsync(hello.AsMemory(), _agentUdp!, ct);
            }
        }

        return true;
    }

    private async Task<int> EndPhaseAsync(TestPhaseEndParams end, CancellationToken ct)
    {
        _activePhases.TryRemove(end.StreamId, out TestPhaseStartParams? start);
        if (start is null || DropFragments)
        {
            _trains.TryRemove(end.StreamId, out _);
            return 0;
        }

        if (start.Kind == "udp_train" && _trains.TryRemove(end.StreamId, out TrainCollector? train))
        {
            await Task.Delay(200, ct); // 最后一批包还在路上
            // 单块的记录数有上限（ProtocolLimits.MaxRecordsPerChunk）：按序号范围切成多个列块，Agent 一侧按同样的范围取行
            const int PerFragment = 40_000;
            uint total = (uint)end.SentCount;
            int fragments = 0;
            for (uint from = 0; from < total || fragments == 0; from += PerFragment)
            {
                uint to = (uint)Math.Min((long)from + PerFragment, total);
                var pkt = new StreamBuilder(new StreamSpec("pkt", [Col.Seq("seq"), Col.TsProbe("t_recv", nullable: true)]));
                foreach ((uint seq, long recv) in train.Received.Where(p => p.Key >= from && p.Key < to).OrderBy(p => p.Key))
                {
                    pkt.Add(seq, recv);
                }

                var streams = new List<RawStream> { pkt.Build() };
                if (fragments == 0)
                {
                    var step = new StreamBuilder(Kinds.UdpTrain.FindStream("step")!);
                    step.Add(0, (long)start.Params["mbps"].GetDouble(), NowUs() - (start.DurationSeconds * 1_000_000L), NowUs());
                    streams.Add(step.Build());
                }

                RawChunkHeader header = Header("udp_train", null, start.TestId, start.Phase, train.StartUtcMs, start.DurationSeconds * 1000, start.Params);
                header.Params[FragmentKeys.SeqFrom] = JsonSerializer.SerializeToElement(from);
                header.Params[FragmentKeys.SeqTo] = JsonSerializer.SerializeToElement(to);
                // 不能在这里等确认：确认是 Agent 的通知，由同一个入站队列顺序处理，而当前正处在 phase_end 的处理函数里
                await ShipAsync(new RawChunk(header, streams), waitForAck: false, ct: ct);
                fragments++;
                if (total == 0)
                {
                    break;
                }
            }

            return fragments;
        }

        return 0;
    }

    private sealed class EchoCollector
    {
        public ConcurrentDictionary<uint, long> Sent { get; } = new();

        public ConcurrentDictionary<uint, long> Received { get; } = new();
    }

    private sealed class TrainCollector
    {
        public ConcurrentDictionary<uint, long> Received { get; } = new();

        public int Size { get; set; }

        public long StartUtcMs { get; init; }
    }
}