using System.ComponentModel;
using System.Buffers.Binary;
using System.Net;
using System.Runtime.InteropServices;
using System.Security.Principal;
using Windows.Win32;
using Windows.Win32.Foundation;
using Windows.Win32.NetworkManagement.WindowsFilteringPlatform;
using Windows.Win32.Security;

namespace Agent.Admin.Host;

public sealed record NetworkEndpoint(IPAddress Address, ushort Port);

public sealed unsafe class WfpDynamicSession : IDisposable
{
    private const FWPM_FILTER_FLAGS FwpmFilterFlagClearActionRight = (FWPM_FILTER_FLAGS)0x00000008;
    private readonly FwpmEngineClose0SafeHandle _engine;
    private Guid? _subLayerKey;
    private readonly List<ulong> _filterIds = [];

    private WfpDynamicSession(FwpmEngineClose0SafeHandle engine) => _engine = engine;
    public static WfpDynamicSession Open(string displayName = "Agent Sandbox MCP dynamic task session")
    {
        FWPM_SESSION0 session = default;
        session.flags = PInvoke.FWPM_SESSION_FLAG_DYNAMIC;
        uint result = PInvoke.FwpmEngineOpen0(10, null, session, out FwpmEngineClose0SafeHandle engine);
        ThrowIfError(result, "FwpmEngineOpen0");
        return new(engine);
    }

    public void ProbeWriteTransaction()
    {
        uint result = PInvoke.FwpmTransactionBegin0(_engine, 0);
        ThrowIfError(result, "FwpmTransactionBegin0");
        result = PInvoke.FwpmTransactionAbort0(_engine);
        ThrowIfError(result, "FwpmTransactionAbort0");
    }

    public IReadOnlyList<ulong> ApplyDefaultBlock(SecurityIdentifier appContainerSid, IReadOnlyList<NetworkEndpoint>? allowedEndpoints = null)
    {
        if (_subLayerKey is not null) throw new InvalidOperationException("WFP policy is already applied in this session.");
        byte[] sidBytes = new byte[appContainerSid.BinaryLength];
        appContainerSid.GetBinaryForm(sidBytes, 0);
        uint result = PInvoke.FwpmTransactionBegin0(_engine, 0);
        ThrowIfError(result, "FwpmTransactionBegin0");
        bool committed = false;
        try
        {
            Guid subLayerKey = Guid.NewGuid();
            fixed (char* subLayerName = "Agent Sandbox MCP task sublayer")
            {
                FWPM_SUBLAYER0 subLayer = default;
                subLayer.subLayerKey = subLayerKey;
                subLayer.displayData.name = new PWSTR(subLayerName);
                subLayer.weight = 0x7FFF;
                ThrowIfError(PInvoke.FwpmSubLayerAdd0(_engine, in subLayer, default), "FwpmSubLayerAdd0");
            }

            fixed (byte* sidPointer = sidBytes)
            {
                foreach (NetworkEndpoint endpoint in allowedEndpoints ?? []) _filterIds.Add(AddPackagePermit(endpoint, subLayerKey, (SID*)sidPointer));
                _filterIds.Add(AddPackageBlock(PInvoke.FWPM_LAYER_ALE_AUTH_CONNECT_V4, subLayerKey, (SID*)sidPointer, "Agent Sandbox MCP IPv4 default block"));
                _filterIds.Add(AddPackageBlock(PInvoke.FWPM_LAYER_ALE_AUTH_CONNECT_V6, subLayerKey, (SID*)sidPointer, "Agent Sandbox MCP IPv6 default block"));
            }
            ThrowIfError(PInvoke.FwpmTransactionCommit0(_engine), "FwpmTransactionCommit0");
            committed = true;
            _subLayerKey = subLayerKey;
            return _filterIds.ToArray();
        }
        finally
        {
            if (!committed)
            {
                _filterIds.Clear();
                _ = PInvoke.FwpmTransactionAbort0(_engine);
            }
        }
    }

    public void RemovePolicy()
    {
        if (_subLayerKey is null) return;
        ThrowIfError(PInvoke.FwpmTransactionBegin0(_engine, 0), "FwpmTransactionBegin0(remove)");
        bool committed = false;
        try
        {
            foreach (ulong filterId in _filterIds.AsEnumerable().Reverse()) ThrowIfError(PInvoke.FwpmFilterDeleteById0(_engine, filterId), "FwpmFilterDeleteById0");
            Guid subLayerKey = _subLayerKey.Value;
            ThrowIfError(PInvoke.FwpmSubLayerDeleteByKey0(_engine, in subLayerKey), "FwpmSubLayerDeleteByKey0");
            ThrowIfError(PInvoke.FwpmTransactionCommit0(_engine), "FwpmTransactionCommit0(remove)");
            committed = true;
            _filterIds.Clear();
            _subLayerKey = null;
        }
        finally
        {
            if (!committed) _ = PInvoke.FwpmTransactionAbort0(_engine);
        }
    }

    public void Dispose()
    {
        _engine.Dispose();
    }

    private ulong AddPackageBlock(Guid layerKey, Guid subLayerKey, SID* sid, string displayName)
    {
        FWPM_FILTER_CONDITION0 condition = default;
        condition.fieldKey = PInvoke.FWPM_CONDITION_ALE_PACKAGE_ID;
        condition.matchType = FWP_MATCH_TYPE.FWP_MATCH_EQUAL;
        condition.conditionValue.type = FWP_DATA_TYPE.FWP_SID;
        condition.conditionValue.sid = sid;
        fixed (char* name = displayName)
        {
            FWPM_FILTER0 filter = default;
            filter.filterKey = Guid.NewGuid();
            filter.displayData.name = new PWSTR(name);
            filter.layerKey = layerKey;
            filter.subLayerKey = subLayerKey;
            filter.weight.type = FWP_DATA_TYPE.FWP_UINT8;
            filter.weight.uint8 = 8;
            filter.numFilterConditions = 1;
            filter.filterCondition = &condition;
            filter.action.type = FWP_ACTION_TYPE.FWP_ACTION_BLOCK;
            ThrowIfError(PInvoke.FwpmFilterAdd0(_engine, in filter, default, out ulong id), "FwpmFilterAdd0");
            return id;
        }
    }

    private ulong AddPackagePermit(NetworkEndpoint endpoint, Guid subLayerKey, SID* sid)
    {
        FWPM_FILTER_CONDITION0* conditions = stackalloc FWPM_FILTER_CONDITION0[3];
        conditions[0].fieldKey = PInvoke.FWPM_CONDITION_ALE_PACKAGE_ID;
        conditions[0].matchType = FWP_MATCH_TYPE.FWP_MATCH_EQUAL;
        conditions[0].conditionValue.type = FWP_DATA_TYPE.FWP_SID;
        conditions[0].conditionValue.sid = sid;
        conditions[1].fieldKey = PInvoke.FWPM_CONDITION_IP_REMOTE_ADDRESS;
        conditions[1].matchType = FWP_MATCH_TYPE.FWP_MATCH_EQUAL;
        Guid layer;
        FWP_BYTE_ARRAY16 ipv6 = default;
        if (endpoint.Address.AddressFamily == System.Net.Sockets.AddressFamily.InterNetwork)
        {
            layer = PInvoke.FWPM_LAYER_ALE_AUTH_CONNECT_V4;
            conditions[1].conditionValue.type = FWP_DATA_TYPE.FWP_UINT32;
            // FWP_UINT32 uses host-order numeric values. For 127.0.0.1 the
            // numeric value is 0x7F000001; parse the network-order octets into
            // that value rather than reinterpreting their in-memory byte order.
            conditions[1].conditionValue.uint32 = BinaryPrimitives.ReadUInt32BigEndian(endpoint.Address.GetAddressBytes());
        }
        else
        {
            layer = PInvoke.FWPM_LAYER_ALE_AUTH_CONNECT_V6;
            ipv6.byteArray16 = endpoint.Address.GetAddressBytes();
            conditions[1].conditionValue.type = FWP_DATA_TYPE.FWP_BYTE_ARRAY16_TYPE;
            conditions[1].conditionValue.byteArray16 = &ipv6;
        }
        conditions[2].fieldKey = PInvoke.FWPM_CONDITION_IP_REMOTE_PORT;
        conditions[2].matchType = FWP_MATCH_TYPE.FWP_MATCH_EQUAL;
        conditions[2].conditionValue.type = FWP_DATA_TYPE.FWP_UINT16;
        conditions[2].conditionValue.uint16 = endpoint.Port;
        string displayName = $"Agent Sandbox MCP allow {endpoint.Address}:{endpoint.Port}";
        fixed (char* name = displayName)
        {
            FWPM_FILTER0 filter = default;
            filter.filterKey = Guid.NewGuid();
            filter.displayData.name = new PWSTR(name);
            filter.layerKey = layer;
            filter.subLayerKey = subLayerKey;
            filter.weight.type = FWP_DATA_TYPE.FWP_UINT8;
            filter.weight.uint8 = 15;
            filter.numFilterConditions = 3;
            filter.filterCondition = conditions;
            // A plain permit is a soft permit and can be overridden by a block
            // from another sub-layer (for example Windows Firewall). The exact
            // task allowance must be a hard permit; our higher-weight match in
            // this sub-layer still prevents the task default-block from running.
            filter.flags = FwpmFilterFlagClearActionRight;
            filter.action.type = FWP_ACTION_TYPE.FWP_ACTION_PERMIT;
            ThrowIfError(PInvoke.FwpmFilterAdd0(_engine, in filter, default, out ulong id), "FwpmFilterAdd0(permit)");
            return id;
        }
    }

    private static void ThrowIfError(uint result, string operation)
    {
        if (result != 0) throw new Win32Exception(unchecked((int)result), $"{operation} failed with WFP/Win32 error 0x{result:X8}.");
    }
}