using System.Diagnostics;
using System.Globalization;
using System.Runtime.InteropServices;
using System.Text;
using Microsoft.Extensions.Logging;

namespace ClawdDotNet.Core.Sandbox;

/// <summary>
/// #654: The sandbox under Linux, built on bubblewrap.
///
/// <para><b>What holds, and where it comes from</b> (each with a test that tries to break it):</para>
/// <list type="bullet">
/// <item><b>File system:</b> own mount namespace. <c>/usr</c> and the few files of
/// <c>/etc</c> a program needs are mounted read-only; the workspace at <c>/work</c> is the
/// only writable place besides a private <c>/tmp</c>. <c>$HOME</c>, the instance, the key
/// ring and other workspaces are not mounted at all.</item>
/// <item><b>Inheritance:</b> the environment of bwrap itself is empty, the sandbox gets only
/// what the request names. .NET opens its files close-on-exec; the sandbox sees fds 0–2.</item>
/// <item><b>Network:</b> own network namespace with nothing but loopback.</item>
/// <item><b>Limits:</b> wall-clock timeout (the whole namespace dies with bwrap), address
/// space, CPU time, file size, open files — as rlimits; processes as RLIMIT_NPROC, set
/// inside the sandbox's own user namespace, where it counts the sandbox's processes and not
/// the instance's threads. Output is capped while it is read.</item>
/// <item><b>Privileges:</b> all capabilities dropped, <c>no_new_privs</c> (bwrap sets it
/// always), no nested user namespaces where bwrap can refuse them, and the
/// <see cref="SeccompFilter"/>.</item>
/// </list>
/// </summary>
public sealed class BwrapSandbox : ISandboxRunner
{
    private const string Shell = "/bin/sh";

    private readonly string _bwrap;
    private readonly BwrapFeatures _features;
    private readonly ILogger _logger;

    private BwrapSandbox(string bwrap, BwrapFeatures features, ILogger logger)
    {
        _bwrap = bwrap;
        _features = features;
        _logger = logger;
    }

    public bool IsAvailable => true;
    public string? UnavailableReason => null;

    /// <summary>
    /// #712: The launcher for runs that need network or secrets (<see cref="SandboxRequest.NeedsLauncher"/>).
    /// Without it such a run is refused, not run without them.
    /// </summary>
    public SandboxHelper? Helper { get; set; }

    /// <summary>#712: The logger the egress proxy reports to.</summary>
    public ILoggerFactory? LoggerFactory { get; set; }

    /// <summary>
    /// The runner of this machine: bwrap under Linux if it can actually build a sandbox here
    /// (tried once — a kernel or AppArmor policy that forbids user namespaces shows only when
    /// trying), otherwise <see cref="NoSandbox"/> with the reason.
    /// </summary>
    public static ISandboxRunner Detect(ILogger logger)
    {
        if (!OperatingSystem.IsLinux())
            return new NoSandbox("code runs only under Linux, in a bubblewrap sandbox");
        if (!SeccompFilter.Supports(RuntimeInformation.ProcessArchitecture))
            return new NoSandbox($"no seccomp filter for {RuntimeInformation.ProcessArchitecture}");

        var bwrap = new[] { "/usr/bin/bwrap", "/bin/bwrap", "/usr/local/bin/bwrap" }.FirstOrDefault(File.Exists);
        if (bwrap is null)
            return new NoSandbox("bubblewrap (bwrap) is not installed — 'apt install bubblewrap'");

        var help = RunQuiet(bwrap, ["--help"]);
        // The probe mounts the system exactly as a real run does — a missing /lib64 (the
        // dynamic loader) looks like "no such file" and would pass for a broken sandbox.
        var probeArgs = new List<string> { "--unshare-all", "--die-with-parent", "--ro-bind", "/usr", "/usr" };
        foreach (var (path, link) in HostLayout.Read().TopLevel)
            probeArgs.AddRange(link is not null ? ["--symlink", link, path] : ["--ro-bind-try", path, path]);
        probeArgs.AddRange(["--proc", "/proc", "--dev", "/dev", "/usr/bin/true"]);
        var probe = RunQuiet(bwrap, probeArgs);
        if (probe.ExitCode != 0)
            return new NoSandbox($"bwrap cannot create a sandbox here ({probe.Output.Trim()}) — user namespaces may be disabled");

        var features = new BwrapFeatures(
            DisableUserns: help.Output.Contains("--disable-userns", StringComparison.Ordinal),
            ClearEnv: help.Output.Contains("--clearenv", StringComparison.Ordinal));
        logger.LogInformation("Sandbox ready: {Bwrap}{Userns}", bwrap,
            features.DisableUserns ? "" : " (older bwrap: nested user namespaces are refused by the seccomp filter only)");
        return new BwrapSandbox(bwrap, features, logger);
    }

    public async Task<SandboxResult> RunAsync(SandboxRequest request, CancellationToken ct)
    {
        Validate(request);
        if (request.NeedsLauncher && Helper is null)
            throw new SandboxUnavailableException("this run needs network or secrets, and the sandbox launcher is not available");

        // A directory only this user can enter: the filter, the proxy's socket, the secrets.
        var privateDir = Path.Combine(Path.GetTempPath(), $"clawddotnet_sbx_{Guid.NewGuid():N}");
        Directory.CreateDirectory(privateDir);
        if (!OperatingSystem.IsWindows())
            File.SetUnixFileMode(privateDir, UnixFileMode.UserRead | UnixFileMode.UserWrite | UnixFileMode.UserExecute);

        EgressProxy? proxy = null;
        try
        {
            var filterFile = Path.Combine(privateDir, "seccomp.bpf");
            await File.WriteAllBytesAsync(filterFile, SeccompFilter.Build(RuntimeInformation.ProcessArchitecture), ct);

            var launch = LaunchFiles.None;
            if (request.Egress.Count > 0)
            {
                var socket = Path.Combine(privateDir, "egress.sock");
                proxy = EgressProxy.Start(socket, request.Egress,
                    LoggerFactory?.CreateLogger<EgressProxy>() ?? (ILogger)Microsoft.Extensions.Logging.Abstractions.NullLogger.Instance);
                launch = launch with { EgressSocket = socket };
            }
            if (request.SecretEnvironment.Count > 0)
            {
                var secrets = Path.Combine(privateDir, "env.json");
                await SandboxLauncher.WriteSecretFileAsync(secrets, request.SecretEnvironment, ct);
                launch = launch with { SecretFile = secrets };
            }

            var start = new ProcessStartInfo(Shell)
            {
                UseShellExecute = false,
                RedirectStandardInput = true,
                RedirectStandardOutput = true,
                RedirectStandardError = true,
                StandardOutputEncoding = Encoding.UTF8,
                StandardErrorEncoding = Encoding.UTF8
            };
            // Nothing of the main process's environment — not even to bwrap.
            start.Environment.Clear();
            foreach (var argument in CommandLine(_bwrap, filterFile, request, HostLayout.Read(), _features, Helper, launch))
                start.ArgumentList.Add(argument);

            return await RunProcessAsync(start, request, ct);
        }
        finally
        {
            proxy?.Dispose();
            try { Directory.Delete(privateDir, recursive: true); }
            catch (IOException) { /* in /tmp, harmless */ }
        }
    }

    /// <summary>
    /// The whole command line: <c>/bin/sh -c WRAPPER sh FILTER LIMITS… bwrap ARGS… -- sh -c INNER sh NPROC PROGRAM ARGS…</c>.
    /// Pure — the tests read it on any platform.
    /// </summary>
    public static IReadOnlyList<string> CommandLine(
        string bwrap, string filterFile, SandboxRequest request, HostLayout host, BwrapFeatures features,
        SandboxHelper? helper = null, LaunchFiles? launch = null)
    {
        launch ??= LaunchFiles.None;
        if (request.NeedsLauncher && helper is null)
            throw new ArgumentException("This run needs the sandbox launcher.");
        var limits = request.Limits;
        var args = new List<string>
        {
            "-c",
            // Outer: rlimits that bind the sandbox as a whole, then bwrap with the filter on fd 3.
            """f=$1; ulimit -v "$2" && ulimit -t "$3" && ulimit -f "$4" && ulimit -n "$5" || exit 126; shift 5; exec "$@" 3<"$f" """,
            "sh",
            filterFile,
            limits.MaxMemoryBytes > 0 ? Kilobytes(limits.MaxMemoryBytes) : "unlimited",
            ((long)Math.Ceiling(limits.Timeout.TotalSeconds) + 5).ToString(CultureInfo.InvariantCulture),
            Blocks(limits.MaxFileBytes),
            limits.MaxOpenFiles.ToString(CultureInfo.InvariantCulture),
            bwrap,
            "--unshare-all",
            "--unshare-user",
            "--die-with-parent",
            "--new-session",
            "--cap-drop", "ALL",
            "--hostname", "sandbox"
        };
        // bwrap's own environment is already empty; where it can, it says so once more.
        if (features.ClearEnv)
            args.Add("--clearenv");
        if (features.DisableUserns)
            args.Add("--disable-userns");

        // The system, read-only. Merged /usr: /bin, /lib … are links into it.
        args.AddRange(["--ro-bind", "/usr", "/usr"]);
        foreach (var (path, link) in host.TopLevel)
        {
            if (link is not null)
                args.AddRange(["--symlink", link, path]);
            else
                args.AddRange(["--ro-bind-try", path, path]);
        }
        foreach (var etc in EtcFiles)
            args.AddRange(["--ro-bind-try", etc, etc]);

        args.AddRange(["--proc", "/proc", "--dev", "/dev", "--tmpfs", "/tmp", "--dir", "/run"]);

        // The workspace, and inside it what must stay as it was checked.
        var workspace = Path.GetFullPath(request.WorkspacePath);
        args.AddRange(["--bind", workspace, SandboxPaths.Work]);
        foreach (var relative in request.ProtectedPaths)
        {
            var hostPath = Path.GetFullPath(Path.Combine(workspace, relative));
            if (SandboxPaths.InSandbox(workspace, hostPath) is { } inside && inside != SandboxPaths.Work)
                args.AddRange(["--ro-bind-try", hostPath, inside]);
        }
        foreach (var path in request.ReadOnlyPaths)
            args.AddRange(["--ro-bind", Path.GetFullPath(path), Path.GetFullPath(path)]);

        // #712: a runtime's home, the MCP socket — each at its own place, never inside /work,
        // where it would stand among the files that come back as the run's result.
        foreach (var mount in request.Mounts)
        {
            if (!mount.Target.StartsWith('/') || mount.Target == SandboxPaths.Work
                || mount.Target.StartsWith(SandboxPaths.Work + "/", StringComparison.Ordinal)
                || mount.Target.Split('/').Contains(".."))
                throw new ArgumentException($"'{mount.Target}' is not a place for a mount.");
            args.AddRange([mount.Writable ? "--bind" : "--ro-bind", Path.GetFullPath(mount.HostPath), mount.Target]);
        }

        // #712: what the launcher needs — the helper itself, the proxy's socket, the secrets.
        if (request.NeedsLauncher && helper is not null)
        {
            foreach (var path in helper.ReadOnlyPaths.Except(request.ReadOnlyPaths.Select(Path.GetFullPath)))
                args.AddRange(["--ro-bind", path, path]);
            if (launch.EgressSocket is { } socket)
                args.AddRange(["--bind", socket, SandboxPaths.EgressSocket]);
            if (launch.SecretFile is { } secrets)
                args.AddRange(["--ro-bind", secrets, SandboxPaths.SecretFile]);
        }

        // The root bwrap builds is a tmpfs of its own — writable unless it is remounted. Only
        // /work and /tmp, each a mount of its own, stay writable.
        args.AddRange(["--remount-ro", "/", "--chdir", SandboxPaths.Work]);

        var environment = new Dictionary<string, string>(StringComparer.Ordinal)
        {
            ["PATH"] = "/usr/local/bin:/usr/bin:/bin",
            ["HOME"] = "/tmp",
            ["LANG"] = "C.UTF-8",
            ["TMPDIR"] = "/tmp"
        };
        if (request.NeedsLauncher && helper is not null)
        {
            foreach (var (key, value) in helper.Environment)
                environment[key] = value;
            if (request.Egress.Count > 0)
            {
                // Everything goes to the proxy; the sandbox's loopback has nothing else.
                var proxy = $"http://127.0.0.1:{SandboxPaths.ProxyPort}";
                environment["HTTPS_PROXY"] = environment["https_proxy"] = proxy;
                environment["HTTP_PROXY"] = environment["http_proxy"] = proxy;
            }
        }
        foreach (var (key, value) in request.Environment)
            environment[key] = value;
        foreach (var (key, value) in environment.OrderBy(e => e.Key, StringComparer.Ordinal))
            args.AddRange(["--setenv", key, value]);

        args.AddRange(["--seccomp", "3", "--"]);

        // Inner: fd 3 — the filter bwrap read — is closed; bwrap hands it on otherwise. Then
        // the process limit, set inside the sandbox's own user namespace. bash calls it -u,
        // dash (Debian's and Ubuntu's sh) -p; if neither takes, nothing runs.
        args.AddRange([Shell, "-c", """exec 3<&-; ulimit -u "$1" 2>/dev/null || ulimit -p "$1" || exit 126; shift; exec "$@" """, "sh",
            limits.MaxProcesses.ToString(CultureInfo.InvariantCulture)]);
        if (request.NeedsLauncher && helper is not null)
        {
            args.AddRange([helper.Executable, SandboxLauncher.Argument]);
            if (request.Egress.Count > 0)
                args.AddRange([SandboxLauncher.EgressOption, SandboxPaths.EgressSocket]);
            if (request.SecretEnvironment.Count > 0)
                args.AddRange([SandboxLauncher.SecretsOption, SandboxPaths.SecretFile]);
            args.Add("--");
        }
        args.Add(request.Executable);
        args.AddRange(request.Arguments);
        return args;
    }

    /// <summary>What a program commonly needs of <c>/etc</c> — and nothing that says who uses this machine.</summary>
    private static readonly string[] EtcFiles =
    [
        "/etc/ld.so.cache", "/etc/ld.so.conf", "/etc/ld.so.conf.d", "/etc/alternatives",
        "/etc/localtime", "/etc/ssl/certs", "/etc/ca-certificates", "/etc/fonts", "/etc/mime.types"
    ];

    /// <summary>The workspace must exist and must not be a place that contains everything else.</summary>
    private static void Validate(SandboxRequest request)
    {
        var workspace = Path.GetFullPath(request.WorkspacePath).TrimEnd('/');
        if (!Directory.Exists(workspace))
            throw new ArgumentException($"The workspace '{workspace}' does not exist.");
        var home = System.Environment.GetFolderPath(System.Environment.SpecialFolder.UserProfile).TrimEnd('/');
        if (workspace.Length == 0 || workspace == home || workspace.Count(c => c == '/') < 2)
            throw new ArgumentException($"'{workspace}' is not a workspace — it would expose far more than one agent's files.");
        if (string.IsNullOrWhiteSpace(request.Executable))
            throw new ArgumentException("No program given.");
    }

    private async Task<SandboxResult> RunProcessAsync(ProcessStartInfo start, SandboxRequest request, CancellationToken ct)
    {
        using var process = await Spawner.StartAsync(start);
        using var timeout = CancellationTokenSource.CreateLinkedTokenSource(ct);
        timeout.CancelAfter(request.Limits.Timeout);

        var output = new CappedOutput(request.Limits.MaxOutputBytes);
        var stdout = PumpAsync(process.StandardOutput, output, isError: false, process);
        var stderr = PumpAsync(process.StandardError, output, isError: true, process);

        try
        {
            if (request.StandardInput is { } input)
                await process.StandardInput.WriteAsync(input.AsMemory(), timeout.Token);
            process.StandardInput.Close();
        }
        catch (IOException)
        {
            // The program did not read its input — that is its business.
        }

        var timedOut = false;
        try
        {
            await process.WaitForExitAsync(timeout.Token);
        }
        catch (OperationCanceledException)
        {
            timedOut = !ct.IsCancellationRequested;
            Kill(process);
            if (!timedOut)
                throw;
        }

        // After a stop, the streams end with the processes — but a reader must never hang on a
        // pipe something still holds. Bounded, then the result is what arrived.
        var pumps = Task.WhenAll(stdout, stderr);
        if (timedOut || output.Truncated)
            await Task.WhenAny(pumps, Task.Delay(TimeSpan.FromSeconds(5), CancellationToken.None));
        else
            await pumps;
        if (timedOut)
            _logger.LogWarning("Sandbox run of {Program} stopped after {Timeout}", request.Executable, request.Limits.Timeout);

        return new SandboxResult
        {
            ExitCode = process.HasExited ? process.ExitCode : -1,
            StandardOutput = output.Out,
            StandardError = output.Err,
            TimedOut = timedOut,
            OutputTruncated = output.Truncated
        };
    }

    /// <summary>Reads one stream into the shared cap; past the cap the run is stopped.</summary>
    private static async Task PumpAsync(StreamReader reader, CappedOutput output, bool isError, Process process)
    {
        var buffer = new char[8192];
        int read;
        while ((read = await reader.ReadAsync(buffer.AsMemory())) > 0)
        {
            if (!output.Append(buffer.AsSpan(0, read), isError))
            {
                Kill(process);
                return;
            }
        }
    }

    /// <summary>
    /// Stops the run: everything below bwrap is collected first, then all of it is killed.
    ///
    /// <para><b>Not trusting <c>--die-with-parent</c>.</b> The kernel clears the parent-death
    /// signal when a process changes its credentials — and bwrap does exactly that while it
    /// builds the user namespace. Seen on the test machine: bwrap's monitor was gone, the
    /// sandbox lived on under PID 1 and kept the output pipe open. Killing the sandbox's own
    /// PID 1 (the init of its PID namespace) takes every process inside with it.</para>
    /// </summary>
    private static void Kill(Process process)
    {
        int root;
        try
        {
            root = process.Id;
        }
        catch (InvalidOperationException)
        {
            return;
        }

        foreach (var pid in Descendants(root).Append(root))
        {
            try
            {
                using var victim = Process.GetProcessById(pid);
                victim.Kill();
            }
            catch (ArgumentException) { }            // already gone
            catch (InvalidOperationException) { }
            catch (System.ComponentModel.Win32Exception) { }
        }
    }

    /// <summary>All processes below <paramref name="root"/>, from /proc — collected before anything is killed.</summary>
    private static IEnumerable<int> Descendants(int root)
    {
        var children = new Dictionary<int, List<int>>();
        foreach (var dir in Directory.EnumerateDirectories("/proc"))
        {
            if (!int.TryParse(Path.GetFileName(dir), out var pid))
                continue;
            try
            {
                // /proc/<pid>/stat: "pid (comm) state ppid …" — comm may hold spaces and parentheses.
                var stat = File.ReadAllText(Path.Combine(dir, "stat"));
                var fields = stat[(stat.LastIndexOf(')') + 2)..].Split(' ');
                if (int.TryParse(fields[1], out var parent))
                {
                    if (!children.TryGetValue(parent, out var list))
                        children[parent] = list = [];
                    list.Add(pid);
                }
            }
            catch (IOException) { }                  // it ended meanwhile
            catch (UnauthorizedAccessException) { }
        }

        var result = new List<int>();
        var queue = new Queue<int>([root]);
        while (queue.Count > 0 && result.Count < 10_000)
        {
            foreach (var child in children.GetValueOrDefault(queue.Dequeue()) ?? [])
            {
                result.Add(child);
                queue.Enqueue(child);
            }
        }
        return result;
    }

    private static string Kilobytes(long bytes) => Math.Max(1, bytes / 1024).ToString(CultureInfo.InvariantCulture);

    /// <summary><c>ulimit -f</c> counts 512-byte blocks in POSIX sh.</summary>
    private static string Blocks(long bytes) => Math.Max(1, bytes / 512).ToString(CultureInfo.InvariantCulture);

    private static (int ExitCode, string Output) RunQuiet(string program, IReadOnlyList<string> arguments)
    {
        try
        {
            var start = new ProcessStartInfo(program)
            {
                UseShellExecute = false,
                RedirectStandardOutput = true,
                RedirectStandardError = true
            };
            start.Environment.Clear();
            foreach (var argument in arguments)
                start.ArgumentList.Add(argument);
            using var process = Process.Start(start)!;
            var text = process.StandardOutput.ReadToEnd() + process.StandardError.ReadToEnd();
            return process.WaitForExit(5000) ? (process.ExitCode, text) : (-1, "timed out");
        }
        catch (Exception ex) when (ex is System.ComponentModel.Win32Exception or IOException)
        {
            return (-1, ex.Message);
        }
    }

    private sealed class CappedOutput(int max)
    {
        private readonly StringBuilder _out = new();
        private readonly StringBuilder _err = new();
        private int _total;

        public bool Truncated { get; private set; }
        public string Out { get { lock (this) return _out.ToString(); } }
        public string Err { get { lock (this) return _err.ToString(); } }

        public bool Append(ReadOnlySpan<char> text, bool isError)
        {
            lock (this)
            {
                var room = max - _total;
                var take = Math.Min(room, text.Length);
                (isError ? _err : _out).Append(text[..Math.Max(0, take)]);
                _total += Math.Max(0, take);
                if (take < text.Length)
                    Truncated = true;
                return !Truncated;
            }
        }
    }
}

/// <summary>
/// #654: Starts every sandbox from one thread that lives as long as the process.
///
/// <para><b>Why.</b> <c>--die-with-parent</c> uses the kernel's parent-death signal, and that
/// is bound to the <i>thread</i> that forked, not to the process. .NET forks on the calling
/// thread — a thread-pool thread, and the pool retires idle threads. Seen on the test machine
/// under load: the pool dropped a thread, bwrap got SIGKILL (exit 137), and when it happened
/// while bwrap was still building the sandbox, the inside lived on as an orphan holding the
/// output pipes. Started from a thread that never ends, the signal fires only when the
/// instance itself dies — which is what it is for.</para>
/// </summary>
internal static class Spawner
{
    private static readonly System.Collections.Concurrent.BlockingCollection<(ProcessStartInfo Start, TaskCompletionSource<Process> Done)> Queue = new();

    static Spawner()
    {
        var thread = new Thread(() =>
        {
            foreach (var (start, done) in Queue.GetConsumingEnumerable())
            {
                try
                {
                    done.SetResult(Process.Start(start) ?? throw new InvalidOperationException("The sandbox did not start."));
                }
                catch (Exception ex)
                {
                    done.SetException(ex);
                }
            }
        })
        {
            IsBackground = true,
            Name = "ClawdDotNet sandbox spawner"
        };
        thread.Start();
    }

    public static Task<Process> StartAsync(ProcessStartInfo start)
    {
        var done = new TaskCompletionSource<Process>(TaskCreationOptions.RunContinuationsAsynchronously);
        Queue.Add((start, done));
        return done.Task;
    }
}

/// <summary>#712: Files the runner prepared for the launcher, on the host.</summary>
public sealed record LaunchFiles(string? EgressSocket, string? SecretFile)
{
    public static LaunchFiles None { get; } = new(null, null);
}

/// <summary>#654: What the installed bwrap can do beyond the basics.</summary>
public sealed record BwrapFeatures(bool DisableUserns, bool ClearEnv)
{
    public static BwrapFeatures All { get; } = new(true, true);
}

/// <summary>
/// #654: How the top level of this machine looks — which of <c>/bin</c>, <c>/lib</c> … are
/// links into <c>/usr</c> (merged /usr) and which are directories. Read once per run; a
/// record, so the tests can describe any machine.
/// </summary>
public sealed record HostLayout(IReadOnlyList<(string Path, string? Link)> TopLevel)
{
    private static readonly string[] Candidates = ["/bin", "/sbin", "/lib", "/lib64", "/lib32", "/libx32"];

    public static HostLayout Read()
    {
        var entries = new List<(string, string?)>();
        foreach (var path in Candidates)
        {
            var info = new FileInfo(path);
            if (info.LinkTarget is { } target)
                entries.Add((path, target.TrimStart('/')));
            else if (Directory.Exists(path))
                entries.Add((path, null));
        }
        return new HostLayout(entries);
    }

    /// <summary>A merged-/usr system, as Debian and Ubuntu are today.</summary>
    public static HostLayout MergedUsr { get; } = new([("/bin", "usr/bin"), ("/sbin", "usr/sbin"), ("/lib", "usr/lib"), ("/lib64", "usr/lib64")]);
}