using System.ComponentModel;
using System.Runtime.InteropServices;
using System.Security.Principal;
using Windows.Win32;
using Windows.Win32.Foundation;
using Windows.Win32.Security;
using Windows.Win32.System.Memory;

namespace Agent.Admin.Host;

internal static unsafe class LoopbackExemption
{
    private const string MutexName = "Local\\AgentSandboxMcpLoopbackExemption";

    public static bool Contains(SecurityIdentifier sid)
    {
        using MutexLease mutex = AcquireMutex();
        return Read().Any(item => item.Sid.Equals(sid));
    }

    public static bool Ensure(SecurityIdentifier sid)
    {
        using MutexLease mutex = AcquireMutex();
        List<Entry> entries = Read();
        if (entries.Any(item => item.Sid.Equals(sid))) return false;
        entries.Add(new(sid, 0));
        Write(entries);
        return true;
    }

    public static void Remove(SecurityIdentifier sid)
    {
        using MutexLease mutex = AcquireMutex();
        List<Entry> entries = Read();
        if (entries.RemoveAll(item => item.Sid.Equals(sid)) == 0) return;
        Write(entries);
    }

    private static MutexLease AcquireMutex()
    {
        var mutex = new Mutex(false, MutexName);
        try
        {
            try
            {
                if (!mutex.WaitOne(TimeSpan.FromSeconds(30))) throw new TimeoutException("Timed out waiting for the loopback exemption lock.");
            }
            catch (AbandonedMutexException)
            {
                // Ownership is granted to this thread when the prior owner died.
            }
            return new(mutex);
        }
        catch
        {
            mutex.Dispose();
            throw;
        }
    }

    private static List<Entry> Read()
    {
        uint result = PInvoke.NetworkIsolationGetAppContainerConfig(out uint count, out SID_AND_ATTRIBUTES* native);
        ThrowIfError(result, "NetworkIsolationGetAppContainerConfig");
        try
        {
            var entries = new List<Entry>(checked((int)count));
            for (uint index = 0; index < count; index++)
            {
                SID_AND_ATTRIBUTES item = native[index];
                entries.Add(new(new SecurityIdentifier((nint)item.Sid.Value), item.Attributes));
            }
            return entries;
        }
        finally
        {
            HANDLE heap = PInvoke.GetProcessHeap();
            for (uint index = 0; index < count; index++)
                if (native[index].Sid.Value is not null) _ = PInvoke.HeapFree(heap, HEAP_FLAGS.HEAP_NONE, native[index].Sid.Value);
            if (native is not null) _ = PInvoke.HeapFree(heap, HEAP_FLAGS.HEAP_NONE, native);
        }
    }

    private static void Write(IReadOnlyList<Entry> entries)
    {
        var native = new SID_AND_ATTRIBUTES[entries.Count];
        var allocations = new List<nint>(entries.Count);
        try
        {
            for (int index = 0; index < entries.Count; index++)
            {
                byte[] bytes = new byte[entries[index].Sid.BinaryLength];
                entries[index].Sid.GetBinaryForm(bytes, 0);
                void* allocation = NativeMemory.Alloc((nuint)bytes.Length);
                Marshal.Copy(bytes, 0, (nint)allocation, bytes.Length);
                allocations.Add((nint)allocation);
                native[index].Sid = new PSID(allocation);
                native[index].Attributes = entries[index].Attributes;
            }
            ThrowIfError(PInvoke.NetworkIsolationSetAppContainerConfig(native), "NetworkIsolationSetAppContainerConfig");
        }
        finally
        {
            foreach (nint allocation in allocations) NativeMemory.Free((void*)allocation);
        }
    }

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

    private sealed record Entry(SecurityIdentifier Sid, uint Attributes);

    private sealed class MutexLease(Mutex mutex) : IDisposable
    {
        public void Dispose()
        {
            mutex.ReleaseMutex();
            mutex.Dispose();
        }
    }
}