using System.Text.Json;
using Fleet.Core;
using MoShiro.Protocol;

namespace MoShiro.Measurement.Tests;

public sealed class MeasurementTests
{
    private static JsonElement Json(string raw) => JsonDocument.Parse(raw).RootElement.Clone();

    private static TaskSpec Line(string kind, params (string Key, string Raw)[] parameters) => new()
    {
        Id = new TaskId("tsk_" + kind),
        Scope = "line",
        Kind = kind,
        Params = parameters.ToDictionary(p => p.Key, p => Json(p.Raw)),
    };

    private static BudgetGuard Guard(NodeLimits? limits = null) => new(limits ?? new NodeLimits { AllowedQuicPorts = [7443], AllowedTcpPorts = [7444] }, BlockCatalog.AgentKinds);

    [Fact]
    public void BudgetGuard_AcceptsDefaults_RejectsUnknownKindsBadParamsAndForeignPorts()
    {
        BudgetGuard guard = Guard();
        double mbps = 0;
        long pps = 0;
        Assert.True(guard.Check(Line("udp_echo"), ref mbps, ref pps).Accepted);
        Assert.Equal(RejectReasons.Unsupported, guard.Check(Line("no_such_kind"), ref mbps, ref pps).Reason);
        Assert.Equal(RejectReasons.InvalidParams, guard.Check(Line("udp_echo", ("rate_hz", "100000")), ref mbps, ref pps).Reason);
        Assert.Equal(RejectReasons.InvalidParams, guard.Check(Line("udp_echo", ("nope", "1")), ref mbps, ref pps).Reason);
        Assert.Equal(RejectReasons.PolicyDenied, guard.Check(Line("udp_echo", ("port", "9999")), ref mbps, ref pps).Reason);
    }

    [Fact]
    public void BudgetGuard_AccumulatesTheBudget_AndRejectsWhatDoesNotFit()
    {
        BudgetGuard guard = Guard(new NodeLimits { MaxTotalMbps = 15, AllowedQuicPorts = [7443] });
        double mbps = 0;
        long pps = 0;
        Assert.True(guard.CheckPhase("udp_train", new Dictionary<string, JsonElement> { ["mbps"] = Json("10") }).Accepted);
        Assert.Equal(RejectReasons.BudgetExceeded, guard.CheckPhase("udp_train", new Dictionary<string, JsonElement> { ["mbps"] = Json("20") }).Reason);

        Assert.True(guard.Check(Line("udp_train", ("mbps", "10")), ref mbps, ref pps).Accepted);
        Assert.Equal(RejectReasons.BudgetExceeded, guard.Check(Line("udp_train", ("mbps", "10")), ref mbps, ref pps).Reason);
    }

    [Fact]
    public void BudgetGuard_TargetTasks_AreDeniedUnlessTheAddressIsInAnAllowedRange()
    {
        Assert.False(Guard().TargetAllowed("1.1.1.1"));

        BudgetGuard some = Guard(new NodeLimits { AllowedTargetCidrs = ["192.0.2.0/24", "2001:db8::/32"] });
        Assert.True(some.TargetAllowed("192.0.2.17"));
        Assert.True(some.TargetAllowed("192.0.2.17:443"));
        Assert.True(some.TargetAllowed("2001:db8::1"));
        Assert.False(some.TargetAllowed("192.0.3.1"));
        Assert.False(some.TargetAllowed("example.com"));
        Assert.False(some.TargetAllowed(null));
    }

    // ---- RowStore ----

    [Fact]
    public void RowStore_RecordsOnlyItsOwnSidesColumns_AndCutsExistingRowsInKeyOrder()
    {
        StreamSpec msg = Kinds.UdpEcho.FindStream("msg")!;
        var probe = new RowStore(msg, Side.Probe, 64);
        Assert.Equal(["size", "t1", "t4"], msg.ColumnsOf(Side.Probe).Select(c => c.Name));
        Assert.Equal(["t2", "t3"], msg.ColumnsOf(Side.Agent).Select(c => c.Name));

        for (int i = 0; i < 4; i++)
        {
            long key = probe.Append();
            probe.Set(key, probe.ColumnIndex("size"), 64);
            probe.Set(key, probe.ColumnIndex("t1"), 1000 + i);
        }

        probe.Set(2, probe.ColumnIndex("t4"), 5000);
        RawStream cut = probe.Cut(1, 4, includeKey: true);
        Assert.Equal(3, cut.Count);
        Assert.Equal(["seq", "size", "t1", "t4"], cut.Columns.Select(c => c.Name));
        Assert.Equal([1L, 2L, 3L], cut.Find("seq")!.Values);
        Assert.Equal([false, true, false], cut.Find("t4")!.Valid!);

        Assert.Throws<ArgumentException>(() => probe.ColumnIndex("t2"));
    }

    [Fact]
    public void RowStore_RingOverwritesOldRows_AndNeverReturnsAStaleOne()
    {
        var store = new RowStore(Kinds.UdpEcho.FindStream("msg")!, Side.Agent, 16);
        store.SetRow(3, [30, 31]);
        Assert.True(store.TryGet(3, 0, out long t2));
        Assert.Equal(30, t2);

        store.SetRow(3 + 16, [40, 41]); // same slot, newer key
        Assert.False(store.Contains(3));
        Assert.False(store.TryGet(3, 0, out _));
        Assert.True(store.TryGet(19, 1, out long t3));
        Assert.Equal(41, t3);
    }

    // ---- 声明驱动的合并 ----

    private sealed class Source : IRecordSource
    {
        public Dictionary<string, RowStore> Stores { get; } = [];

        public RowStore? Records(uint streamId, string streamName) => Stores.GetValueOrDefault(streamName);
    }

    private static RawChunk EchoFragment(params (uint Seq, long T1, long? T4)[] rows)
    {
        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, long? t4) in rows)
        {
            msg.Add(seq, 64, t1, t4);
        }

        return new RawChunk(new RawChunkHeader { Kind = "udp_echo", Clock = new ChunkClock { ProbeOffsetUs = 1000 } }, [msg.Build()]);
    }

    [Fact]
    public void FragmentMerger_ProbeDriven_AddsTheAgentColumnsByKey_AndLeavesMissingRowsEmpty()
    {
        var agent = new RowStore(Kinds.UdpEcho.FindStream("msg")!, Side.Agent, 64);
        agent.SetRow(0, [2000, 2050]);
        agent.SetRow(2, [4000, 4040]);
        var source = new Source { Stores = { ["msg"] = agent } };
        RawChunk fragment = EchoFragment((0, 1000, 3000), (1, 2000, null), (2, 3000, 5000));

        Result<IReadOnlyList<RawStream>> merged = FragmentMerger.Merge(Kinds.UdpEcho, fragment, 1, source);
        RawStream msg = Assert.Single(merged.Value);
        Assert.Equal(["seq", "size", "t1", "t2", "t3", "t4"], msg.Columns.Select(c => c.Name));
        Assert.Equal([true, false, true], Enumerable.Range(0, 3).Select(i => msg.Find("t2")!.IsValid(i)));
        Assert.Equal(2050, msg.Find("t3")!.Values[0]);
        Assert.False(msg.Find("t4")!.IsValid(1));
    }

    [Fact]
    public void FragmentMerger_ProbeDriven_WithoutAnyAgentRecords_KeepsTheProbeRows()
    {
        RawChunk fragment = EchoFragment((0, 1000, 3000), (1, 2000, null));
        RawStream msg = Assert.Single(FragmentMerger.Merge(Kinds.UdpEcho, fragment, 1, new Source()).Value);
        Assert.Equal(2, msg.Count);
        Assert.False(msg.Find("t2")!.IsValid(0));
        Assert.Equal(3000, msg.Find("t4")!.Values[0]);
    }

    [Fact]
    public void FragmentMerger_RefusesAFragmentWithoutTheProbeColumns()
    {
        var bad = new StreamBuilder(new StreamSpec("msg", [Col.Seq("seq")]));
        bad.Add(0);
        var fragment = new RawChunk(new RawChunkHeader { Kind = "udp_echo" }, [bad.Build()]);
        Assert.False(FragmentMerger.Merge(Kinds.UdpEcho, fragment, 1, new Source()).IsSuccess);
    }

    [Fact]
    public void FragmentMerger_AgentDriven_TakesRowsFromTheAgent_AndMarksUnreceivedPacketsAsLost()
    {
        var agent = new RowStore(Kinds.UdpTrain.FindStream("pkt")!, Side.Agent, 64);
        for (uint seq = 0; seq < 5; seq++)
        {
            agent.SetRow(seq, [1000, 100 + seq]);
        }

        var pkt = new StreamBuilder(new StreamSpec("pkt", [Col.Seq("seq"), Col.TsProbe("t_recv")]));
        pkt.Add(1, 5000).Add(1, 9999).Add(3, 6000); // sequence 1 arrives twice (the first one counts); 0, 2 and 4 are lost
        var header = new RawChunkHeader { Kind = "udp_train" };
        header.Params[FragmentKeys.SeqFrom] = Json("0");
        header.Params[FragmentKeys.SeqTo] = Json("5");
        var source = new Source { Stores = { ["pkt"] = agent } };

        RawStream merged = Assert.Single(FragmentMerger.Merge(Kinds.UdpTrain, new RawChunk(header, [pkt.Build()]), 1, source).Value);
        Assert.Equal(5, merged.Count);
        Assert.Equal(["seq", "size", "t_send", "t_recv"], merged.Columns.Select(c => c.Name));
        RawColumn recv = merged.Find("t_recv")!;
        Assert.Equal([false, true, false, true, false], Enumerable.Range(0, 5).Select(recv.IsValid));
        Assert.Equal(5000, recv.Values[1]);
        Assert.Equal(103, merged.Find("t_send")!.Values[3]);

        // A one-way fragment that does not declare its sequence range is invalid
        Assert.False(FragmentMerger.Merge(Kinds.UdpTrain, new RawChunk(new RawChunkHeader { Kind = "udp_train" }, [pkt.Build()]), 1, source).IsSuccess);
    }

    [Fact]
    public void FragmentMerger_StreamsWithoutAgentColumns_PassThroughInDeclaredOrder()
    {
        var conn = new StreamBuilder(new StreamSpec("conn", [Col.U8("outcome"), Col.Seq("attempt"), Col.TsProbe("ts")]));
        conn.Add(0, 0, 100);
        var fragment = new RawChunk(new RawChunkHeader { Kind = "connect" }, [conn.Build()]);

        RawStream merged = Assert.Single(FragmentMerger.Merge(Kinds.Connect, fragment, 1, new Source()).Value);
        Assert.Equal(["attempt", "ts", "tcp_us", "tls_us", "outcome"], merged.Columns.Select(c => c.Name));
        Assert.False(merged.Find("tcp_us")!.IsValid(0));
    }

    // ---- 两侧各自记录完整行的流（内核采样）----

    private static RawStream InfoRows(StreamSpec spec, params (long Ts, long Side, long Retrans)[] rows)
    {
        var builder = new StreamBuilder(spec);
        foreach ((long ts, long side, long retrans) in rows)
        {
            builder.Add(ts, side, 40_000, 5_000, 10, retrans, 100, 1_000);
        }

        return builder.Build();
    }

    [Fact]
    public void FragmentMerger_UnionStreams_ConcatenateBothSidesRows_AndEachAgentRowIsDeliveredOnce()
    {
        StreamSpec info = Kinds.TcpEcho.FindStream("tcp_info")!;
        Assert.Equal(info.Columns.Count, info.ColumnsOf(Side.Agent).Count); // each side owns every column
        var agent = new RowStore(info, Side.Agent, 64);
        var recorder = new Source { Stores = { ["tcp_info"] = agent } };
        agent.SetRow(agent.Append(), [500, 1, 40_000, 5_000, 10, 7, 100, 1_000]);
        agent.SetRow(agent.Append(), [600, 1, 40_000, 5_000, 10, 8, 120, 1_100]);

        var msg = EchoFragment((0, 1000, 3000));
        var fragment = new RawChunk(msg.Header, [msg.Streams[0], InfoRows(info, (100, 0, 1))]);

        IReadOnlyList<RawStream> first = FragmentMerger.Merge(Kinds.TcpEcho, fragment, 1, recorder).Value;
        RawStream merged = first.Single(x => x.Name == "tcp_info");
        Assert.Equal(3, merged.Count);
        Assert.Equal([0L, 1L, 1L], merged.Find("side")!.Values);
        Assert.Equal([1L, 7L, 8L], merged.Find("retrans_total")!.Values);

        // The agent's rows were consumed: the next chunk carries only the probe's own rows
        IReadOnlyList<RawStream> second = FragmentMerger.Merge(Kinds.TcpEcho, fragment, 1, recorder).Value;
        Assert.Equal(1, second.Single(x => x.Name == "tcp_info").Count);
    }

    [Fact]
    public void FragmentMerger_UnionStreams_AgentRowsAreDeliveredEvenWhenTheProbeHasNone()
    {
        StreamSpec info = Kinds.TcpEcho.FindStream("tcp_info")!;
        var agent = new RowStore(info, Side.Agent, 64);
        agent.SetRow(agent.Append(), [500, 1, 40_000, 5_000, 10, 7, 100, 1_000]);

        // A probe on a platform without kernel sampling sends the echo rows only
        RawChunk fragment = EchoFragment((0, 1000, 3000));
        IReadOnlyList<RawStream> merged = FragmentMerger.Merge(Kinds.TcpEcho, fragment, 1, new Source { Stores = { ["tcp_info"] = agent } }).Value;
        RawStream rows = merged.Single(x => x.Name == "tcp_info");
        Assert.Equal(1, rows.Count);
        Assert.Equal(7, rows.Find("retrans_total")!.Values[0]);
    }

    [Fact]
    public void TcpInfo_OnLinux_ReadsTheKernelCountersOfALiveConnection()
    {
        Assert.SkipUnless(OperatingSystem.IsLinux(), "TCP_INFO is Linux only");
        using var listener = new System.Net.Sockets.TcpListener(System.Net.IPAddress.Loopback, 0);
        listener.Start();
        using var client = new System.Net.Sockets.TcpClient();
        client.Connect((System.Net.IPEndPoint)listener.LocalEndpoint);
        using System.Net.Sockets.TcpClient server = listener.AcceptTcpClient();
        client.GetStream().Write(new byte[1000]);
        Assert.True(TcpInfo.TrySample(client.Client, out TcpInfoSample sample));
        Assert.True(sample.Cwnd > 0);
        Assert.True(sample.SegsOut > 0);
    }

    // ---- 回显行的交出时机 ----

    [Fact]
    public void EchoRows_Tcp_WaitsForTheReplyInsteadOfCallingASeveralSecondStallALoss()
    {
        var rows = new EchoRows(Kinds.TcpEcho, 64, EchoRows.TcpGraceUs, settleOnReply: true);
        uint a = rows.Add(64, 0);
        uint b = rows.Add(64, 100_000);
        uint c = rows.Add(64, 200_000);
        rows.Reply(a, 1_000);
        rows.Reply(c, 9_000_000); // the third echo came back first in wall-clock terms, the second one is stuck behind a stall

        // 10 s later: row b has no reply yet. TCP does not lose data, so it is not given up on and it holds back the rows behind it
        RawStream first = rows.Take(10_000_000, all: false)!;
        Assert.Equal([0L], first.Find("seq")!.Values);
        Assert.Null(rows.Take(10_000_000, all: false));

        // The stall ends and the reply arrives with a multi-second RTT: it is recorded, not lost
        rows.Reply(b, 8_000_000);
        RawStream rest = rows.Take(10_000_000, all: false)!;
        Assert.Equal([1L, 2L], rest.Find("seq")!.Values);
        Assert.Equal([true, true], rest.Find("t4")!.Valid!);
        Assert.Equal(8_000_000, rest.Find("t4")!.Values[0]);

        // A connection that is really gone: after the long grace the row is delivered with the reply missing
        uint d = rows.Add(64, 20_000_000);
        Assert.Null(rows.Take(40_000_000, all: false));
        RawStream gone = rows.Take(51_000_000, all: false)!;
        Assert.Equal([(long)d], gone.Find("seq")!.Values);
        Assert.Equal([false], gone.Find("t4")!.Valid!);
    }

    [Fact]
    public void EchoRows_Udp_GivesUpOnAMissingReplyAfterTwoSeconds()
    {
        var rows = new EchoRows(Kinds.UdpEcho, 64, EchoRows.DatagramGraceUs, settleOnReply: false);
        rows.Add(64, 0);
        rows.Add(64, 1_000_000);
        RawStream cut = rows.Take(2_500_000, all: false)!;
        Assert.Equal([0L], cut.Find("seq")!.Values);
        Assert.Equal([false], cut.Find("t4")!.Valid!); // never replied: lost
        Assert.Equal([1L], rows.Take(0, all: true)!.Find("seq")!.Values);
    }

    // ---- 积木目录 ----

    [Fact]
    public void BlockCatalog_EveryKindHasADeclaration_AndEveryListenerKeepsTheRecordsItsStreamsNeed()
    {
        var env = new BlockEnv { Clock = new MonoClock(TimeProvider.System), Log = NullNodeLog.Instance, Session = () => null };
        foreach (BlockPair pair in BlockCatalog.All)
        {
            KindSpec kind = Kinds.Find(pair.Kind)!;
            Assert.NotNull(kind);
            Assert.Contains(pair.Kind, BlockCatalog.AgentKinds);

            if (pair.Listener is null)
            {
                Assert.True(pair.RunsAsTarget, $"{pair.Kind} has no listener, so it can only be a target task");
                continue;
            }

            Assert.Contains(pair.Kind, BlockCatalog.ProbeKinds);
            BlockSpec spec = new(kind, ParamResolver.Resolve(kind.Params, null).Values!, 1, 7443, new FragmentScope(null, null, null));
            IBlock listener = pair.Listener(env, spec);

            // 声明了连接键、Agent 一侧有列的流：监听方必须为它保留按键记录，否则合并不出完整的行
            foreach (StreamSpec stream in kind.Streams.Where(s => s.Key is not null && s.ColumnsOf(Side.Agent).Count > 0))
            {
                Assert.NotNull(listener.Records(stream.Name));
            }
        }
    }

    [Fact]
    public void BlockSpec_ConfigSignatureIgnoresTheStreamId_ButSignatureDoesNot()
    {
        BlockSpec Create(string json, uint stream = 1, int port = 7443) => BlockSpec.Create(
            "udp_echo",
            JsonDocument.Parse(json).RootElement.EnumerateObject().ToDictionary(p => p.Name, p => p.Value.Clone()),
            stream,
            port,
            new FragmentScope(null, null, null))!;

        BlockSpec baseline = Create("""{"rate_hz":10,"size":64}""");
        Assert.Equal(baseline.Signature, Create("""{"size":64,"rate_hz":10}""").Signature);
        Assert.NotEqual(baseline.Signature, Create("""{"rate_hz":20,"size":64}""").Signature);
        Assert.NotEqual(baseline.Signature, Create("""{"rate_hz":10,"size":64}""", stream: 2).Signature);
        Assert.Equal(baseline.ConfigSignature, Create("""{"rate_hz":10,"size":64}""", stream: 2).ConfigSignature);
        Assert.NotEqual(baseline.ConfigSignature, Create("""{"rate_hz":10,"size":64}""", port: 9000).ConfigSignature);
        Assert.Null(BlockSpec.Create("udp_echo", new Dictionary<string, JsonElement> { ["rate_hz"] = Json("100000") }, 1, 7443, new FragmentScope(null, null, null)));
    }

    // ---- 起止时间 ----

    [Fact]
    public void ChunkFinalizer_TakesStartAndDurationFromTheDataTimestamps_NotFromTheDeliveryTime()
    {
        var agent = new RowStore(Kinds.UdpEcho.FindStream("msg")!, Side.Agent, 64);
        agent.SetRow(0, [2000, 2050]);
        RawChunk fragment = EchoFragment((0, 1000, 3000));
        IReadOnlyList<RawStream> streams = FragmentMerger.Merge(Kinds.UdpEcho, fragment, 1, new Source { Stores = { ["msg"] = agent } }).Value;

        var anchor = new ClockAnchor(AgentMonoUs: 1_000_000, UtcMs: 1_700_000_000_000);
        RawChunk chunk = ChunkFinalizer.Finalize(Kinds.UdpEcho, fragment, streams, new NodeId("agt_x"), anchor, new Dictionary<string, JsonElement> { ["port"] = Json("7443") }, 7443, new ChunkId(Guid.NewGuid()));

        // t1 = 1000 us (Probe clock) + offset 1000 = 2000 us on the Agent clock, which is 998 ms before the anchor
        Assert.Equal(1_700_000_000_000 - 998, chunk.Header.StartUtcMs);
        Assert.Equal(7443, chunk.Header.Port);
        Assert.False(chunk.Header.Partial);
        Assert.Equal(1_000_000, chunk.Header.Clock.AgentMonoUs);
    }
}