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

namespace MoShiro.Probe;

/// <summary>所有连接共用的东西。</summary>
public sealed record ConnectionEnv(NodeIdentity Identity, string ProbeId, ProbeLimits Limits, AgentDirectory Directory, MonoClock Clock, TimeProvider Time, INodeLog Log);

/// <summary>
/// 与一个 Agent 的点对点会话：认证握手 → 同步时钟 → 按 Agent 下发的 <c>line_task_set</c> 运行角色 →
/// 每 10 秒把已完成的行切成列块放进发件箱 → 响应按需测试的阶段。会话断开时所有角色停止（凭据属于会话），按退避重连。
/// 会话内的所有状态（套接字、发件箱、角色）随会话一起创建和释放。
/// </summary>
public sealed class AgentConnection
{
    private static readonly TimeSpan s_flushInterval = TimeSpan.FromSeconds(ProtocolLimits.ChunkSeconds);
    private static readonly TimeSpan s_clockInterval = TimeSpan.FromSeconds(30);
    private static readonly RpcRouter<ConnectionSession> s_router = BuildRouter();

    /// <summary>对等面入站方法的路由器（只读，测试用）。</summary>
    public static RpcRouter<ConnectionSession> Router => s_router;

    private readonly ConnectionEnv _env;
    private volatile ConnectionSession? _session;

    public AgentConnection(KnownAgent endpoint, ConnectionEnv env)
    {
        Endpoint = endpoint;
        _env = env;
    }

    public KnownAgent Endpoint { get; }

    public bool IsConnected => _session is not null;

    public int RoleCount => _session?.RoleCount ?? 0;

    public long ClockOffsetUs => _session?.Clock.OffsetUs ?? 0;

    public async Task RunAsync(CancellationToken ct)
    {
        var backoff = new Backoff(TimeSpan.FromSeconds(1), TimeSpan.FromSeconds(20));
        while (!ct.IsCancellationRequested)
        {
            try
            {
                string reason = await RunOnceAsync(backoff, ct);
                _env.Log.Warn("agent session ended", ("agent", Endpoint.AgentId), ("err", reason));
            }
            catch (OperationCanceledException) when (ct.IsCancellationRequested)
            {
                return;
            }
            catch (Exception ex) when (ex is IOException or System.Net.Sockets.SocketException or InvalidOperationException or System.Security.Authentication.AuthenticationException or OperationCanceledException or FormatException)
            {
                _env.Log.Warn("agent connection failed", ("agent", Endpoint.AgentId), ("err", ex.Message));
            }

            try
            {
                await Task.Delay(backoff.Next(), _env.Time, ct);
            }
            catch (OperationCanceledException)
            {
                return;
            }
        }
    }

    private async Task<string> RunOnceAsync(Backoff backoff, CancellationToken ct)
    {
        AddressFamily family = Endpoint.IpVersion == IpVersions.V6 ? AddressFamily.InterNetworkV6 : AddressFamily.InterNetwork;
        using var resolveTimeout = new CancellationTokenSource(TimeSpan.FromSeconds(10), _env.Time);
        using CancellationTokenSource resolveCt = CancellationTokenSource.CreateLinkedTokenSource(ct, resolveTimeout.Token);
        IPAddress[] addresses = IPAddress.TryParse(Endpoint.Host, out IPAddress? literal)
            ? [literal]
            : await Dns.GetHostAddressesAsync(Endpoint.Host, family, resolveCt.Token);
        IPAddress? address = addresses.FirstOrDefault(a => a.AddressFamily == family && !a.IsIPv4MappedToIPv6);
        if (address is null)
        {
            throw new SocketException((int)SocketError.HostNotFound);
        }

        (PeerSessionInfo info, TlsFrameTransport transport) = await PeerConnector.ConnectAsync(
            address.ToString(), Endpoint.TcpPort, Convert.FromBase64String(Endpoint.Spki), _env.Identity, _env.ProbeId, ProbeApp.Version, 8 << 20, _env.Time, ct);
        var remote = (IPEndPoint)info.Remote!;
        var data = new DataSession(info.Credentials.SessionIdValue, info.Credentials.Key);
        using var udp = new UdpMux(new SingleSession(data), null, _env.Clock, new DataPlaneCounters(), _env.Log, remote.AddressFamily == AddressFamily.InterNetworkV6 ? IPAddress.IPv6Any : IPAddress.Any, []);
        udp.Bind();
        var outbox = new Outbox(8 << 20);
        var clock = new ClockEstimator();
        var shipper = new FragmentShipper(outbox, new NodeId(Endpoint.AgentId), new ProbeId(_env.ProbeId), clock, _env.Time);
        var blockEnv = new BlockEnv
        {
            Clock = _env.Clock,
            Log = _env.Log,
            Session = () => data,
            Udp = udp,
            PeerAddress = remote.Address,
            PeerSpkiSha256 = Convert.FromBase64String(Endpoint.Spki),
            Fragments = shipper,
            MaxPacketsPerSecond = _env.Limits.MaxTotalPps,
            MaxBulkMbps = _env.Limits.MaxTestMbps,
        };
        var session = new ConnectionSession(blockEnv, clock, outbox, _env);
        var rpc = new RpcSession<ConnectionSession>(transport, s_router, session, _env.Time);
        rpc.OnBinary = (_, _) => ValueTask.CompletedTask;

        using CancellationTokenSource cts = CancellationTokenSource.CreateLinkedTokenSource(ct);
        Task run = rpc.RunAsync(cts.Token);
        Task receive = udp.RunAsync(cts.Token);
        Task pump = PumpAsync(rpc, outbox, cts.Token);
        try
        {
            // 先同步时钟（块头里要带偏移）；Agent 随后推送线路任务集
            await SyncClockAsync(rpc, clock, rounds: 5, cts.Token);
            backoff.Reset();
            _session = session;
            _env.Log.Info("agent session up", ("agent", Endpoint.AgentId), ("offset_us", clock.OffsetUs));
            Task maintenance = MaintainAsync(rpc, session, clock, cts.Token);
            Task finished = await Task.WhenAny(run, receive, pump, maintenance);
            await cts.CancelAsync();
            await Task.WhenAll(Quiet(run), Quiet(receive), Quiet(pump), Quiet(maintenance));
            return finished == run ? "closed by the agent" : "session task failed";
        }
        finally
        {
            _session = null;
            await cts.CancelAsync();
            await session.StopAllAsync();
            await Quiet(run);
            await Quiet(receive);
            await Quiet(pump);
            await rpc.DisposeAsync();
        }
    }

    private static async Task Quiet(Task task)
    {
        try
        {
            await task;
        }
        catch (Exception ex) when (ex is OperationCanceledException or IOException or ObjectDisposedException or System.Net.Sockets.SocketException or InvalidOperationException)
        {
            // 会话结束时的正常噪声
        }
    }

    private static async Task PumpAsync(RpcSession<ConnectionSession> rpc, Outbox outbox, CancellationToken ct)
    {
        while (!ct.IsCancellationRequested)
        {
            byte[] frame = await outbox.TakeNextAsync(2 << 20, ct);
            await rpc.SendBinaryAsync(frame, ct);
        }
    }

    private async Task SyncClockAsync(RpcSession<ConnectionSession> rpc, ClockEstimator clock, int rounds, CancellationToken ct)
    {
        for (int i = 0; i < rounds; i++)
        {
            long sent = _env.Clock.NowUs();
            Result<ClockResult> reply = await rpc.RequestAsync(MoShiroMethods.Clock, new ClockParams { ProbeMonoUs = sent }, MoShiroJsonContext.Default.ClockParams, MoShiroJsonContext.Default.ClockResult, ct);
            if (reply.TryGet(out ClockResult? value, out _))
            {
                clock.Add(sent, value.AgentMonoUs, _env.Clock.NowUs());
            }
        }
    }

    /// <summary>切块（每 10 秒）和时钟同步（每 30 秒）。</summary>
    private async Task MaintainAsync(RpcSession<ConnectionSession> rpc, ConnectionSession session, ClockEstimator clock, CancellationToken ct)
    {
        using var flush = new PeriodicTimer(s_flushInterval, _env.Time);
        long lastClockUs = _env.Clock.NowUs();
        while (await flush.WaitForNextTickAsync(ct))
        {
            session.FlushLineBlocks();
            if (_env.Clock.NowUs() - lastClockUs >= s_clockInterval.TotalMicroseconds)
            {
                lastClockUs = _env.Clock.NowUs();
                await SyncClockAsync(rpc, clock, rounds: 1, ct);
            }
        }
    }

    private static RpcRouter<ConnectionSession> BuildRouter() =>
        new RpcRouter<ConnectionSession>(RpcSide.Dialer, MoShiroMethodTable.Table)
            .On(
                FleetMethods.StateSync,
                FleetRpcContext.Default.StateSyncParams,
                FleetRpcContext.Default.StateSyncResult,
                async (session, sync, _) =>
                {
                    var result = new StateSyncResult { AppliedVersion = sync.Version };
                    if (sync.Type == LineTaskSet.StateType)
                    {
                        LineTaskSet? set = sync.Doc.Deserialize(MoShiroJsonContext.Default.LineTaskSet);
                        if (set is null)
                        {
                            return Result<StateSyncResult>.Fail(FleetErrors.InvalidParams, "line_task_set could not be parsed");
                        }

                        result.Rejected.AddRange(await session.ApplyAsync(set));
                    }

                    return Result<StateSyncResult>.Ok(result);
                })
            .On(
                MoShiroMethods.TestPhase,
                MoShiroJsonContext.Default.TestPhaseStartParams,
                MoShiroJsonContext.Default.TestPhaseStartResult,
                (session, start, _) => ValueTask.FromResult(Result<TestPhaseStartResult>.Ok(new TestPhaseStartResult { Accepted = session.StartPhase(start) })))
            .On(
                MoShiroMethods.TestPhaseEnd,
                MoShiroJsonContext.Default.TestPhaseEndParams,
                MoShiroJsonContext.Default.TestPhaseEndResult,
                async (session, end, ct) => Result<TestPhaseEndResult>.Ok(new TestPhaseEndResult { Fragments = await session.EndPhaseAsync(end, ct) }))
            .OnNotify(
                MoShiroMethods.TestAbort,
                MoShiroJsonContext.Default.TestAbortParams,
                async (session, abort, _) => await session.AbortAsync(abort.TestId))
            .OnNotify(
                FleetMethods.OutboxAck,
                FleetRpcContext.Default.OutboxAckParams,
                (session, ack, _) =>
                {
                    session.Outbox.Acknowledge(ack.Ids.Concat(ack.Duplicates).Select(Guid.Parse));
                    return ValueTask.CompletedTask;
                });
}

/// <summary>一次会话内的运行状态：线路任务的积木、按需测试的阶段。方法由路由器在单个入站队列里顺序调用。</summary>
public sealed class ConnectionSession
{
    private const int MaxBlocks = 16;

    private readonly object _gate = new();
    private readonly Dictionary<string, RunningBlock> _lineBlocks = [];
    private readonly ConcurrentDictionary<uint, PhaseState> _phases = new();

    public ConnectionSession(BlockEnv blockEnv, ClockEstimator clock, Outbox outbox, ConnectionEnv env)
    {
        BlockEnv = blockEnv;
        Clock = clock;
        Outbox = outbox;
        Env = env;
    }

    public BlockEnv BlockEnv { get; }

    public ClockEstimator Clock { get; }

    public Outbox Outbox { get; }

    public ConnectionEnv Env { get; }

    public int RoleCount
    {
        get
        {
            lock (_gate)
            {
                return _lineBlocks.Count;
            }
        }
    }

    /// <summary>调和线路任务：参数没有变化的积木继续运行，其余启动或停止。返回被拒绝的任务（不认识的 Kind、参数非法、超过上限）。</summary>
    public async Task<List<RejectedItem>> ApplyAsync(LineTaskSet set)
    {
        var rejected = new List<RejectedItem>();
        var wanted = new Dictionary<string, BlockSpec>();
        foreach (LineTask task in set.Tasks)
        {
            BlockSpec? spec = BlockCatalog.ProbeKinds.Contains(task.Kind)
                ? BlockSpec.Create(task.Kind, task.Params, task.StreamId, task.Port, new FragmentScope(task.Id, null, null))
                : null;
            if (spec is null)
            {
                rejected.Add(new RejectedItem { Id = task.Id.Value, Reason = BlockCatalog.ProbeKinds.Contains(task.Kind) ? "invalid_params" : "unsupported" });
            }
            else if (wanted.Count >= MaxBlocks)
            {
                rejected.Add(new RejectedItem { Id = task.Id.Value, Reason = "budget_exceeded" });
            }
            else
            {
                wanted[task.Id.Value] = spec;
            }
        }

        var stopped = new List<RunningBlock>();
        lock (_gate)
        {
            foreach ((string id, RunningBlock running) in _lineBlocks.ToList())
            {
                if (!wanted.TryGetValue(id, out BlockSpec? spec) || spec.Signature != running.Signature)
                {
                    stopped.Add(running);
                    _lineBlocks.Remove(id);
                }
            }

            foreach ((string id, BlockSpec spec) in wanted)
            {
                if (!_lineBlocks.ContainsKey(id))
                {
                    _lineBlocks[id] = RunningBlock.Start(BlockCatalog.Find(spec.KindId)!.Dialer(BlockEnv, spec), spec.Signature);
                }
            }
        }

        // 被停止的积木把已经稳定的行交出去（刚发出、还在等回复的行不交：它们不是丢失，只是没有等到）
        foreach (RunningBlock running in stopped)
        {
            await running.StopAsync();
            running.Block.Flush(FlushMode.Settled);
        }

        return rejected;
    }

    public void FlushLineBlocks()
    {
        List<RunningBlock> blocks;
        lock (_gate)
        {
            blocks = [.. _lineBlocks.Values];
        }

        foreach (RunningBlock running in blocks)
        {
            running.Block.Flush(FlushMode.Settled);
        }
    }

    /// <summary>按需测试的一个阶段：起对应的拨号方积木（单向数据流是接收端）。超过本地上限（或不认识）则拒绝。</summary>
    public bool StartPhase(TestPhaseStartParams start)
    {
        string? direction = start.Params.TryGetValue("direction", out System.Text.Json.JsonElement d) ? d.GetString() : null;
        if (_phases.ContainsKey(start.StreamId) || !BlockCatalog.ProbeKinds.Contains(start.Kind) || BlockCatalog.Find(start.Kind)!.Supports(direction) is false)
        {
            return false;
        }

        if (start.Kind == "udp_train")
        {
            double mbps = start.Params["mbps"].GetDouble();
            long size = start.Params["size"].GetInt64();
            long pps = (long)(mbps * 1_000_000 / 8 / Math.Max(1, size));
            if (mbps > Env.Limits.MaxTestMbps || pps > Env.Limits.MaxTotalPps)
            {
                return false;
            }
        }

        BlockSpec? spec = BlockSpec.Create(start.Kind, start.Params, start.StreamId, start.Port, new FragmentScope(null, start.TestId, start.Phase));
        if (spec is null)
        {
            return false;
        }

        _phases[start.StreamId] = new PhaseState(start.TestId, start.Kind, RunningBlock.Start(BlockCatalog.Find(start.Kind)!.Dialer(BlockEnv, spec), spec.Signature));
        return true;
    }

    /// <summary>阶段结束：收尾、交出列块，返回块数。</summary>
    public async Task<int> EndPhaseAsync(TestPhaseEndParams end, CancellationToken ct)
    {
        if (!_phases.TryRemove(end.StreamId, out PhaseState? phase))
        {
            return 0;
        }

        // 最后一批包（单向数据流）、最后几个回显的回复还在途中
        await Task.Delay(TimeSpan.FromMilliseconds(phase.Kind == "udp_train" ? 200 : 300), BlockEnv.Time, ct);
        await phase.Running.StopAsync();
        return phase.Running.Block.Flush(FlushMode.All, end.SentCount);
    }

    public async Task AbortAsync(TestId test)
    {
        foreach ((uint stream, PhaseState phase) in _phases.ToList())
        {
            if (phase.Test == test && _phases.TryRemove(stream, out _))
            {
                await phase.Running.StopAsync();
            }
        }
    }

    public async Task StopAllAsync()
    {
        List<RunningBlock> blocks;
        lock (_gate)
        {
            blocks = [.. _lineBlocks.Values];
            _lineBlocks.Clear();
        }

        blocks.AddRange(_phases.Values.Select(p => p.Running));
        _phases.Clear();
        foreach (RunningBlock running in blocks)
        {
            await running.StopAsync();
        }
    }

    private sealed record PhaseState(TestId Test, string Kind, RunningBlock Running);
}