// now: Nix-based distributed command runner
// Copyright (C) 2026 Eric Rodrigues Pires
//
// This program is free software: you can redistribute it and/or modify it under
// the terms of the GNU Affero General Public License as published by the Free
// Software Foundation, either version 3 of the License, or (at your option)
// any later version.
//
// This program is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
// FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for
// more details.
//
// You should have received a copy of the GNU Affero General Public License along
// with this program. If not, see <https://www.gnu.org/licenses/>.

const std = @import("std");
const builtin = @import("builtin");
const now_step = @import("now_step");

var terminating = std.atomic.Value(bool).init(false);
var signalled = std.atomic.Value(bool).init(false);

fn killChild(io: std.Io, child: *std.process.Child) void {
    child.kill(io);
    _ = child.wait(io) catch {};
}

fn waitForChild(child: *std.process.Child, io: std.Io, result: *std.process.Child.Term) void {
    waitForChildInner(child, io, result) catch |err| {
        std.log.err("waitForChild failed: {t}", .{err});
    };
}

fn waitForChildInner(child: *std.process.Child, io: std.Io, result: *std.process.Child.Term) !void {
    result.* = try child.wait(io);
}

fn pipe(reader: *std.Io.Reader, writer: *std.Io.Writer) void {
    pipeInner(reader, writer) catch |err| {
        std.log.err("pipe failed: {t}", .{err});
    };
}

fn pipeInner(reader: *std.Io.Reader, writer: *std.Io.Writer) !void {
    _ = try reader.streamRemaining(writer);
    try writer.flush();
}

fn takeLine(r: *std.Io.Reader) std.Io.Reader.DelimiterError![]u8 {
    var scanned: usize = 0;
    while (true) {
        const contents = r.buffer[r.seek..r.end];
        if (std.mem.findAnyPos(u8, contents, scanned, "\r\n")) |i| {
            if (contents[i] == '\n') {
                r.toss(i + 1);
                return contents[0..i];
            }
            if (i + 1 < contents.len) {
                if (contents[i + 1] == '\n') {
                    r.toss(i + 2);
                } else {
                    r.toss(i + 1);
                }
                return contents[0..i];
            }
            scanned = i;
        } else {
            scanned = contents.len;
        }
        if (contents.len == r.buffer.len) return error.StreamTooLong;
        r.fillMore() catch |err| switch (err) {
            error.ReadFailed => |e| return e,
            error.EndOfStream => {
                const rest = r.buffer[r.seek..r.end];
                if (rest.len != 0 and rest[rest.len - 1] == '\r') {
                    r.toss(rest.len);
                    return rest[0 .. rest.len - 1];
                }
                return error.EndOfStream;
            },
        };
    }
}

fn writeToStderr(mutex: *std.Io.Mutex, io: std.Io, reader: *std.Io.Reader, writer: *std.Io.Writer, secret_collection: now_step.SecretCollection) void {
    writeToStderrInner(mutex, io, reader, writer, secret_collection) catch |err| {
        std.log.err("writeToStderr failed: {t}", .{err});
    };
}

fn writeToStderrInner(mutex: *std.Io.Mutex, io: std.Io, reader: *std.Io.Reader, writer: *std.Io.Writer, secret_collection: now_step.SecretCollection) !void {
    while (true) {
        const line = takeLine(reader) catch |err| switch (err) {
            error.EndOfStream => {
                const rest = reader.buffered();
                if (rest.len > 0) {
                    const redacted = try secret_collection.redactLine(rest);
                    try mutex.lock(io);
                    defer mutex.unlock(io);
                    try writer.writeAll(redacted);
                    try writer.writeByte('\n');
                    try writer.flush();
                }
                break;
            },
            error.ReadFailed => {
                const rest = reader.buffered();
                if (rest.len > 0) {
                    const redacted = try secret_collection.redactLine(rest);
                    try mutex.lock(io);
                    defer mutex.unlock(io);
                    try writer.writeAll(redacted);
                    try writer.writeByte('\n');
                    try writer.flush();
                }
                break;
            },
            else => return err,
        };

        const redacted_line = try secret_collection.redactLine(line);
        try mutex.lock(io);
        defer mutex.unlock(io);
        try writer.writeAll(redacted_line);
        try writer.writeByte('\n');
        try writer.flush();
    }
}

fn terminateGroup(io: std.Io, child: *std.process.Child, timeout_secs: i64) !void {
    if (terminating.swap(true, .acq_rel)) return;
    const pgid = child.id.?;
    if (timeout_secs > 0) {
        std.posix.kill(-pgid, std.posix.SIG.TERM) catch {};
        io.sleep(.fromSeconds(timeout_secs), .awake) catch {};
    }
    std.posix.kill(-pgid, std.posix.SIG.KILL) catch {};
}

fn watchLifeline(io: std.Io, reader: *std.Io.Reader, child: *std.process.Child, timeout_secs: i64) !void {
    _ = reader.discardRemaining() catch {};
    try terminateGroup(io, child, timeout_secs);
}

fn onSignal(_: std.posix.SIG) callconv(.c) void {
    signalled.store(true, .release);
}

fn installSignalHandlers() !void {
    const act: std.posix.Sigaction = .{
        .handler = .{ .handler = onSignal },
        .mask = std.posix.sigemptyset(),
        .flags = std.posix.SA.RESTART,
    };
    std.posix.sigaction(std.posix.SIG.TERM, &act, null);
    std.posix.sigaction(std.posix.SIG.INT, &act, null);
    std.posix.sigaction(std.posix.SIG.HUP, &act, null);
}

fn watchSignals(io: std.Io, child: *std.process.Child, timeout_secs: i64) !void {
    while (!signalled.load(.acquire)) try io.sleep(.fromMilliseconds(50), .awake);
    try terminateGroup(io, child, timeout_secs);
}

pub fn main(init: std.process.Init) !void {
    try installSignalHandlers();
    const allocator: std.mem.Allocator = init.arena.allocator();

    var secret_names: std.ArrayList([]const u8) = .empty;
    defer secret_names.deinit(allocator);

    var arg_iterator = try init.minimal.args.iterateAllocator(allocator);
    defer arg_iterator.deinit();
    _ = arg_iterator.next().?; // Discard program name
    const timeout_secs = try std.fmt.parseInt(i64, arg_iterator.next().?, 10);
    const derivation_or_flag = arg_iterator.next().?;
    var derivation: [:0]const u8 = undefined;
    const preserve_stdout = std.mem.eql(u8, derivation_or_flag, "--preserve-stdout");
    if (preserve_stdout) {
        derivation = arg_iterator.next().?;
    } else {
        derivation = derivation_or_flag;
    }
    while (arg_iterator.next()) |secret| {
        try secret_names.append(allocator, secret);
    }

    var child_env = std.process.Environ.Map.init(allocator);
    defer child_env.deinit();
    var env_iterator = init.environ_map.iterator();
    while (env_iterator.next()) |entry| {
        try child_env.put(entry.key_ptr.*, entry.value_ptr.*);
    }

    var io_impl: std.Io.Threaded = .init(allocator, .{});
    defer io_impl.deinit();
    const io = io_impl.io();

    var lifeline = false;
    const stdin_buf = try allocator.alloc(u8, 128 * 1024);
    defer allocator.free(stdin_buf);
    var stdin: std.Io.File.Reader = .init(.stdin(), io, stdin_buf);
    const stdin_reader = &stdin.interface;

    while (try stdin_reader.takeDelimiter(0)) |entry| {
        if (entry.len == 0) {
            lifeline = true;
            break;
        }
        const sep = std.mem.findScalar(u8, entry, '=') orelse
            return error.MalformedEnvEntry;
        if (sep == 0) return error.EmptyEnvKey;
        try child_env.put(entry[0..sep], entry[sep + 1 ..]);
    }

    var secret_collection = try now_step.SecretCollection.init(allocator, &child_env, secret_names.items);
    defer secret_collection.deinit();

    try child_env.put("CI", "true");
    try child_env.put("TERM", "xterm-256color");

    var child = try std.process.spawn(io, .{
        .argv = &.{derivation},
        .environ_map = &child_env,
        .pgid = 0,
        .stdin = .ignore,
        .stdout = .pipe,
        .stderr = .pipe,
    });
    errdefer killChild(io, &child);

    // If stdin gets dropped, terminate the process gracefully
    if (lifeline) {
        const t = try std.Thread.spawn(.{}, watchLifeline, .{ io, stdin_reader, &child, timeout_secs });
        t.detach();
    }
    const t = try std.Thread.spawn(.{}, watchSignals, .{ io, &child, timeout_secs });
    t.detach();

    const child_stdout_buffer = try allocator.alloc(u8, 128 * 1024);
    defer allocator.free(child_stdout_buffer);
    var child_stdout_file_reader = child.stdout.?.reader(io, child_stdout_buffer);
    const child_stdout_reader = &child_stdout_file_reader.interface;

    const child_stderr_buffer = try allocator.alloc(u8, 128 * 1024);
    defer allocator.free(child_stderr_buffer);
    var child_stderr_file_reader = child.stderr.?.reader(io, child_stderr_buffer);
    const child_stderr_reader = &child_stderr_file_reader.interface;

    const stderr_buffer = try allocator.alloc(u8, 128 * 1024);
    defer allocator.free(stderr_buffer);
    var stderr_file_writer: std.Io.File.Writer = .init(.stderr(), io, stderr_buffer);
    const stderr_writer = &stderr_file_writer.interface;

    var result: std.process.Child.Term = undefined;
    var mutex = std.Io.Mutex.init;
    var task_group = std.Io.Group.init;
    if (preserve_stdout) {
        const stdout_buffer = try allocator.alloc(u8, 128 * 1024);
        defer allocator.free(stdout_buffer);
        var stdout_file_writer: std.Io.File.Writer = .init(.stdout(), io, stdout_buffer);
        const stdout_writer = &stdout_file_writer.interface;
        try task_group.concurrent(io, pipe, .{ child_stdout_reader, stdout_writer });
    } else {
        try task_group.concurrent(io, writeToStderr, .{ &mutex, io, child_stdout_reader, stderr_writer, secret_collection });
    }
    try task_group.concurrent(io, writeToStderr, .{ &mutex, io, child_stderr_reader, stderr_writer, secret_collection });
    try task_group.concurrent(io, waitForChild, .{ &child, io, &result });
    try task_group.await(io);

    const code: u8 = switch (result) {
        .exited => |c| c,
        .signal => |sig| 128 +% @as(u8, @truncate(@intFromEnum(sig))),
        else => 1,
    };
    std.process.exit(code);
}