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>
    /// 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"]);
        var probe = RunQuiet(bwrap, ["--unshare-all", "--die-with-parent", "--ro-bind", "/usr", "/usr",
            "--symlink", "usr/bin", "/bin", "--symlink", "usr/lib", "/lib", "--proc", "/proc", "--dev", "/dev", "/bin/true"]);
        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);

        var filterFile = Path.Combine(Path.GetTempPath(), $"clawddotnet_seccomp_{Guid.NewGuid():N}.bpf");
        await File.WriteAllBytesAsync(filterFile, SeccompFilter.Build(RuntimeInformation.ProcessArchitecture), ct);
        try
        {
            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))
                start.ArgumentList.Add(argument);

            return await RunProcessAsync(start, request, ct);
        }
        finally
        {
            try { File.Delete(filterFile); }
            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)
    {
        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,
            Kilobytes(limits.MaxMemoryBytes),
            ((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)]);

        args.AddRange(["--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"
        };
        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: the process limit, set inside the sandbox's own user namespace.
        args.AddRange([Shell, "-c", """ulimit -u "$1" || exit 126; shift; exec "$@" """, "sh",
            limits.MaxProcesses.ToString(CultureInfo.InvariantCulture), 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 = Process.Start(start) ?? throw new InvalidOperationException("The sandbox did not 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;
        }

        await Task.WhenAll(stdout, stderr);
        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>bwrap dies, and with it — die-with-parent, own PID namespace — everything inside.</summary>
    private static void Kill(Process process)
    {
        try
        {
            if (!process.HasExited)
                process.Kill(entireProcessTree: true);
        }
        catch (InvalidOperationException) { }
        catch (System.ComponentModel.Win32Exception) { }
    }

    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: 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")]);
}