using System.Diagnostics;
using System.IO.Pipes;
using System.Net;
using System.Runtime.InteropServices;
using System.Security.Principal;
using Agent.Admin.Client;
using Microsoft.Win32.SafeHandles;

namespace Agent.Admin.Host;

internal static partial class Program
{
    private const string FixedProfileName = "AgentSandboxMcp.V1";

    public static async Task<int> Main(string[] args)
    {
        try
        {
            Arguments options = Arguments.Parse(args);
            using WindowsIdentity identity = WindowsIdentity.GetCurrent();
            if (!new WindowsPrincipal(identity).IsInRole(WindowsBuiltInRole.Administrator))
                throw new InvalidOperationException("The elevated helper requires an administrator token.");
            using Process parent = Process.GetProcessById(checked((int)options.ParentPid));
            if (parent.HasExited || !string.Equals(Path.GetFileName(parent.MainModule?.FileName), "agent-mcp.exe", StringComparison.OrdinalIgnoreCase))
                throw new InvalidOperationException("The parent Process is not the fixed agent-mcp executable.");

            using var pipe = new NamedPipeClientStream(".", options.PipeName, PipeDirection.InOut,
                PipeOptions.Asynchronous | PipeOptions.WriteThrough, TokenImpersonationLevel.Identification);
            await pipe.ConnectAsync(30_000).ConfigureAwait(false);
            if (!Native.GetNamedPipeServerProcessId(pipe.SafePipeHandle, out uint serverPid) || serverPid != options.ParentPid)
                throw new InvalidOperationException("The Named Pipe server Process does not match the expected parent.");

            AdminRequest hello = await AdminProtocol.ReadAsync<AdminRequest>(pipe, CancellationToken.None).ConfigureAwait(false);
            ValidateEnvelope(hello, options, "hello");
            await AdminProtocol.WriteAsync(pipe, new AdminResponse(true), CancellationToken.None).ConfigureAwait(false);

            using var state = new NetworkState();
            while (true)
            {
                AdminRequest request = await AdminProtocol.ReadAsync<AdminRequest>(pipe, CancellationToken.None).ConfigureAwait(false);
                if (request.Version != AdminProtocol.Version || request.TaskId != options.TaskId)
                {
                    await AdminProtocol.WriteAsync(pipe, new AdminResponse(false, "sandbox.ipc_rejected", "Protocol or Task identity mismatch."), CancellationToken.None).ConfigureAwait(false);
                    return 3;
                }
                try
                {
                    AdminResponse response = request.Operation switch
                    {
                        "inspect-loopback" => Inspect(request),
                        "apply-network" => Apply(request, state),
                        "cleanup-network" => Cleanup(state),
                        "recover-loopback" => Recover(request),
                        "shutdown" => new(true),
                        _ => new(false, "sandbox.ipc_rejected", "Unsupported administrative pipe operation.")
                    };
                    await AdminProtocol.WriteAsync(pipe, response, CancellationToken.None).ConfigureAwait(false);
                    if (request.Operation == "shutdown") return 0;
                    if (!response.Success && response.ErrorCode == "sandbox.ipc_rejected") return 3;
                }
                catch (Exception ex) when (ex is not (OutOfMemoryException or StackOverflowException or AccessViolationException))
                {
                    await AdminProtocol.WriteAsync(pipe, new AdminResponse(false, "sandbox.network_apply_failed", ex.Message), CancellationToken.None).ConfigureAwait(false);
                }
            }
        }
        catch (Exception ex) when (ex is not (OutOfMemoryException or StackOverflowException or AccessViolationException))
        {
            Console.Error.WriteLine($"agent-mcp-admin: {ex.Message}");
            return 2;
        }
    }

    private static AdminResponse Inspect(AdminRequest request)
    {
        SecurityIdentifier sid = ValidateSid(request.ApplicationSid);
        return new(true, LoopbackExempt: LoopbackExemption.Contains(sid));
    }

    private static AdminResponse Apply(AdminRequest request, NetworkState state)
    {
        if (state.Applied) return new(false, "sandbox.ipc_rejected", "Network Task Resource is already applied.");
        SecurityIdentifier sid = ValidateSid(request.ApplicationSid);
        IReadOnlyList<NetworkEndpoint> endpoints = ValidateNetwork(request.NetworkAccess, request.Endpoints);
        bool needsLoopback = endpoints.Any(static item => IPAddress.IsLoopback(item.Address));
        bool currentLoopback = LoopbackExemption.Contains(sid);
        if (needsLoopback && currentLoopback != request.OriginallyLoopbackExempt)
            return new(false, "sandbox.network_apply_failed", "Loopback state changed after Recovery Journal preparation.");

        state.Session = WfpDynamicSession.Open();
        try
        {
            state.Session.ApplyDefaultBlock(sid, endpoints);
            if (needsLoopback && !currentLoopback) state.LoopbackAdded = LoopbackExemption.Ensure(sid);
            state.Sid = sid;
            state.Applied = true;
            return new(true);
        }
        catch
        {
            if (state.LoopbackAdded) LoopbackExemption.Remove(sid);
            state.Session.Dispose();
            state.Session = null;
            state.LoopbackAdded = false;
            throw;
        }
    }

    private static AdminResponse Cleanup(NetworkState state)
    {
        state.Cleanup();
        return new(true);
    }

    private static AdminResponse Recover(AdminRequest request)
    {
        SecurityIdentifier sid = ValidateSid(request.ApplicationSid);
        if (request.OriginallyLoopbackExempt == true) _ = LoopbackExemption.Ensure(sid);
        else LoopbackExemption.Remove(sid);
        return new(true);
    }

    private static SecurityIdentifier ValidateSid(string? value)
    {
        if (string.IsNullOrWhiteSpace(value)) throw new InvalidDataException("Application SID is required.");
        var requested = new SecurityIdentifier(value);
        int result = Native.DeriveAppContainerSidFromAppContainerName(FixedProfileName, out nint pointer);
        if (result < 0) Marshal.ThrowExceptionForHR(result);
        try
        {
            var fixedSid = new SecurityIdentifier(pointer);
            if (!fixedSid.Equals(requested)) throw new InvalidDataException("Only the fixed Application Profile SID is accepted.");
            return requested;
        }
        finally { if (pointer != 0) _ = Native.FreeSid(pointer); }
    }

    private static IReadOnlyList<NetworkEndpoint> ValidateNetwork(string? access, IReadOnlyList<AdminNetworkEndpoint>? input)
    {
        if (access is not ("none" or "local" or "restricted")) throw new InvalidDataException("Unsupported Network Access Mode.");
        if (input is null || input.Count > 4096) throw new InvalidDataException("Network endpoint list is invalid.");
        var endpoints = input.Select(item =>
        {
            if (!IPAddress.TryParse(item.Address, out IPAddress? address) || item.Port is < 1 or > 65535)
                throw new InvalidDataException("Network endpoint is invalid.");
            return new NetworkEndpoint(address, (ushort)item.Port);
        }).Distinct().ToArray();
        if (access == "none" && endpoints.Length != 0) throw new InvalidDataException("Network none must not contain endpoints.");
        if (access == "local" && (endpoints.Length == 0 || endpoints.Any(static item => !IPAddress.IsLoopback(item.Address))))
            throw new InvalidDataException("Network local accepts only IPv4/IPv6 loopback endpoints.");
        if (access == "restricted" && endpoints.Length == 0) throw new InvalidDataException("Network restricted requires endpoints.");
        return endpoints;
    }

    private static void ValidateEnvelope(AdminRequest request, Arguments options, string operation)
    {
        if (request.Version != options.Protocol || request.Version != AdminProtocol.Version ||
            request.TaskId != options.TaskId || request.Operation != operation)
            throw new InvalidDataException("Administrative pipe handshake mismatch.");
    }

    private sealed class NetworkState : IDisposable
    {
        public WfpDynamicSession? Session { get; set; }
        public SecurityIdentifier? Sid { get; set; }
        public bool LoopbackAdded { get; set; }
        public bool Applied { get; set; }
        public void Cleanup()
        {
            if (LoopbackAdded && Sid is not null) LoopbackExemption.Remove(Sid);
            LoopbackAdded = false;
            if (Session is not null)
            {
                try { Session.RemovePolicy(); }
                finally { Session.Dispose(); Session = null; }
            }
            Applied = false;
        }
        public void Dispose() { try { Cleanup(); } catch { Session?.Dispose(); } }
    }

    private sealed record Arguments(string PipeName, uint ParentPid, Guid TaskId, int Protocol)
    {
        public static Arguments Parse(string[] args)
        {
            var values = new Dictionary<string, string>(StringComparer.Ordinal);
            for (int i = 0; i < args.Length; i += 2)
            {
                if (i + 1 >= args.Length || args[i] is not ("--pipe" or "--parent-pid" or "--task" or "--protocol"))
                    throw new ArgumentException("The elevated helper accepts only its fixed internal arguments.");
                values.Add(args[i], args[i + 1]);
            }
            if (values.Count != 4 || !uint.TryParse(values["--parent-pid"], out uint parentPid) ||
                !Guid.TryParseExact(values["--task"], "N", out Guid taskId) || !int.TryParse(values["--protocol"], out int protocol) ||
                protocol != AdminProtocol.Version || string.IsNullOrWhiteSpace(values["--pipe"]))
                throw new ArgumentException("The elevated helper arguments are invalid.");
            return new(values["--pipe"], parentPid, taskId, protocol);
        }
    }

    private static partial class Native
    {
        [LibraryImport("kernel32.dll", SetLastError = true)]
        [return: MarshalAs(UnmanagedType.Bool)]
        internal static partial bool GetNamedPipeServerProcessId(SafePipeHandle pipe, out uint serverProcessId);
        [LibraryImport("userenv.dll", StringMarshalling = StringMarshalling.Utf16)]
        internal static partial int DeriveAppContainerSidFromAppContainerName(string name, out nint sid);
        [LibraryImport("advapi32.dll")]
        internal static partial nint FreeSid(nint sid);
    }
}