using System.Net;
using System.Security.Principal;
using Agent.Admin.Client;
using Agent.Windows.Recovery;
using Agent.Windows.Security;

namespace Agent.Windows.Network;

public sealed class TaskNetworkPolicy : IAsyncDisposable
{
    private readonly IElevatedNetworkSession _session;
    private readonly TaskResourceJournal _journal;
    private readonly ISecurityAuditSink _audit;
    private readonly SecurityIdentifier _appContainerSid;
    private readonly Guid? _loopbackEntryId;
    private readonly Guid? _sessionId;
    private bool _disposed;

    private TaskNetworkPolicy(IElevatedNetworkSession session, TaskResourceJournal journal, ISecurityAuditSink audit,
        SecurityIdentifier appContainerSid, Guid? loopbackEntryId, Guid? sessionId)
    {
        _session = session;
        _journal = journal;
        _audit = audit;
        _appContainerSid = appContainerSid;
        _loopbackEntryId = loopbackEntryId;
        _sessionId = sessionId;
    }

    public static async Task<TaskNetworkPolicy> ApplyAsync(IElevatedNetworkSession session, TaskResourceJournal journal,
        SecurityIdentifier appContainerSid, string networkAccess, IReadOnlyList<NetworkEndpoint> allowances,
        ISecurityAuditSink audit, Guid? sessionId, CancellationToken cancellationToken)
    {
        bool needsLoopback = allowances.Any(static item => IPAddress.IsLoopback(item.Address));
        bool originallyExempt = needsLoopback && await session.InspectLoopbackAsync(appContainerSid.Value, cancellationToken).ConfigureAwait(false);
        Guid? loopbackEntryId = needsLoopback ? journal.Prepare("network.loopback_exemption", appContainerSid.Value,
            new Dictionary<string, string> { ["originally_present"] = originallyExempt.ToString() }) : null;
        try
        {
            AdminNetworkEndpoint[] endpoints = allowances.Select(static item => new AdminNetworkEndpoint(item.Address.ToString(), item.Port)).ToArray();
            await session.ApplyAsync(appContainerSid.Value, networkAccess, endpoints, originallyExempt, cancellationToken).ConfigureAwait(false);
            if (loopbackEntryId is not null) journal.MarkApplied(loopbackEntryId.Value);
            audit.Write(new(DateTimeOffset.UtcNow, journal.TaskId, "network.rule.apply", "applied", new Dictionary<string, string>
            {
                ["ipv4_ipv6_default_block"] = "true", ["endpoint_count"] = endpoints.Length.ToString(), ["session"] = "dynamic"
            }, SessionId: sessionId));
            if (loopbackEntryId is not null)
                audit.Write(new(DateTimeOffset.UtcNow, journal.TaskId, "network.loopback.apply", originallyExempt ? "preserved" : "added",
                    new Dictionary<string, string> { ["app_container_sid"] = appContainerSid.Value }, SessionId: sessionId));
            return new(session, journal, audit, appContainerSid, loopbackEntryId, sessionId);
        }
        catch
        {
            bool cleaned = false;
            try { await session.CleanupAsync(CancellationToken.None).ConfigureAwait(false); cleaned = true; } catch { }
            if (cleaned && loopbackEntryId is not null) journal.MarkCleaned(loopbackEntryId.Value);
            try { await session.DisposeAsync().ConfigureAwait(false); } catch { }
            throw;
        }
    }

    public async ValueTask DisposeAsync()
    {
        if (_disposed) return;
        try
        {
            await _session.CleanupAsync(CancellationToken.None).ConfigureAwait(false);
            if (_loopbackEntryId is not null) _journal.MarkCleaned(_loopbackEntryId.Value);
        }
        finally { await _session.DisposeAsync().ConfigureAwait(false); }
        _audit.Write(new(DateTimeOffset.UtcNow, _journal.TaskId, "network.rule.remove", "removed",
            new Dictionary<string, string> { ["residual_filters"] = "0" }, SessionId: _sessionId));
        if (_loopbackEntryId is not null)
            _audit.Write(new(DateTimeOffset.UtcNow, _journal.TaskId, "network.loopback.remove", "restored",
                new Dictionary<string, string> { ["app_container_sid"] = _appContainerSid.Value }, SessionId: _sessionId));
        _disposed = true;
    }
}