using System.ComponentModel;
using System.Runtime.InteropServices;
using System.Security.Principal;
using System.Text;
using Microsoft.Win32.SafeHandles;
using Windows.Win32;
using Windows.Win32.Foundation;
using Windows.Win32.Security;
using Windows.Win32.System.JobObjects;
using Windows.Win32.System.Threading;

namespace Agent.Windows.Security;

public sealed record RestrictedProcessResult(uint ProcessId, uint ExitCode, string StandardOutput, string StandardError);
public sealed record SandboxJobLimits(uint MaxProcesses = 16, ulong MaxJobMemoryBytes = 4UL * 1024 * 1024 * 1024);

public sealed unsafe class RestrictedProcessLauncher(ExecutablePolicy executablePolicy, ISecurityAuditSink audit, SecurityIdentifier? appContainerSid = null, SandboxJobLimits? jobLimits = null)
{
    private const uint ProcessCreationChildProcessRestricted = 0x1;
    private readonly SandboxJobLimits _jobLimits = jobLimits ?? new();

    public RestrictedProcessResult Run(Guid taskId, string executable, IReadOnlyList<string> arguments, string workingDirectory, TimeSpan timeout, string? standardInput = null)
    {
        VerifiedExecutable verified;
        VerifiedExecutable? managedEntry = null;
        try
        {
            verified = executablePolicy.Verify(executable);
            string? managedEntryPath = ExecutablePolicy.ManagedEntryPoint(verified.CanonicalPath, arguments);
            if (managedEntryPath is not null) managedEntry = executablePolicy.Verify(managedEntryPath);
        }
        catch (SandboxPolicyException ex)
        {
            audit.Write(new(DateTimeOffset.UtcNow, taskId, "policy.deny", "denied", new Dictionary<string, string> { ["code"] = ex.Code, ["executable"] = executable }));
            throw;
        }

        string cwd = CanonicalPath.Resolve(workingDirectory);
        using FileStream stdin = CaptureFile();
        using FileStream stdout = CaptureFile();
        using FileStream stderr = CaptureFile();
        if (!string.IsNullOrEmpty(standardInput))
        {
            byte[] input = Encoding.UTF8.GetBytes(standardInput);
            stdin.Write(input);
            stdin.Flush(true);
        }
        stdin.Position = 0;
        MakeInheritable(stdin.SafeFileHandle);
        MakeInheritable(stdout.SafeFileHandle);
        MakeInheritable(stderr.SafeFileHandle);
        using SafeFileHandle job = PInvoke.CreateJobObject(null, null);
        if (job.IsInvalid) throw LastError("CreateJobObject");
        try
        {
            JOBOBJECT_EXTENDED_LIMIT_INFORMATION limits = default;
            limits.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE |
                JOB_OBJECT_LIMIT.JOB_OBJECT_LIMIT_ACTIVE_PROCESS | JOB_OBJECT_LIMIT.JOB_OBJECT_LIMIT_JOB_MEMORY;
            limits.BasicLimitInformation.ActiveProcessLimit = _jobLimits.MaxProcesses;
            limits.JobMemoryLimit = (nuint)_jobLimits.MaxJobMemoryBytes;
            if (!PInvoke.SetInformationJobObject(new HANDLE(job.DangerousGetHandle()), JOBOBJECTINFOCLASS.JobObjectExtendedLimitInformation, &limits, (uint)sizeof(JOBOBJECT_EXTENDED_LIMIT_INFORMATION)))
                throw LastError("SetInformationJobObject");

            nuint attributeSize = 0;
            uint attributeCount = appContainerSid is null ? 2u : 3u;
            _ = PInvoke.InitializeProcThreadAttributeList(default, attributeCount, 0, &attributeSize);
            LPPROC_THREAD_ATTRIBUTE_LIST attributes = (LPPROC_THREAD_ATTRIBUTE_LIST)NativeMemory.Alloc(attributeSize);
            void* sidBuffer = null;
            void* internetClientSidBuffer = null;
            void* privateNetworkSidBuffer = null;
            try
            {
                if (!PInvoke.InitializeProcThreadAttributeList(attributes, attributeCount, 0, &attributeSize)) throw LastError("InitializeProcThreadAttributeList");
                uint childPolicy = ProcessCreationChildProcessRestricted;
                if (!PInvoke.UpdateProcThreadAttribute(attributes, 0, PInvoke.PROC_THREAD_ATTRIBUTE_CHILD_PROCESS_POLICY, &childPolicy, sizeof(uint), null, null))
                    throw LastError("UpdateProcThreadAttribute");
                HANDLE[] inheritedHandles =
                [
                    new(stdin.SafeFileHandle.DangerousGetHandle()),
                    new(stdout.SafeFileHandle.DangerousGetHandle()),
                    new(stderr.SafeFileHandle.DangerousGetHandle())
                ];
                fixed (HANDLE* inheritedHandlePointer = inheritedHandles)
                {
                    if (!PInvoke.UpdateProcThreadAttribute(attributes, 0, PInvoke.PROC_THREAD_ATTRIBUTE_HANDLE_LIST,
                        inheritedHandlePointer, (nuint)(sizeof(HANDLE) * inheritedHandles.Length), null, null))
                        throw LastError("UpdateProcThreadAttribute(HandleList)");
                }
                SecurityCapabilities securityCapabilities = default;
                if (appContainerSid is not null)
                {
                    byte[] sidBytes = new byte[appContainerSid.BinaryLength];
                    appContainerSid.GetBinaryForm(sidBytes, 0);
                    sidBuffer = NativeMemory.Alloc((nuint)sidBytes.Length);
                    Marshal.Copy(sidBytes, 0, (nint)sidBuffer, sidBytes.Length);
                    securityCapabilities.AppContainerSid = (nint)sidBuffer;
                    internetClientSidBuffer = AllocateSid(new SecurityIdentifier("S-1-15-3-1"));
                    privateNetworkSidBuffer = AllocateSid(new SecurityIdentifier("S-1-15-3-3"));
                    SID_AND_ATTRIBUTES* capabilities = stackalloc SID_AND_ATTRIBUTES[2];
                    capabilities[0].Sid = new PSID(internetClientSidBuffer);
                    capabilities[0].Attributes = 0x00000004;
                    capabilities[1].Sid = new PSID(privateNetworkSidBuffer);
                    capabilities[1].Attributes = 0x00000004;
                    securityCapabilities.Capabilities = (nint)capabilities;
                    securityCapabilities.CapabilityCount = 2;
                    if (!PInvoke.UpdateProcThreadAttribute(attributes, 0, PInvoke.PROC_THREAD_ATTRIBUTE_SECURITY_CAPABILITIES,
                        &securityCapabilities, (nuint)sizeof(SecurityCapabilities), null, null)) throw LastError("UpdateProcThreadAttribute(SecurityCapabilities)");
                }
                STARTUPINFOEXW startup = default;
                startup.StartupInfo.cb = (uint)sizeof(STARTUPINFOEXW);
                startup.StartupInfo.dwFlags = STARTUPINFOW_FLAGS.STARTF_USESTDHANDLES;
                startup.StartupInfo.hStdInput = inheritedHandles[0];
                startup.StartupInfo.hStdOutput = inheritedHandles[1];
                startup.StartupInfo.hStdError = inheritedHandles[2];
                startup.lpAttributeList = attributes;
                PROCESS_INFORMATION process = default;
                string commandLine = BuildCommandLine(verified.CanonicalPath, arguments);
                Span<char> commandBuffer = new char[commandLine.Length + 1];
                commandLine.AsSpan().CopyTo(commandBuffer);
                if (!PInvoke.CreateProcess(null, ref commandBuffer, null, null, true,
                    PROCESS_CREATION_FLAGS.EXTENDED_STARTUPINFO_PRESENT | PROCESS_CREATION_FLAGS.CREATE_SUSPENDED,
                    null, cwd, in startup.StartupInfo, out process)) throw LastError("CreateProcess");
                try
                {
                    using var processView = new SafeFileHandle(process.hProcess, ownsHandle: false);
                    if (!PInvoke.AssignProcessToJobObject(job, processView)) throw LastError("AssignProcessToJobObject");
                    try
                    {
                        VerifiedExecutable postCreate = executablePolicy.Verify(executable);
                        if (!string.Equals(postCreate.CanonicalPath, verified.CanonicalPath, StringComparison.OrdinalIgnoreCase) ||
                            !string.Equals(postCreate.Sha256, verified.Sha256, StringComparison.Ordinal))
                            throw new SandboxPolicyException("process.identity_changed", "Executable identity changed during process creation.");
                        if (managedEntry is not null)
                        {
                            VerifiedExecutable postCreateManaged = executablePolicy.Verify(managedEntry.CanonicalPath);
                            if (!string.Equals(postCreateManaged.Sha256, managedEntry.Sha256, StringComparison.Ordinal))
                                throw new SandboxPolicyException("process.managed_identity_changed", "Managed entry assembly identity changed during process creation.");
                        }
                    }
                    catch (SandboxPolicyException ex)
                    {
                        audit.Write(new(DateTimeOffset.UtcNow, taskId, "policy.deny", "denied", new Dictionary<string, string>
                        {
                            ["code"] = ex.Code, ["executable"] = executable, ["stage"] = "post_create_pre_resume"
                        }));
                        throw;
                    }
                    if (PInvoke.ResumeThread(process.hThread) == uint.MaxValue) throw LastError("ResumeThread");
                    var processAudit = new Dictionary<string, string>
                    {
                        ["path"] = verified.CanonicalPath, ["sha256"] = verified.Sha256, ["child_process_policy"] = "restricted",
                        ["app_container_sid"] = appContainerSid?.Value ?? "none", ["network_capabilities"] = appContainerSid is null ? "none" : "internetClient,privateNetworkClientServer",
                        ["pid"] = process.dwProcessId.ToString()
                    };
                    if (managedEntry is not null)
                    {
                        processAudit["managed_entry_path"] = managedEntry.CanonicalPath;
                        processAudit["managed_entry_sha256"] = managedEntry.Sha256;
                    }
                    audit.Write(new(DateTimeOffset.UtcNow, taskId, "process.start", "allowed", processAudit));
                    WAIT_EVENT wait = PInvoke.WaitForSingleObject(process.hProcess, checked((uint)Math.Clamp(timeout.TotalMilliseconds, 1, uint.MaxValue - 1)));
                    if (wait == WAIT_EVENT.WAIT_TIMEOUT)
                    {
                        _ = PInvoke.TerminateJobObject(job, 1460);
                        audit.Write(new(DateTimeOffset.UtcNow, taskId, "process.timeout", "terminated", new Dictionary<string, string> { ["pid"] = process.dwProcessId.ToString() }));
                        throw new TimeoutException("Sandbox process exceeded its timeout.");
                    }
                    if (!PInvoke.GetExitCodeProcess(processView, out uint exitCode)) throw LastError("GetExitCodeProcess");
                    audit.Write(new(DateTimeOffset.UtcNow, taskId, "process.exit", "observed", new Dictionary<string, string> { ["pid"] = process.dwProcessId.ToString(), ["exit_code"] = exitCode.ToString() }));
                    if (exitCode != 0)
                    {
                        audit.Write(new(DateTimeOffset.UtcNow, taskId, "policy.deny", "observed", new Dictionary<string, string>
                        {
                            ["code"] = "restricted_process_failure", ["pid"] = process.dwProcessId.ToString(),
                            ["child_process_policy"] = "restricted", ["exit_code"] = exitCode.ToString()
                        }));
                    }
                    return new(process.dwProcessId, exitCode, ReadCapture(stdout), ReadCapture(stderr));
                }
                finally
                {
                    PInvoke.CloseHandle(process.hThread);
                    PInvoke.CloseHandle(process.hProcess);
                }
            }
            finally
            {
                if (attributes.Value is not null) PInvoke.DeleteProcThreadAttributeList(attributes);
                NativeMemory.Free(sidBuffer);
                NativeMemory.Free(internetClientSidBuffer);
                NativeMemory.Free(privateNetworkSidBuffer);
                NativeMemory.Free(attributes.Value);
            }
        }
        finally { }
    }

    private static Win32Exception LastError(string operation)
    {
        int code = Marshal.GetLastWin32Error();
        return new(code, $"{operation} failed with Win32 error {code}.");
    }
    private static FileStream CaptureFile()
    {
        string path = Path.Combine(Path.GetTempPath(), "agent-mcp-process-" + Guid.NewGuid().ToString("N") + ".tmp");
        return new FileStream(path, FileMode.CreateNew, FileAccess.ReadWrite, FileShare.ReadWrite | FileShare.Delete, 4096,
            FileOptions.DeleteOnClose | FileOptions.SequentialScan);
    }
    private static void* AllocateSid(SecurityIdentifier sid)
    {
        byte[] bytes = new byte[sid.BinaryLength];
        sid.GetBinaryForm(bytes, 0);
        void* buffer = NativeMemory.Alloc((nuint)bytes.Length);
        Marshal.Copy(bytes, 0, (nint)buffer, bytes.Length);
        return buffer;
    }
    private static void MakeInheritable(SafeFileHandle handle)
    {
        if (!PInvoke.SetHandleInformation(handle, (uint)HANDLE_FLAGS.HANDLE_FLAG_INHERIT, HANDLE_FLAGS.HANDLE_FLAG_INHERIT))
            throw LastError("SetHandleInformation");
    }
    private static string ReadCapture(FileStream stream)
    {
        const int maxCharacters = 1024 * 1024;
        stream.Position = 0;
        using var reader = new StreamReader(stream, Encoding.UTF8, detectEncodingFromByteOrderMarks: true, leaveOpen: true);
        char[] buffer = new char[maxCharacters + 1];
        int count = reader.ReadBlock(buffer, 0, buffer.Length);
        return count > maxCharacters ? new string(buffer, 0, maxCharacters) + "\n[output truncated]" : new string(buffer, 0, count);
    }
    private static string BuildCommandLine(string executable, IReadOnlyList<string> arguments) => string.Join(' ', [Quote(executable), .. arguments.Select(Quote)]);
    private static string Quote(string value)
    {
        if (value.Length > 0 && !value.Any(static character => char.IsWhiteSpace(character) || character == '"')) return value;
        var result = new System.Text.StringBuilder().Append('"');
        int slashes = 0;
        foreach (char character in value)
        {
            if (character == '\\') { slashes++; continue; }
            if (character == '"') result.Append('\\', (slashes * 2) + 1).Append('"');
            else result.Append('\\', slashes).Append(character);
            slashes = 0;
        }
        return result.Append('\\', slashes * 2).Append('"').ToString();
    }

    [StructLayout(LayoutKind.Sequential)]
    private struct SecurityCapabilities
    {
        public nint AppContainerSid;
        public nint Capabilities;
        public uint CapabilityCount;
        public uint Reserved;
    }
}