using System.Text.Json;
using Fleet.Hosting;
using Fleet.Core;
using Fleet.Rpc;
using MoShiro.Measurement;
using MoShiro.Protocol;

namespace MoShiro.Agent;

/// <summary>
/// 期望任务集（<c>task_set</c>）的调和器：每个任务先经过 <see cref="BudgetGuard"/>（参数校验 + 本地硬上限）。
/// 通过的线路任务<b>为每个拨号方会话各登记一组流</b>（流编号和监听方积木都属于会话，互不干扰），并汇总成 <c>line_task_set</c> 下发给这个拨号方；
/// 通过的目标任务在本地运行，与拨号方无关。被拒绝的任务带着原因（<c>rejected[]</c>）回报给 Backend，其余任务照常运行。
/// 每次调用都应用完整的快照，所以是幂等的。预算（码率、包速率）按单个拨号方会话计算，会话总数由 <c>PeerListenerOptions.MaxSessions</c> 限制。
/// </summary>
public sealed class TaskManager
{
    // 线路任务 = 有监听方积木、并且不是只用于按需测试的 Kind（quic_stream、tcp_transfer 会占满链路）；目标任务 = 可以由 Agent 自己运行的 Kind
    private static readonly HashSet<string> s_lineKinds = [.. BlockCatalog.All.Where(p => p.Listener is not null && !p.OnDemandOnly).Select(p => p.Kind)];
    private static readonly HashSet<string> s_targetKinds = [.. BlockCatalog.All.Where(p => p.RunsAsTarget).Select(p => p.Kind)];

    private readonly object _gate = new();
    private readonly Func<BudgetGuard> _guardFactory;
    private readonly StreamRegistry _streams;
    private readonly BlockEnv _env;
    private readonly TargetRunners _targets;
    private readonly Func<IReadOnlyList<int>> _quicPorts;
    private readonly Func<IReadOnlyList<int>> _tcpPorts;
    private readonly TrustAnchor _trust;
    private readonly bool _acceptLineTasks;
    private readonly Dictionary<DialerSession, Dictionary<string, LineEntry>> _sessionLines = [];
    private volatile Dictionary<string, byte[]> _peers = new(StringComparer.Ordinal);
    private Dictionary<string, (TaskSpec Task, TaskDecision Decision)> _lineSpecs = [];
    private long _version;
    private BudgetGuard? _guard;

    public TaskManager(Func<BudgetGuard> guardFactory, StreamRegistry streams, BlockEnv env, TargetRunners targets, Func<IReadOnlyList<int>> quicPorts, Func<IReadOnlyList<int>> tcpPorts, TrustAnchor? trust = null, bool acceptLineTasks = true)
    {
        _trust = trust ?? TrustAnchor.None;
        _acceptLineTasks = acceptLineTasks;
        _guardFactory = guardFactory;
        _streams = streams;
        _env = env;
        _targets = targets;
        _quicPorts = quicPorts;
        _tcpPorts = tcpPorts;
    }

    public BudgetGuard Guard => _guard ??= _guardFactory();

    /// <summary>某个拨号方的线路任务集变化（版本递增）：由 <see cref="PeerHub"/> 推送给这个拨号方。</summary>
    public Action<DialerSession, LineTaskSet>? OnLineTasksChanged { get; set; }

    /// <summary>允许连接的拨号方变了（期望状态里的 peers）：已经不在名单里的会话由 <see cref="PeerHub"/> 关闭。</summary>
    public Action<IReadOnlySet<string>>? OnPeersChanged { get; set; }

    /// <summary>点对点握手时查找拨号方的公钥（SPKI）。</summary>
    public byte[]? LookupPeerKey(string peerId) => _peers.TryGetValue(peerId, out byte[]? key) ? key : null;

    /// <summary>名册里全部拨号方的公钥：敲门关卡用它们派生配对密钥。</summary>
    public IReadOnlyList<byte[]> PeerKeys() => [.. _peers.Values];

    /// <summary>一个新的拨号方会话开始：为它登记全部线路任务的流并启动监听方积木，返回要下发给它的线路任务集。</summary>
    public LineTaskSet OpenSession(DialerSession dialer)
    {
        lock (_gate)
        {
            var entries = new Dictionary<string, LineEntry>();
            foreach ((string id, (TaskSpec task, TaskDecision decision)) in _lineSpecs)
            {
                LineEntry entry = BuildLineEntry(task, decision, dialer);
                Start(entry);
                entries[id] = entry;
            }

            _sessionLines[dialer] = entries;
            return SetFor(entries);
        }
    }

    /// <summary>拨号方 会话结束：释放它的所有流和积木（记录随之丢弃，下一个会话从零开始）。</summary>
    public async Task CloseSessionAsync(DialerSession dialer)
    {
        List<LineEntry> entries;
        lock (_gate)
        {
            if (!_sessionLines.Remove(dialer, out Dictionary<string, LineEntry>? lines))
            {
                return;
            }

            entries = [.. lines.Values];
        }

        foreach (LineEntry entry in entries)
        {
            await Retire(entry);
        }
    }

    public async ValueTask<StateSyncResult> ApplyAsync(long version, TaskSet set, CancellationToken ct)
    {
        _ = ct;
        var result = new StateSyncResult { AppliedVersion = version };

        // 固定了操作员创世的节点：面板带来的新版本操作员集合，从固定的创世验过才采用（名册本身由面板决定）
        _trust.Update(set.TrustChain);
        var peers = new Dictionary<string, byte[]>(StringComparer.Ordinal);
        foreach (PeerSpec peer in set.Peers)
        {
            try
            {
                peers[peer.NodeId.Value] = Convert.FromBase64String(peer.PublicKey);
            }
            catch (FormatException)
            {
                result.Rejected.Add(new RejectedItem { Id = peer.NodeId.Value, Reason = RejectReasons.InvalidParams });
            }
        }

        _peers = peers;

        BudgetGuard guard = Guard;
        double usedMbps = 0;
        long usedPps = 0;
        var line = new Dictionary<string, (TaskSpec Task, TaskDecision Decision)>();
        var targets = new List<(TaskSpec Spec, TaskDecision Decision)>();

        foreach (TaskSpec task in set.Tasks)
        {
            if (!task.Enabled)
            {
                continue;
            }

            bool shapeOk = task.Scope == "target" ? s_targetKinds.Contains(task.Kind) : _acceptLineTasks && s_lineKinds.Contains(task.Kind);
            TaskDecision decision = shapeOk ? guard.Check(task, ref usedMbps, ref usedPps) : TaskDecision.Reject(RejectReasons.Unsupported);
            if (!decision.Accepted)
            {
                result.Rejected.Add(new RejectedItem { Id = task.Id.Value, Reason = decision.Reason! });
                continue;
            }

            if (task.Scope == "target")
            {
                targets.Add((task, decision));
            }
            else
            {
                line[task.Id.Value] = (task, decision);
            }
        }

        _targets.Reconcile(targets, guard.AddressAllowed);

        // 参数没有变化的任务沿用原来的流和积木，数据面上的会话（重放窗口、记录）不被打断
        var retired = new List<LineEntry>();
        var changes = new List<(DialerSession Dialer, LineTaskSet Set)>();
        lock (_gate)
        {
            _lineSpecs = line;
            _version = version;
            foreach ((DialerSession dialer, Dictionary<string, LineEntry> current) in _sessionLines)
            {
                var next = new Dictionary<string, LineEntry>();
                bool lineTasksChanged = current.Count != line.Count;
                foreach ((string id, (TaskSpec task, TaskDecision decision)) in line)
                {
                    LineEntry candidate = BuildLineEntry(task, decision, dialer);
                    if (current.TryGetValue(id, out LineEntry? existing) && existing.Signature == candidate.Signature)
                    {
                        next[id] = existing;
                    }
                    else
                    {
                        lineTasksChanged = true;
                        Start(candidate);
                        next[id] = candidate;
                    }
                }

                foreach ((string id, LineEntry entry) in current)
                {
                    if (!next.TryGetValue(id, out LineEntry? kept) || !ReferenceEquals(kept, entry))
                    {
                        retired.Add(entry);
                    }
                }

                current.Clear();
                foreach ((string id, LineEntry entry) in next)
                {
                    current[id] = entry;
                }

                // The Agent receives a full task snapshot for peer/trust changes and target-only tasks too.
                // Those changes do not affect a 拨号方's line blocks, so avoid retransmitting the same full
                // line_task_set to every 拨号方. New sessions still get their initial snapshot in OpenSession.
                if (lineTasksChanged)
                {
                    changes.Add((dialer, SetFor(current)));
                }
            }
        }

        foreach (LineEntry entry in retired)
        {
            await Retire(entry);
        }

        OnPeersChanged?.Invoke(new HashSet<string>(peers.Keys, StringComparer.Ordinal));
        foreach ((DialerSession dialer, LineTaskSet next) in changes)
        {
            OnLineTasksChanged?.Invoke(dialer, next);
        }

        return result;
    }

    /// <summary>停止时释放所有拨号方会话的流和所有目标任务（Backend 失联的时候任务继续运行：节点自治）。</summary>
    public async Task StopAsync()
    {
        _targets.Reconcile([]);
        List<LineEntry> entries;
        lock (_gate)
        {
            entries = [.. _sessionLines.Values.SelectMany(l => l.Values)];
            _sessionLines.Clear();
        }

        foreach (LineEntry entry in entries)
        {
            await Retire(entry);
        }
    }

    private LineTaskSet SetFor(Dictionary<string, LineEntry> entries) =>
        new()
        {
            Version = _version,
            Tasks = [.. entries.Values.Select(e => e.Task)],
            QuicPorts = [.. _quicPorts()],
            TcpPorts = [.. _tcpPorts()],
        };

    private LineEntry BuildLineEntry(TaskSpec task, TaskDecision decision, DialerSession dialer)
    {
        var info = new StreamInfo(_streams.Allocate(), task.Kind, decision.Resolved!) { TaskId = task.Id, Owner = dialer };
        var lineTask = new LineTask
        {
            Id = task.Id,
            Kind = task.Kind,
            Params = new Dictionary<string, JsonElement>(decision.Resolved!.Values),
            StreamId = info.Id,
            Port = info.Port,
        };
        var spec = new BlockSpec(decision.Kind!, decision.Resolved, info.Id, info.Port, new FragmentScope(task.Id, null, null));
        return new LineEntry(lineTask, info, spec, _env.WithSession(dialer));
    }

    /// <summary>在数据面登记这个流并启动它的监听方积木。</summary>
    private void Start(LineEntry entry)
    {
        IBlock block = BlockCatalog.Find(entry.Stream.Kind)!.Listener!(entry.Env, entry.Spec);
        entry.Stream.Running = RunningBlock.Start(block, entry.Signature);
        _streams.Add(entry.Stream);
    }

    private async Task Retire(LineEntry entry)
    {
        _streams.Remove(entry.Stream.Id);
        if (entry.Stream.Running is { } running)
        {
            await running.StopAsync();
        }

        entry.Stream.Owner?.ForgetStream(entry.Stream.Id);
    }

    private sealed record LineEntry(LineTask Task, StreamInfo Stream, BlockSpec Spec, BlockEnv Env)
    {
        public string Signature => Spec.ConfigSignature;
    }
}