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(EntryExecutablePolicy executablePolicy, ISecurityAuditSink audit, SecurityIdentifier? appContainerSid = null,
    SandboxJobLimits? jobLimits = null, bool allowNetwork = false)
{
    private const uint ProcessCreationAllApplicationPackagesOptOut = 0x1;
    private readonly SandboxJobLimits _jobLimits = jobLimits ?? new();

    public RestrictedProcessResult Run(Guid taskId, string executable, IReadOnlyList<string> arguments, string workingDirectory, TimeSpan timeout,
        string? standardInput = null, Guid? sessionId = null) =>
        Run(taskId, executable, arguments, workingDirectory, timeout, standardInput, sessionId, CancellationToken.None);

    public RestrictedProcessResult Run(Guid taskId, string executable, IReadOnlyList<string> arguments, string workingDirectory, TimeSpan timeout,
        string? standardInput, Guid? sessionId, CancellationToken cancellationToken)
    {
        string cwd = ResolvedPath.Resolve(workingDirectory);
        VerifiedEntryExecutable verified;
        string? managedEntry = null;
        try
        {
            verified = executablePolicy.VerifyEntry(executable, cwd);
            managedEntry = executablePolicy.ManagedEntryPoint(verified.ResolvedPath, arguments, cwd);
            if (managedEntry is not null) executablePolicy.VerifyManagedEntryPoint(managedEntry);
        }
        catch (SandboxPolicyException ex)
        {
            audit.Write(new(DateTimeOffset.UtcNow, taskId, "policy.deny", "denied", new Dictionary<string, string> { ["code"] = ex.Code, ["executable"] = executable }, SessionId: sessionId, ErrorCode: ex.Code));
            throw;
        }

        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 ? 1u : 3u;
            _ = PInvoke.InitializeProcThreadAttributeList(default, attributeCount, 0, &attributeSize);
            LPPROC_THREAD_ATTRIBUTE_LIST attributes = (LPPROC_THREAD_ATTRIBUTE_LIST)NativeMemory.Alloc(attributeSize);
            void* sidBuffer = null;
            var capabilityBuffers = new List<nint>();
            try
            {
                if (!PInvoke.InitializeProcThreadAttributeList(attributes, attributeCount, 0, &attributeSize)) throw LastError("InitializeProcThreadAttributeList");
                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;
                    SecurityIdentifier[] capabilitySids = allowNetwork
                        ? [DeriveCapabilitySid("registryRead"), DeriveCapabilitySid("lpacInstrumentation"), new("S-1-15-3-1"), new("S-1-15-3-3")]
                        : [DeriveCapabilitySid("registryRead"), DeriveCapabilitySid("lpacInstrumentation")];
                    SID_AND_ATTRIBUTES* capabilities = stackalloc SID_AND_ATTRIBUTES[capabilitySids.Length];
                    for (int i = 0; i < capabilitySids.Length; i++)
                    {
                        void* buffer = AllocateSid(capabilitySids[i]);
                        capabilityBuffers.Add((nint)buffer);
                        capabilities[i].Sid = new PSID(buffer);
                        capabilities[i].Attributes = 0x00000004;
                    }
                    securityCapabilities.Capabilities = (nint)capabilities;
                    securityCapabilities.CapabilityCount = (uint)capabilitySids.Length;
                    if (!PInvoke.UpdateProcThreadAttribute(attributes, 0, PInvoke.PROC_THREAD_ATTRIBUTE_SECURITY_CAPABILITIES,
                        &securityCapabilities, (nuint)sizeof(SecurityCapabilities), null, null)) throw LpacFailure("UpdateProcThreadAttribute(SecurityCapabilities)", verified.ResolvedPath);
                    uint allApplicationPackagesPolicy = ProcessCreationAllApplicationPackagesOptOut;
                    if (!PInvoke.UpdateProcThreadAttribute(attributes, 0, PInvoke.PROC_THREAD_ATTRIBUTE_ALL_APPLICATION_PACKAGES_POLICY,
                        &allApplicationPackagesPolicy, sizeof(uint), null, null)) throw LpacFailure("UpdateProcThreadAttribute(AllApplicationPackagesPolicy)", verified.ResolvedPath);
                }
                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.ResolvedPath, arguments);
                Span<char> commandBuffer = new char[commandLine.Length + 1];
                commandLine.AsSpan().CopyTo(commandBuffer);
                char[] environmentBlock = BuildEnvironmentBlock(cwd);
                fixed (char* environment = environmentBlock)
                {
                    if (!PInvoke.CreateProcess(null, ref commandBuffer, null, null, true,
                        PROCESS_CREATION_FLAGS.EXTENDED_STARTUPINFO_PRESENT | PROCESS_CREATION_FLAGS.CREATE_SUSPENDED |
                        PROCESS_CREATION_FLAGS.CREATE_UNICODE_ENVIRONMENT,
                        environment, 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
                    {
                        VerifiedEntryExecutable postCreate = executablePolicy.VerifyEntry(executable, cwd);
                        if (!string.Equals(postCreate.ResolvedPath, verified.ResolvedPath, 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)
                            executablePolicy.VerifyManagedEntryPoint(managedEntry);
                    }
                    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"
                        }, SessionId: sessionId, ErrorCode: ex.Code));
                        throw;
                    }
                    (bool tokenIsAppContainer, string? tokenAppContainerSid, bool processInJob) = InspectConfinement(processView, job);
                    if (appContainerSid is not null && (!tokenIsAppContainer ||
                        !string.Equals(tokenAppContainerSid, appContainerSid.Value, StringComparison.Ordinal) || !processInJob))
                        throw new SandboxPolicyException("process.lpac_mismatch",
                            "Entry Process did not receive the required LPAC identity and Job Object confinement.",
                            "process", verified.ResolvedPath, "execute", "Verify LPAC support on this Windows host; no fallback is permitted.");
                    if (PInvoke.ResumeThread(process.hThread) == uint.MaxValue) throw LastError("ResumeThread");
                    var processAudit = new Dictionary<string, string>
                    {
                        ["path"] = verified.ResolvedPath, ["sha256"] = verified.Sha256, ["entry_classification"] = verified.Classification.ToString(),
                        ["alias"] = verified.Alias ?? string.Empty, ["descendant_policy"] = "lpac_inherited",
                        ["app_container_sid"] = appContainerSid?.Value ?? "none", ["lpac"] = (appContainerSid is not null).ToString(),
                        ["token_is_app_container"] = tokenIsAppContainer.ToString(), ["token_app_container_sid"] = tokenAppContainerSid ?? "none",
                        ["process_in_job"] = processInJob.ToString(), ["all_application_packages_opt_out"] = (appContainerSid is not null).ToString(),
                        ["network_capabilities"] = appContainerSid is null || !allowNetwork ? "none" : "internetClient,privateNetworkClientServer",
                        ["baseline_capabilities"] = appContainerSid is null ? "none" : "registryRead,lpacInstrumentation",
                        ["environment_policy"] = "minimal_no_broker_credentials",
                        ["arguments"] = "redacted", ["argument_count"] = arguments.Count.ToString(System.Globalization.CultureInfo.InvariantCulture),
                        ["pid"] = process.dwProcessId.ToString()
                    };
                    if (managedEntry is not null)
                    {
                        processAudit["managed_entry_path"] = ResolvedPath.Resolve(managedEntry);
                    }
                    audit.Write(new(DateTimeOffset.UtcNow, taskId, "process.start", "allowed", processAudit, SessionId: sessionId));
                    long deadline = Environment.TickCount64 + checked((long)Math.Clamp(timeout.TotalMilliseconds, 1, int.MaxValue));
                    WAIT_EVENT wait;
                    while (true)
                    {
                        cancellationToken.ThrowIfCancellationRequested();
                        long remaining = deadline - Environment.TickCount64;
                        if (remaining <= 0) { wait = WAIT_EVENT.WAIT_TIMEOUT; break; }
                        wait = PInvoke.WaitForSingleObject(process.hProcess, (uint)Math.Min(remaining, 100));
                        if (wait != WAIT_EVENT.WAIT_TIMEOUT) break;
                    }
                    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() }, SessionId: sessionId));
                        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() }, SessionId: sessionId));
                    if (exitCode != 0)
                    {
                        audit.Write(new(DateTimeOffset.UtcNow, taskId, "policy.deny", "observed", new Dictionary<string, string>
                        {
                            ["code"] = "sandbox_process_failure", ["pid"] = process.dwProcessId.ToString(),
                            ["descendant_policy"] = "lpac_inherited", ["exit_code"] = exitCode.ToString()
                        }, SessionId: sessionId, ErrorCode: "process.nonzero_exit"));
                    }
                    return new(process.dwProcessId, exitCode, ReadCapture(stdout), ReadCapture(stderr));
                }
                catch (OperationCanceledException)
                {
                    _ = PInvoke.TerminateJobObject(job, 1223);
                    _ = PInvoke.WaitForSingleObject(process.hProcess, 5000);
                    audit.Write(new(DateTimeOffset.UtcNow, taskId, "process.cancel", "terminated",
                        new Dictionary<string, string> { ["pid"] = process.dwProcessId.ToString() }, SessionId: sessionId));
                    throw;
                }
                finally
                {
                    PInvoke.CloseHandle(process.hThread);
                    PInvoke.CloseHandle(process.hProcess);
                }
            }
            finally
            {
                if (attributes.Value is not null) PInvoke.DeleteProcThreadAttributeList(attributes);
                NativeMemory.Free(sidBuffer);
                foreach (nint buffer in capabilityBuffers) NativeMemory.Free((void*)buffer);
                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 SandboxPolicyException LpacFailure(string operation, string executable)
    {
        int code = Marshal.GetLastWin32Error();
        return new("process.lpac_unsupported", $"{operation} failed with Win32 error {code}; no fallback is permitted.",
            "process", executable, "execute", "Use a supported Windows 11 host and verify LPAC capability derivation and process-attribute support.");
    }

    private static (bool IsAppContainer, string? AppContainerSid, bool InJob) InspectConfinement(SafeFileHandle process, SafeFileHandle job)
    {
        if (!PInvoke.IsProcessInJob(process, job, out BOOL inJob)) throw LastError("IsProcessInJob");
        if (!PInvoke.OpenProcessToken(process, TOKEN_ACCESS_MASK.TOKEN_QUERY, out SafeFileHandle token)) throw LastError("OpenProcessToken");
        using (token)
        {
            Span<byte> flag = stackalloc byte[sizeof(uint)];
            if (!PInvoke.GetTokenInformation(token, TOKEN_INFORMATION_CLASS.TokenIsAppContainer, flag, out _))
                throw LastError("GetTokenInformation(TokenIsAppContainer)");
            bool isAppContainer = MemoryMarshal.Read<uint>(flag) != 0;
            _ = PInvoke.GetTokenInformation(token, TOKEN_INFORMATION_CLASS.TokenAppContainerSid, Span<byte>.Empty, out uint sidInformationLength);
            if (sidInformationLength < IntPtr.Size) throw LastError("GetTokenInformation(TokenAppContainerSid size)");
            Span<byte> sidInformation = sidInformationLength <= 256 ? stackalloc byte[(int)sidInformationLength] : new byte[sidInformationLength];
            if (!PInvoke.GetTokenInformation(token, TOKEN_INFORMATION_CLASS.TokenAppContainerSid, sidInformation, out _))
                throw LastError("GetTokenInformation(TokenAppContainerSid)");
            nint sidPointer = MemoryMarshal.Read<nint>(sidInformation);
            string? sid = sidPointer == 0 ? null : new SecurityIdentifier(sidPointer).Value;
            return (isAppContainer, sid, inJob);
        }
    }
    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 SecurityIdentifier DeriveCapabilitySid(string name)
    {
        if (!PInvoke.DeriveCapabilitySidsFromName(name, out PSID* groupSids, out uint groupCount, out PSID* capabilitySids, out uint capabilityCount) || capabilityCount != 1)
            throw LastError($"DeriveCapabilitySidsFromName({name})");
        try { return new SecurityIdentifier((nint)capabilitySids[0].Value); }
        finally
        {
            for (uint i = 0; i < groupCount; i++) _ = PInvoke.LocalFree(new HLOCAL((nint)groupSids[i].Value));
            for (uint i = 0; i < capabilityCount; i++) _ = PInvoke.LocalFree(new HLOCAL((nint)capabilitySids[i].Value));
            _ = PInvoke.LocalFree(new HLOCAL((nint)groupSids));
            _ = PInvoke.LocalFree(new HLOCAL((nint)capabilitySids));
        }
    }
    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 char[] BuildEnvironmentBlock(string cwd)
    {
        string systemRoot = Environment.GetFolderPath(Environment.SpecialFolder.Windows);
        var values = new SortedDictionary<string, string>(StringComparer.OrdinalIgnoreCase)
        {
            ["ComSpec"] = Path.Combine(systemRoot, "System32", "cmd.exe"),
            ["DOTNET_CLI_HOME"] = cwd,
            ["DOTNET_CLI_TELEMETRY_OPTOUT"] = "1",
            ["DOTNET_CLI_USE_MSBUILD_SERVER"] = "0",
            ["DOTNET_CLI_WORKLOAD_UPDATE_NOTIFY_DISABLE"] = "true",
            ["DOTNET_NOLOGO"] = "1",
            ["DOTNET_SKIP_FIRST_TIME_EXPERIENCE"] = "1",
            ["DOTNET_SKIP_WORKLOAD_INTEGRITY_CHECK"] = "true",
            ["MSBuildEnableWorkloadResolver"] = "false",
            ["MSBUILDDISABLENODEREUSE"] = "1",
            ["NUMBER_OF_PROCESSORS"] = Environment.ProcessorCount.ToString(System.Globalization.CultureInfo.InvariantCulture),
            ["OS"] = "Windows_NT",
            ["PATH"] = Path.Combine(systemRoot, "System32"),
            ["PATHEXT"] = ".COM;.EXE;.BAT;.CMD",
            ["PROCESSOR_ARCHITECTURE"] = RuntimeInformation.ProcessArchitecture.ToString().ToUpperInvariant(),
            ["SystemDrive"] = Path.GetPathRoot(systemRoot)!.TrimEnd(Path.DirectorySeparatorChar),
            ["SystemRoot"] = systemRoot,
            ["TEMP"] = cwd,
            ["TMP"] = cwd,
            ["windir"] = systemRoot
        };
        foreach (string name in new[]
        {
            "ALLUSERSPROFILE", "APPDATA", "CommonProgramFiles", "CommonProgramFiles(x86)", "CommonProgramW6432",
            "HOMEDRIVE", "HOMEPATH", "LOCALAPPDATA", "ProgramData", "ProgramFiles", "ProgramFiles(x86)",
            "ProgramW6432", "PUBLIC", "USERDOMAIN", "USERNAME", "USERPROFILE"
        })
        {
            string? value = Environment.GetEnvironmentVariable(name);
            if (!string.IsNullOrEmpty(value)) values[name] = value;
        }
        string block = string.Concat(values.Select(pair => $"{pair.Key}={pair.Value}\0")) + "\0";
        return block.ToCharArray();
    }
    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;
    }
}