using System.Collections.Concurrent;
using Fleet.Core;
using Fleet.Rpc;
using MoShiro.Measurement;
using MoShiro.Protocol;

namespace MoShiro.Agent;

/// <summary>
/// 按需测试的编排（<c>moshiro.test.run</c>）：Agent 是时间的主人。每个阶段：
/// 登记数据面的流 → <c>moshiro.test.phase</c>（Probe 准备好自己一侧的角色）→ 运行（单向数据流由 Agent 发包，回显类等待 Probe 发包）
/// → <c>moshiro.test.phase_end</c>（Probe 收尾，回复它交出的列块个数）→ 等所有列块合并上传 → <c>phase_done</c>。
/// 同一时间只能运行 <c>max_concurrent_tests</c> 个测试；Probe 掉线或被取消时测试中止，已经上传的阶段数据保留。
/// </summary>
public sealed class TestRunner
{
    private static readonly TimeSpan s_phaseSlack = TimeSpan.FromSeconds(15);
    private static readonly TimeSpan s_fragmentWait = TimeSpan.FromSeconds(10);

    private readonly ConcurrentDictionary<string, RunningTest> _running = new();
    private readonly TaskManager _tasks;
    private readonly SessionRegistry _sessions;
    private readonly StreamRegistry _streams;
    private readonly BlockEnv _env;
    private readonly BackendUplink _uplink;
    private readonly int _maxConcurrent;
    private readonly Func<string, IReadOnlyList<int>> _ports;
    private readonly TimeProvider _time;
    private readonly INodeLog _log;

    public TestRunner(TaskManager tasks, SessionRegistry sessions, StreamRegistry streams, BlockEnv env, BackendUplink uplink, int maxConcurrent, Func<string, IReadOnlyList<int>> ports, TimeProvider time, INodeLog log)
    {
        _tasks = tasks;
        _sessions = sessions;
        _streams = streams;
        _env = env;
        _uplink = uplink;
        _maxConcurrent = Math.Max(1, maxConcurrent);
        _ports = ports;
        _time = time;
        _log = log;
    }

    public int Running => _running.Count;

    /// <summary>校验并接受一个测试（立即返回），实际运行在后台进行，状态通过 <c>moshiro.test.status</c> 回报。</summary>
    public Result<TestRunResult> Start(TestRunParams run)
    {
        if (!IpVersions.IsValid(run.IpVersion))
        {
            return Result<TestRunResult>.Fail(MoShiroErrors.ValidationFailed, "ip_version must be ipv4 or ipv6");
        }

        if (_sessions.Get(run.ProbeId.Value, run.IpVersion) is not { } probe)
        {
            return Result<TestRunResult>.Fail(MoShiroErrors.ProbeOffline, $"probe {run.ProbeId.Value} is not connected to this agent");
        }

        if (run.Phases.Count == 0)
        {
            return Result<TestRunResult>.Fail(MoShiroErrors.ValidationFailed, "a test needs at least one phase");
        }

        var resolved = new List<(TestPhase Phase, TaskDecision Decision)>();
        foreach (TestPhase phase in run.Phases)
        {
            WithDefaultPort(phase);
            TaskDecision decision = _tasks.Guard.CheckPhase(phase.Kind, phase.Params);
            if (!decision.Accepted)
            {
                return Result<TestRunResult>.Fail(
                    decision.Reason == RejectReasons.Unsupported ? FleetErrors.Unsupported : MoShiroErrors.BudgetExceeded,
                    $"phase {phase.Kind} refused: {decision.Reason}");
            }

            string? direction = decision.Resolved!.Values.ContainsKey("direction") ? decision.Resolved.GetString("direction") : null;
            if (BlockCatalog.Find(phase.Kind) is not { Listener: not null } pair || !pair.Supports(direction))
            {
                return Result<TestRunResult>.Fail(FleetErrors.Unsupported, $"{phase.Kind} does not support direction={direction}");
            }
        }

        foreach (TestPhase phase in run.Phases)
        {
            resolved.Add((phase, _tasks.Guard.CheckPhase(phase.Kind, phase.Params)));
        }

        var running = new RunningTest(probe);
        if (_running.Count >= _maxConcurrent || !_running.TryAdd(run.TestId.Value, running))
        {
            running.Dispose();
            return Result<TestRunResult>.Fail(MoShiroErrors.TestBusy, "this agent is already running a test");
        }

        _ = Task.Run(() => RunAsync(run.TestId, resolved, running), CancellationToken.None);
        return Result<TestRunResult>.Ok(new TestRunResult { Accepted = true });
    }

    /// <summary>阶段没有指定端口时用本 Agent 自己监听的第一个端口（Backend 的模板不知道每个 Agent 的实际端口）。</summary>
    private void WithDefaultPort(TestPhase phase)
    {
        KindSpec? kind = Kinds.Find(phase.Kind);
        if (kind is null || !kind.Params.Any(p => p.Name == "port"))
        {
            return;
        }

        // 按需测试的数据端口由 Agent 决定：Backend 给的端口如果不是本机监听的端口（模板里是出厂默认值，每个 Agent 可以改，
        // 而 Backend 并不知道 TCP 数据端口），就用本机监听的第一个端口。块头里记录的是实际用的端口。
        IReadOnlyList<int> ports = _ports(kind.Blocks.Transport == "udp" ? "udp" : "tcp");
        bool listening = phase.Params.TryGetValue("port", out System.Text.Json.JsonElement requested) && requested.TryGetInt32(out int port) && ports.Contains(port);
        if (!listening)
        {
            phase.Params["port"] = System.Text.Json.JsonDocument.Parse(ports[0].ToString(System.Globalization.CultureInfo.InvariantCulture)).RootElement;
        }
    }

    public bool Cancel(TestId test)
    {
        if (_running.TryGetValue(test.Value, out RunningTest? running))
        {
            running.Cancel.Cancel();
            return true;
        }

        return false;
    }

    /// <summary>某个 Probe 掉线：它上面的测试中止，状态是 failed / probe_offline；别的 Probe 上的测试不受影响。</summary>
    public void ProbeLost(ProbeSession probe)
    {
        foreach (RunningTest running in _running.Values)
        {
            if (ReferenceEquals(running.Probe, probe))
            {
                running.Lost = true;
                running.Cancel.Cancel();
            }
        }
    }

    private async Task RunAsync(TestId test, List<(TestPhase Phase, TaskDecision Decision)> phases, RunningTest running)
    {
        CancellationToken ct = running.Cancel.Token;
        string finalStatus = TestStatuses.Succeeded;
        string? errorCode = null;
        int? lastPhase = null;
        try
        {
            ProbeSession probe = running.Probe;
            for (int index = 0; index < phases.Count; index++)
            {
                lastPhase = index;
                await _uplink.SendTestStatusAsync(test, TestStatuses.Running, index);
                await RunPhaseAsync(probe, test, index, phases[index].Phase, phases[index].Decision, ct);
                await _uplink.SendTestStatusAsync(test, TestStatuses.PhaseDone, index);
            }
        }
        catch (OperationCanceledException)
        {
            finalStatus = running.Lost ? TestStatuses.Failed : TestStatuses.Aborted;
            errorCode = running.Lost ? MoShiroErrors.ProbeOffline.Value : null;
            await TryAbortProbeAsync(running.Probe, test);
        }
        catch (TestFailure failure)
        {
            finalStatus = TestStatuses.Failed;
            errorCode = failure.Code;
            await TryAbortProbeAsync(running.Probe, test);
        }
        finally
        {
            _running.TryRemove(test.Value, out _);
            running.Dispose();
        }

        _log.Info("test finished", ("test", test.Value), ("status", finalStatus));
        await _uplink.SendTestStatusAsync(test, finalStatus, lastPhase, errorCode);
    }

    private async Task RunPhaseAsync(ProbeSession probe, TestId test, int index, TestPhase phase, TaskDecision decision, CancellationToken ct)
    {
        ParamValues values = decision.Resolved!;
        var stream = new StreamInfo(_streams.Allocate(), phase.Kind, values) { TestId = test, Phase = index, Owner = probe };
        int durationSeconds = phase.DurationSeconds > 0 ? phase.DurationSeconds : (int)(values.Values.ContainsKey("duration_s") ? values.GetInt("duration_s") : 10);
        var blockSpec = new BlockSpec(decision.Kind!, values, stream.Id, stream.Port, new FragmentScope(null, test, index));
        stream.Running = RunningBlock.Start(BlockCatalog.Find(phase.Kind)!.Listener!(_env.WithSession(probe), blockSpec), blockSpec.Signature);
        _streams.Add(stream);
        try
        {
            RpcSession<PeerContext> rpc = probe.Rpc ?? throw new TestFailure(MoShiroErrors.ProbeOffline.Value);
            var start = new TestPhaseStartParams
            {
                TestId = test,
                Phase = index,
                Kind = phase.Kind,
                Params = new Dictionary<string, System.Text.Json.JsonElement>(values.Values),
                StreamId = stream.Id,
                DurationSeconds = durationSeconds,
                Port = stream.Port,
            };
            Result<TestPhaseStartResult> accepted = await rpc.RequestAsync(MoShiroMethods.TestPhase, start, MoShiroJsonContext.Default.TestPhaseStartParams, MoShiroJsonContext.Default.TestPhaseStartResult, ct);
            if (!accepted.TryGet(out TestPhaseStartResult? started, out FleetError? error) || !started.Accepted)
            {
                throw new TestFailure(error?.Code.Value ?? FleetErrors.Unsupported.Value);
            }

            // 单向数据流的发送端自己运行到结束（或失败）；其他积木是反应式的，等够时长即可
            if (phase.Kind == "udp_train")
            {
                try
                {
                    await stream.Running.Completion.WaitAsync(ct);
                }
                catch (BlockFaultException)
                {
                    throw new TestFailure(MoShiroErrors.ProbeOffline.Value);
                }
            }
            else
            {
                await Task.Delay(TimeSpan.FromSeconds(durationSeconds), _time, ct);
            }

            long sent = stream.Running.Block.Produced;
            var end = new TestPhaseEndParams { TestId = test, Phase = index, StreamId = stream.Id, SentCount = sent };
            Result<TestPhaseEndResult> ended = await rpc.RequestAsync(MoShiroMethods.TestPhaseEnd, end, MoShiroJsonContext.Default.TestPhaseEndParams, MoShiroJsonContext.Default.TestPhaseEndResult, ct, s_phaseSlack);
            if (!ended.TryGet(out TestPhaseEndResult? finished, out FleetError? endError))
            {
                throw new TestFailure(endError.Code.Value);
            }

            // 列块是经 Probe 的发件箱至少一次送达的：全部合并上传之后才进入下一阶段（超时就带着已有的数据继续）
            DateTimeOffset deadline = _time.GetUtcNow() + s_fragmentWait;
            while (stream.FragmentsDone < finished.Fragments && _time.GetUtcNow() < deadline)
            {
                await Task.Delay(TimeSpan.FromMilliseconds(20), _time, ct);
            }

            if (stream.FragmentsDone < finished.Fragments)
            {
                _log.Warn("fragments missing", ("test", test.Value), ("phase", index), ("expected", finished.Fragments), ("got", stream.FragmentsDone));
            }
        }
        finally
        {
            _streams.Remove(stream.Id);
            await stream.Running.StopAsync();
            probe.ForgetStream(stream.Id);
        }
    }

    private async Task TryAbortProbeAsync(ProbeSession probe, TestId test)
    {
        if (probe.Rpc is not { } rpc)
        {
            return;
        }

        try
        {
            await rpc.NotifyAsync(MoShiroMethods.TestAbort, new TestAbortParams { TestId = test }, MoShiroJsonContext.Default.TestAbortParams, CancellationToken.None);
        }
        catch (Exception ex) when (ex is ObjectDisposedException or InvalidOperationException or IOException)
        {
            // 会话已经结束
        }
    }

    /// <summary>一个正在运行的测试：它在哪个 Probe 上、怎样取消、是不是因为 Probe 掉线而中止。</summary>
    private sealed class RunningTest : IDisposable
    {
        public RunningTest(ProbeSession probe)
        {
            Probe = probe;
        }

        public ProbeSession Probe { get; }

        public CancellationTokenSource Cancel { get; } = new();

        public volatile bool Lost;

        public void Dispose() => Cancel.Dispose();
    }

    private sealed class TestFailure : Exception
    {
        public TestFailure(string code)
            : base(code)
        {
            Code = code;
        }

        public string Code { get; }
    }
}