using System.Diagnostics;
using System.Net;
using MoShiro.Protocol;

namespace MoShiro.Measurement;

/// <summary>
/// 单向数据流（UDP train），监听方（Agent，行的主导方）：按声明的码率均匀发包，发出的每个包 (seq, size, t_send) 记在按序号的记录里，
/// Probe 的接收时刻之后按序号合并（没有到达的就是丢失）。发送在专用线程上完成，不占用线程池。
/// 等不到对端的地址（5 秒内没有任何认证过的包）就以 <see cref="BlockFaultException"/> 结束。
/// </summary>
public sealed class TrainListener : BlockBase
{
    private static readonly TimeSpan s_endpointWait = TimeSpan.FromSeconds(5);

    private readonly RowStore _records;
    private long _sent;

    public TrainListener(BlockEnv env, BlockSpec spec)
        : base(env, spec)
    {
        double pps = spec.Params.GetNumber("mbps") * 1_000_000 / 8 / Math.Max(1, spec.Int("size"));
        _records = new RowStore(spec.Kind.FindStream("pkt")!, Side.Agent, (int)Math.Clamp(pps * spec.Int("duration_s") * 1.25, 1024, 4_000_000));
    }

    public override long Produced => Interlocked.Read(ref _sent);

    public override RowStore? Records(string streamName) => streamName == "pkt" ? _records : null;

    public override async Task RunAsync(CancellationToken ct)
    {
        // Probe 开始前发来的探活包（Data）要被接受，准入策略才能据此学到它当前的源地址（NAT 映射）；它们本身没有内容
        Env.Udp!.Register(Spec.StreamId, static (in UdpDelivery _) => { });
        try
        {
            DataSession session = await WaitForPeerAsync(ct);
            IPEndPoint target = session.UdpEndpoint!;
            TimeSpan duration = TimeSpan.FromSeconds(Spec.Int("duration_s"));
            var done = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
            var thread = new Thread(() =>
            {
                try
                {
                    Send(session, target, duration, ct);
                    done.SetResult();
                }
                catch (Exception ex)
                {
                    done.SetException(ex);
                }
            })
            {
                IsBackground = true,
                Name = "udp-train",
            };
            thread.Start();
            await done.Task;
        }
        finally
        {
            Env.Udp.Unregister(Spec.StreamId);
        }
    }

    private async Task<DataSession> WaitForPeerAsync(CancellationToken ct)
    {
        DateTimeOffset deadline = Env.Time.GetUtcNow() + s_endpointWait;
        while (true)
        {
            DataSession? session = Env.Session();
            if (session?.UdpEndpoint is not null)
            {
                return session;
            }

            if (Env.Time.GetUtcNow() >= deadline)
            {
                throw new BlockFaultException("the peer's address is unknown");
            }

            await Task.Delay(TimeSpan.FromMilliseconds(50), Env.Time, ct);
        }
    }

    private void Send(DataSession session, IPEndPoint target, TimeSpan duration, CancellationToken ct)
    {
        double packetsPerSecond = Spec.Params.GetNumber("mbps") * 1_000_000 / 8 / Spec.Int("size");
        int size = (int)Spec.Int("size");
        byte[] packet = new byte[size];
        using var pause = new ManualResetEventSlim(false);
        uint seq = 0;
        long startTicks = Stopwatch.GetTimestamp();
        long endTicks = startTicks + (long)(duration.TotalSeconds * Stopwatch.Frequency);
        while (!ct.IsCancellationRequested)
        {
            long now = Stopwatch.GetTimestamp();
            if (now >= endTicks)
            {
                break;
            }

            // 到目前为止应该发出的包数（均匀）：落后了就追，领先了就等；追赶的上限是 50 ms 的量，避免长时间停顿之后的突发
            double elapsed = (now - startTicks) / (double)Stopwatch.Frequency;
            long due = Math.Min((long)(elapsed * packetsPerSecond) + 1, seq + Math.Max(1, (long)(packetsPerSecond * 0.05)));
            if (seq >= due)
            {
                // 离下一个包还早就让出 CPU（至少 1 ms 的粒度），临近了再自旋，保证发送间隔均匀
                if ((seq / packetsPerSecond) - elapsed > 0.002)
                {
                    pause.Wait(1, ct);
                }
                else
                {
                    Thread.SpinWait(50);
                }

                continue;
            }

            for (; seq < due; seq++)
            {
                long sendUs = Env.Clock.NowUs();
                UdpPacket.Write(packet, new UdpHeader(UdpPacketType.Data, session.SessionId, Spec.StreamId, seq, sendUs), session.Key);
                Env.Udp!.Send(target, packet);
                Span<long> row = [size, sendUs];
                _records.SetRow(seq, row);
                Interlocked.Exchange(ref _sent, seq + 1L);
            }
        }
    }
}

/// <summary>
/// 单向数据流，拨号方（Probe，行的跟随方）：接收，记录每个序号第一次到达的时刻（按序号索引，热路径不分配）；
/// 序号超过本地包速率上限推算出的范围时忽略。阶段结束时按序号范围切成多个列块交给 Agent，每块不超过 <see cref="ProtocolLimits.MaxRecordsPerChunk"/> 的记录。
/// 开始时先发几个探活包，让 Agent 学到我们当前的源地址（NAT 映射）再开始发送。
/// </summary>
public sealed class TrainDialer : BlockBase
{
    private const int PerFragment = 40_000;

    private readonly RowStore _received;
    private readonly RowStore _step;
    private readonly int _tRecv;
    private readonly long _maxSeq;
    private long _startUs;
    private long _endUs;

    public TrainDialer(BlockEnv env, BlockSpec spec)
        : base(env, spec)
    {
        double pps = spec.Params.GetNumber("mbps") * 1_000_000 / 8 / Math.Max(1, spec.Int("size"));
        long seconds = Math.Max(1, spec.Int("duration_s"));
        _maxSeq = Math.Max(1024, Math.Min((long)(pps * seconds * 1.25), env.MaxPacketsPerSecond * seconds * 3 / 2));
        _received = new RowStore(spec.Kind.FindStream("pkt")!, Side.Probe, (int)_maxSeq);
        _tRecv = _received.ColumnIndex("t_recv");
        _step = new RowStore(spec.Kind.FindStream("step")!, Side.Probe, 16);
    }

    public override async Task RunAsync(CancellationToken ct)
    {
        UdpMux udp = Env.Udp!;
        DataSession session = Env.Session() ?? throw new BlockFaultException("no data session");
        _startUs = Env.Clock.NowUs();
        udp.Register(Spec.StreamId, OnPacket);
        try
        {
            byte[] hello = new byte[UdpPacket.MinBytes + 20];
            IPEndPoint target = PeerEndpoint(Spec.Port);
            for (uint seq = 0; seq < 3; seq++)
            {
                UdpPacket.Write(hello, new UdpHeader(UdpPacketType.Data, session.SessionId, Spec.StreamId, seq, Env.Clock.NowUs()), session.Key);
                udp.Send(target, hello);
            }

            await WaitUntilCancelledAsync(ct);
        }
        finally
        {
            _endUs = Env.Clock.NowUs();
            udp.Unregister(Spec.StreamId);
        }
    }

    public override int Flush(FlushMode mode, long peerProduced = 0)
    {
        if (mode != FlushMode.All)
        {
            return 0;
        }

        long end = _endUs == 0 ? Env.Clock.NowUs() : _endUs;
        long step = _step.Append();
        _step.Set(step, _step.ColumnIndex("index"), 0);
        _step.Set(step, _step.ColumnIndex("mbps"), (long)Spec.Params.GetNumber("mbps"));
        _step.Set(step, _step.ColumnIndex("start_ts"), _startUs);
        _step.Set(step, _step.ColumnIndex("end_ts"), end);

        long total = Math.Min(peerProduced, _maxSeq);
        int fragments = 0;
        for (long from = 0; from < total || fragments == 0; from += PerFragment)
        {
            long to = Math.Min(from + PerFragment, total);
            RawStream pkt = _received.Cut(from, to, includeKey: true);
            var parameters = new Dictionary<string, System.Text.Json.JsonElement>(Spec.Params.Values)
            {
                [FragmentKeys.SeqFrom] = Number(from),
                [FragmentKeys.SeqTo] = Number(to),
            };
            Env.Fragments!.Ship(
                "udp_train",
                Spec.Scope,
                fragments == 0 ? [pkt, _step.Cut(0, 1, includeKey: false)] : [pkt],
                parameters,
                Env.Clock.UtcMsAt(_startUs),
                (int)((end - _startUs) / 1000),
                Spec.Port);
            fragments++;
            if (total == 0)
            {
                break;
            }
        }

        return fragments;
    }

    private static System.Text.Json.JsonElement Number(long value) =>
        System.Text.Json.JsonDocument.Parse(value.ToString(System.Globalization.CultureInfo.InvariantCulture)).RootElement.Clone();

    private void OnPacket(in UdpDelivery delivery)
    {
        long seq = delivery.Header.Seq;
        if (delivery.Header.Type == UdpPacketType.Data && seq < _maxSeq && !_received.IsSet(seq, _tRecv))
        {
            _received.Set(seq, _tRecv, delivery.ReceiveUs); // 只记第一次到达（重复的包不覆盖）
        }
    }
}