using System.Net.Sockets;
using MoShiro.Protocol;

namespace MoShiro.Measurement;

/// <summary>批量传输的收发循环（两端共用）。</summary>
internal static class BulkIo
{
    private const int ChunkBytes = 16 * 1024;

    /// <summary>
    /// 发送：尽可能快，但不超过 <paramref name="mbps"/>（本地硬上限，防止把小机器的 CPU 和链路打满）。
    /// 发送的数据是随机字节（避免被链路上的压缩或重复数据消除"优化"掉）。
    /// </summary>
    public static async Task SendAsync(Stream stream, double mbps, Action<int> onBytes, TimeProvider time, CancellationToken ct)
    {
        byte[] buffer = new byte[ChunkBytes];
        Random.Shared.NextBytes(buffer);
        double bytesPerSecond = mbps * 1_000_000 / 8;
        long start = time.GetTimestamp();
        long sent = 0;
        while (!ct.IsCancellationRequested)
        {
            await stream.WriteAsync(buffer, ct);
            sent += buffer.Length;
            onBytes(buffer.Length);
            if (bytesPerSecond > 0)
            {
                double ahead = (sent / bytesPerSecond) - time.GetElapsedTime(start).TotalSeconds;
                if (ahead > 0.002)
                {
                    await Task.Delay(TimeSpan.FromSeconds(Math.Min(ahead, 0.05)), time, ct);
                }
            }
        }
    }

    /// <summary>接收并丢弃，直到对端关闭或取消。</summary>
    public static async Task ReceiveAsync(Stream stream, Action<int> onBytes, CancellationToken ct)
    {
        byte[] buffer = new byte[4 * ChunkBytes];
        while (!ct.IsCancellationRequested)
        {
            int read = await stream.ReadAsync(buffer, ct);
            if (read <= 0)
            {
                return;
            }

            onBytes(read);
        }
    }

    public static bool IsExpected(Exception ex) =>
        ex is OperationCanceledException or IOException or SocketException or ObjectDisposedException or System.Security.Authentication.AuthenticationException or InvalidOperationException;
}

/// <summary>
/// 批量传输，拨号方（Probe）：开 <c>streams</c> 条 TCP 连接（可选 TLS），按 <c>direction</c> 发送（fwd）或接收（rev），持续 <c>duration_s</c>。
/// 记录每个采样间隔每条连接的<b>累计字节数</b>（<c>bytes</c> 流，Probe 的时钟），Backend 据此求每个间隔的吞吐；
/// 同时采样第一条连接自己这一侧的内核 <c>tcp_info</c>（多条连接的累计计数器不能混在同一个序列里求差）。
/// 发送速率不超过本地上限 <see cref="BlockEnv.MaxBulkMbps"/>。
/// </summary>
public sealed class TcpBulkDialer : BlockBase
{
    private const int MaxRowsPerFragment = 40_000;

    private readonly DriverLog _log;
    private readonly TcpInfoRecorder _info;
    private readonly int _ts;
    private readonly int _conn;
    private readonly int _bytes;
    private long[] _counters = [];

    public TcpBulkDialer(BlockEnv env, BlockSpec spec)
        : base(env, spec)
    {
        long expectedRows = spec.Int("streams") * spec.Int("duration_s") * Math.Max(1, 1000 / Math.Max(1, spec.Int("sample_ms")));
        _log = new DriverLog(spec.Kind.FindStream("bytes")!, Side.Probe, (int)Math.Clamp(expectedRows + 256, 1024, 262_144));
        _info = new TcpInfoRecorder(spec.Kind.FindStream("tcp_info")!, Side.Probe, env.Clock);
        _ts = _log.Store.ColumnIndex("ts");
        _conn = _log.Store.ColumnIndex("conn");
        _bytes = _log.Store.ColumnIndex("bytes");
    }

    public override long Produced => _log.Count;

    public override async Task RunAsync(CancellationToken ct)
    {
        int streams = (int)Spec.Int("streams");
        bool send = Spec.Params.GetString("direction") == "fwd";
        TimeSpan duration = TimeSpan.FromSeconds(Spec.Int("duration_s"));
        TimeSpan sampleEvery = TimeSpan.FromMilliseconds(Spec.Int("sample_ms"));
        _counters = new long[streams];

        var links = new TcpLink?[streams];
        try
        {
            await Task.WhenAll(Enumerable.Range(0, streams).Select(async i =>
            {
                try
                {
                    links[i] = await TcpLink.OpenAsync(Env, Spec.Port, Spec.Bool("tls"), Spec.StreamId, sendAuth: true, ct);
                }
                catch (Exception ex) when (BulkIo.IsExpected(ex))
                {
                    // 这条连接没有建起来：它的字节数一直是 0
                }
            }));

            using var timeout = new CancellationTokenSource(duration, Env.Time);
            using var run = CancellationTokenSource.CreateLinkedTokenSource(ct, timeout.Token);
            var work = new List<Task>();
            double perStream = Env.MaxBulkMbps / streams;
            for (int i = 0; i < streams; i++)
            {
                if (links[i] is not { } link)
                {
                    continue;
                }

                int index = i;
                work.Add(RunConnectionAsync(link.Stream, index, send, perStream, run.Token));
            }

            if (links.FirstOrDefault(l => l is not null) is { } first)
            {
                work.Add(_info.RunAsync(first.Socket, sampleEvery, run.Token));
            }

            Task sampler = SampleAsync(sampleEvery, run.Token);
            await Task.WhenAll(work);
            await run.CancelAsync();
            await sampler;
            Sample(); // 最后一个采样点
        }
        finally
        {
            foreach (TcpLink? link in links)
            {
                link?.Dispose();
            }
        }
    }

    public override int Flush(FlushMode mode, long peerProduced = 0)
    {
        int fragments = 0;
        while (_log.Take(null, all: true, MaxRowsPerFragment) is { } stream)
        {
            RawStream kernel = _info.Rows.Drain(includeKey: false);
            Ship(stream, "ts", kernel.Count > 0 ? [kernel] : null);
            fragments++;
        }

        if (fragments == 0 && _info.Rows.Drain(includeKey: false) is { Count: > 0 } onlyKernel)
        {
            Ship(onlyKernel, "ts");
            fragments++;
        }

        return fragments;
    }

    private async Task RunConnectionAsync(Stream stream, int index, bool send, double mbps, CancellationToken ct)
    {
        try
        {
            if (send)
            {
                await BulkIo.SendAsync(stream, mbps, n => Interlocked.Add(ref _counters[index], n), Env.Time, ct);
            }
            else
            {
                await BulkIo.ReceiveAsync(stream, n => Interlocked.Add(ref _counters[index], n), ct);
            }
        }
        catch (Exception ex) when (BulkIo.IsExpected(ex))
        {
            // 时间到或者连接断开：计数停在这里
        }
    }

    private async Task SampleAsync(TimeSpan every, CancellationToken ct)
    {
        using var timer = new PeriodicTimer(every, Env.Time);
        try
        {
            while (await timer.WaitForNextTickAsync(ct))
            {
                Sample();
            }
        }
        catch (OperationCanceledException)
        {
            // 阶段结束
        }
    }

    private void Sample()
    {
        long now = Env.Clock.NowUs();
        for (int i = 0; i < _counters.Length; i++)
        {
            long key = _log.Store.Append();
            _log.Store.Set(key, _ts, now);
            _log.Store.Set(key, _conn, i);
            _log.Store.Set(key, _bytes, Interlocked.Read(ref _counters[i]));
        }
    }
}

/// <summary>
/// 批量传输，监听方（Agent）：接受连接，<c>rev</c> 时向 Probe 发送（受本地码率上限约束），<c>fwd</c> 时接收并丢弃；
/// 采样第一条连接自己这一侧的内核 <c>tcp_info</c>（Agent 的重传反映下行方向，代理场景里大流量走这个方向）。
/// 连接在 <c>duration_s</c> 之后（加几秒余量）或阶段结束时关闭。
/// </summary>
public sealed class TcpBulkListener : BlockBase
{
    private readonly TcpInfoRecorder _info;
    private CancellationToken _stop;
    private int _sampled;

    public TcpBulkListener(BlockEnv env, BlockSpec spec)
        : base(env, spec)
    {
        _info = new TcpInfoRecorder(spec.Kind.FindStream("tcp_info")!, Side.Agent, env.Clock);
    }

    public override RowStore? Records(string streamName) => streamName == "tcp_info" ? _info.Rows : null;

    public override async Task RunAsync(CancellationToken ct)
    {
        _stop = ct;
        Env.TcpListen!.Register(Spec.StreamId, HandleAsync);
        try
        {
            await WaitUntilCancelledAsync(ct);
        }
        finally
        {
            Env.TcpListen.Unregister(Spec.StreamId);
        }
    }

    private async Task HandleAsync(TcpConnection connection, CancellationToken acceptorToken)
    {
        using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(Spec.Int("duration_s") + 3), Env.Time);
        using var linked = CancellationTokenSource.CreateLinkedTokenSource(acceptorToken, _stop, timeout.Token);
        Task sampler = Task.CompletedTask;
        if (Interlocked.CompareExchange(ref _sampled, 1, 0) == 0)
        {
            sampler = _info.RunAsync(connection.Socket, TimeSpan.FromMilliseconds(Spec.Int("sample_ms")), linked.Token);
        }

        try
        {
            if (Spec.Params.GetString("direction") == "rev")
            {
                await BulkIo.SendAsync(connection.Stream, Env.MaxBulkMbps / Math.Max(1, Spec.Int("streams")), _ => Env.Counters.PacketOut(), Env.Time, linked.Token);
            }
            else
            {
                await BulkIo.ReceiveAsync(connection.Stream, _ => Env.Counters.PacketIn(), linked.Token);
            }
        }
        catch (Exception ex) when (BulkIo.IsExpected(ex))
        {
            // 时间到、阶段结束或对端关闭
        }
        finally
        {
            await linked.CancelAsync();
            await sampler;
        }
    }
}