using System.Net;
using System.Security.Principal;
using Agent.Admin.Client;
using Agent.Admin.Host;

namespace Agent.Windows.Tests;

internal sealed class DirectElevatedNetworkSession : IElevatedNetworkSession
{
    private WfpDynamicSession? _wfp;
    private SecurityIdentifier? _sid;
    private bool _loopbackAdded;

    public Task<bool> InspectLoopbackAsync(string applicationSid, CancellationToken cancellationToken) =>
        Task.FromResult(LoopbackExemption.Contains(new SecurityIdentifier(applicationSid)));

    public Task ApplyAsync(string applicationSid, string networkAccess, IReadOnlyList<AdminNetworkEndpoint> endpoints,
        bool originallyLoopbackExempt, CancellationToken cancellationToken)
    {
        _sid = new(applicationSid);
        NetworkEndpoint[] values = endpoints.Select(static item => new NetworkEndpoint(IPAddress.Parse(item.Address), (ushort)item.Port)).ToArray();
        _wfp = WfpDynamicSession.Open();
        _wfp.ApplyDefaultBlock(_sid, values);
        if (values.Any(static item => IPAddress.IsLoopback(item.Address)) && !originallyLoopbackExempt)
            _loopbackAdded = LoopbackExemption.Ensure(_sid);
        return Task.CompletedTask;
    }

    public Task RecoverLoopbackAsync(string applicationSid, bool originallyLoopbackExempt, CancellationToken cancellationToken)
    {
        var sid = new SecurityIdentifier(applicationSid);
        if (originallyLoopbackExempt) _ = LoopbackExemption.Ensure(sid); else LoopbackExemption.Remove(sid);
        return Task.CompletedTask;
    }

    public Task CleanupAsync(CancellationToken cancellationToken)
    {
        if (_loopbackAdded && _sid is not null) LoopbackExemption.Remove(_sid);
        _loopbackAdded = false;
        if (_wfp is not null) { _wfp.RemovePolicy(); _wfp.Dispose(); _wfp = null; }
        return Task.CompletedTask;
    }

    public async ValueTask DisposeAsync() => await CleanupAsync(CancellationToken.None).ConfigureAwait(false);
}

internal sealed class DirectElevatedNetworkSessionFactory : IElevatedNetworkSessionFactory
{
    public int Starts { get; private set; }
    public Task<IElevatedNetworkSession> StartAsync(Guid taskId, CancellationToken cancellationToken)
    {
        Starts++;
        return Task.FromResult<IElevatedNetworkSession>(new DirectElevatedNetworkSession());
    }
}

internal sealed class NoOpElevatedNetworkSessionFactory : IElevatedNetworkSessionFactory
{
    public int Starts { get; private set; }
    public Task<IElevatedNetworkSession> StartAsync(Guid taskId, CancellationToken cancellationToken)
    {
        Starts++;
        return Task.FromResult<IElevatedNetworkSession>(new NoOpElevatedNetworkSession());
    }

    private sealed class NoOpElevatedNetworkSession : IElevatedNetworkSession
    {
        public Task<bool> InspectLoopbackAsync(string applicationSid, CancellationToken cancellationToken) => Task.FromResult(false);
        public Task ApplyAsync(string applicationSid, string networkAccess, IReadOnlyList<AdminNetworkEndpoint> endpoints,
            bool originallyLoopbackExempt, CancellationToken cancellationToken) => Task.CompletedTask;
        public Task RecoverLoopbackAsync(string applicationSid, bool originallyLoopbackExempt, CancellationToken cancellationToken) => Task.CompletedTask;
        public Task CleanupAsync(CancellationToken cancellationToken) => Task.CompletedTask;
        public ValueTask DisposeAsync() => ValueTask.CompletedTask;
    }
}