pub const Wip = struct {
    gpa: Allocator,
    string_bytes: std.ArrayListUnmanaged(u8),
    /// The first thing in this array is a ErrorMessageList.
    extra: std.ArrayListUnmanaged(u32),
    root_list: std.ArrayListUnmanaged(MessageIndex),

    pub fn init(wip: *Wip, gpa: Allocator) !void {
        wip.* = .{
            .gpa = gpa,
            .string_bytes = .{},
            .extra = .{},
            .root_list = .{},
        };

        // So that 0 can be used to indicate a null string.
        try wip.string_bytes.append(gpa, 0);

        assert(0 == try addExtra(wip, ErrorMessageList{
            .len = 0,
            .start = 0,
            .compile_log_text = 0,
        }));
    }

    pub fn deinit(wip: *Wip) void {
        const gpa = wip.gpa;
        wip.root_list.deinit(gpa);
        wip.string_bytes.deinit(gpa);
        wip.extra.deinit(gpa);
        wip.* = undefined;
    }

    pub fn toOwnedBundle(wip: *Wip, compile_log_text: []const u8) !ErrorBundle {
        const gpa = wip.gpa;
        if (wip.root_list.items.len == 0) {
            assert(compile_log_text.len == 0);
            // Special encoding when there are no errors.
            wip.deinit();
            wip.* = .{
                .gpa = gpa,
                .string_bytes = .{},
                .extra = .{},
                .root_list = .{},
            };
            return empty;
        }

        const compile_log_str_index = if (compile_log_text.len == 0) 0 else str: {
            const str: u32 = @intCast(wip.string_bytes.items.len);
            try wip.string_bytes.ensureUnusedCapacity(gpa, compile_log_text.len + 1);
            wip.string_bytes.appendSliceAssumeCapacity(compile_log_text);
            wip.string_bytes.appendAssumeCapacity(0);
            break :str str;
        };

        wip.setExtra(0, ErrorMessageList{
            .len = @intCast(wip.root_list.items.len),
            .start = @intCast(wip.extra.items.len),
            .compile_log_text = compile_log_str_index,
        });
        try wip.extra.appendSlice(gpa, @as([]const u32, @ptrCast(wip.root_list.items)));
        wip.root_list.clearAndFree(gpa);
        return .{
            .string_bytes = try wip.string_bytes.toOwnedSlice(gpa),
            .extra = try wip.extra.toOwnedSlice(gpa),
        };
    }

    pub fn tmpBundle(wip: Wip) ErrorBundle {
        return .{
            .string_bytes = wip.string_bytes.items,
            .extra = wip.extra.items,
        };
    }

    pub fn addString(wip: *Wip, s: []const u8) Allocator.Error!String {
        const gpa = wip.gpa;
        const index: String = @intCast(wip.string_bytes.items.len);
        try wip.string_bytes.ensureUnusedCapacity(gpa, s.len + 1);
        wip.string_bytes.appendSliceAssumeCapacity(s);
        wip.string_bytes.appendAssumeCapacity(0);
        return index;
    }

    pub fn printString(wip: *Wip, comptime fmt: []const u8, args: anytype) Allocator.Error!String {
        const gpa = wip.gpa;
        const index: String = @intCast(wip.string_bytes.items.len);
        try wip.string_bytes.writer(gpa).print(fmt, args);
        try wip.string_bytes.append(gpa, 0);
        return index;
    }

    pub fn addRootErrorMessage(wip: *Wip, em: ErrorMessage) Allocator.Error!void {
        try wip.root_list.ensureUnusedCapacity(wip.gpa, 1);
        wip.root_list.appendAssumeCapacity(try addErrorMessage(wip, em));
    }

    pub fn addErrorMessage(wip: *Wip, em: ErrorMessage) Allocator.Error!MessageIndex {
        return @enumFromInt(try addExtra(wip, em));
    }

    pub fn addErrorMessageAssumeCapacity(wip: *Wip, em: ErrorMessage) MessageIndex {
        return @enumFromInt(addExtraAssumeCapacity(wip, em));
    }

    pub fn addSourceLocation(wip: *Wip, sl: SourceLocation) Allocator.Error!SourceLocationIndex {
        return @enumFromInt(try addExtra(wip, sl));
    }

    pub fn addReferenceTrace(wip: *Wip, rt: ReferenceTrace) Allocator.Error!void {
        _ = try addExtra(wip, rt);
    }

    pub fn addBundleAsNotes(wip: *Wip, other: ErrorBundle) Allocator.Error!void {
        const gpa = wip.gpa;

        try wip.string_bytes.ensureUnusedCapacity(gpa, other.string_bytes.len);
        try wip.extra.ensureUnusedCapacity(gpa, other.extra.len);

        const other_list = other.getMessages();

        // The ensureUnusedCapacity call above guarantees this.
        const notes_start = wip.reserveNotes(@intCast(other_list.len)) catch unreachable;
        for (notes_start.., other_list) |note, message| {
            // This line can cause `wip.extra.items` to be resized.
            const note_index = @intFromEnum(wip.addOtherMessage(other, message) catch unreachable);
            wip.extra.items[note] = note_index;
        }
    }

    pub fn addBundleAsRoots(wip: *Wip, other: ErrorBundle) !void {
        const gpa = wip.gpa;

        try wip.string_bytes.ensureUnusedCapacity(gpa, other.string_bytes.len);
        try wip.extra.ensureUnusedCapacity(gpa, other.extra.len);

        const other_list = other.getMessages();

        try wip.root_list.ensureUnusedCapacity(gpa, other_list.len);
        for (other_list) |other_msg| {
            // The ensureUnusedCapacity calls above guarantees this.
            wip.root_list.appendAssumeCapacity(wip.addOtherMessage(other, other_msg) catch unreachable);
        }
    }

    pub fn reserveNotes(wip: *Wip, notes_len: u32) !u32 {
        try wip.extra.ensureUnusedCapacity(wip.gpa, notes_len +
            notes_len * @typeInfo(ErrorBundle.ErrorMessage).@"struct".fields.len);
        wip.extra.items.len += notes_len;
        return @intCast(wip.extra.items.len - notes_len);
    }

    pub fn addZirErrorMessages(
        eb: *ErrorBundle.Wip,
        zir: std.zig.Zir,
        tree: std.zig.Ast,
        source: [:0]const u8,
        src_path: []const u8,
    ) !void {
        const Zir = std.zig.Zir;
        const payload_index = zir.extra[@intFromEnum(Zir.ExtraIndex.compile_errors)];
        assert(payload_index != 0);

        const header = zir.extraData(Zir.Inst.CompileErrors, payload_index);
        const items_len = header.data.items_len;
        var extra_index = header.end;
        for (0..items_len) |_| {
            const item = zir.extraData(Zir.Inst.CompileErrors.Item, extra_index);
            extra_index = item.end;
            const err_span = blk: {
                if (item.data.node != 0) {
                    break :blk tree.nodeToSpan(item.data.node);
                }
                const token_starts = tree.tokens.items(.start);
                const start = token_starts[item.data.token] + item.data.byte_offset;
                const end = start + @as(u32, @intCast(tree.tokenSlice(item.data.token).len)) - item.data.byte_offset;
                break :blk std.zig.Ast.Span{ .start = start, .end = end, .main = start };
            };
            const err_loc = std.zig.findLineColumn(source, err_span.main);

            {
                const msg = zir.nullTerminatedString(item.data.msg);
                try eb.addRootErrorMessage(.{
                    .msg = try eb.addString(msg),
                    .src_loc = try eb.addSourceLocation(.{
                        .src_path = try eb.addString(src_path),
                        .span_start = err_span.start,
                        .span_main = err_span.main,
                        .span_end = err_span.end,
                        .line = @intCast(err_loc.line),
                        .column = @intCast(err_loc.column),
                        .source_line = try eb.addString(err_loc.source_line),
                    }),
                    .notes_len = item.data.notesLen(zir),
                });
            }

            if (item.data.notes != 0) {
                const notes_start = try eb.reserveNotes(item.data.notesLen(zir));
                const block = zir.extraData(Zir.Inst.Block, item.data.notes);
                const body = zir.extra[block.end..][0..block.data.body_len];
                for (notes_start.., body) |note_i, body_elem| {
                    const note_item = zir.extraData(Zir.Inst.CompileErrors.Item, body_elem);
                    const msg = zir.nullTerminatedString(note_item.data.msg);
                    const span = blk: {
                        if (note_item.data.node != 0) {
                            break :blk tree.nodeToSpan(note_item.data.node);
                        }
                        const token_starts = tree.tokens.items(.start);
                        const start = token_starts[note_item.data.token] + note_item.data.byte_offset;
                        const end = start + @as(u32, @intCast(tree.tokenSlice(note_item.data.token).len)) - item.data.byte_offset;
                        break :blk std.zig.Ast.Span{ .start = start, .end = end, .main = start };
                    };
                    const loc = std.zig.findLineColumn(source, span.main);

                    // This line can cause `wip.extra.items` to be resized.
                    const note_index = @intFromEnum(try eb.addErrorMessage(.{
                        .msg = try eb.addString(msg),
                        .src_loc = try eb.addSourceLocation(.{
                            .src_path = try eb.addString(src_path),
                            .span_start = span.start,
                            .span_main = span.main,
                            .span_end = span.end,
                            .line = @intCast(loc.line),
                            .column = @intCast(loc.column),
                            .source_line = if (loc.eql(err_loc))
                                0
                            else
                                try eb.addString(loc.source_line),
                        }),
                        .notes_len = 0, // TODO rework this function to be recursive
                    }));
                    eb.extra.items[note_i] = note_index;
                }
            }
        }
    }

    pub fn addZoirErrorMessages(
        eb: *ErrorBundle.Wip,
        zoir: std.zig.Zoir,
        tree: std.zig.Ast,
        source: [:0]const u8,
        src_path: []const u8,
    ) !void {
        assert(zoir.hasCompileErrors());

        for (zoir.compile_errors) |err| {
            const err_span: std.zig.Ast.Span = span: {
                if (err.token == std.zig.Zoir.CompileError.invalid_token) {
                    break :span tree.nodeToSpan(err.node_or_offset);
                }
                const token_start = tree.tokens.items(.start)[err.token];
                const start = token_start + err.node_or_offset;
                const end = token_start + @as(u32, @intCast(tree.tokenSlice(err.token).len));
                break :span .{ .start = start, .end = end, .main = start };
            };
            const err_loc = std.zig.findLineColumn(source, err_span.main);

            try eb.addRootErrorMessage(.{
                .msg = try eb.addString(err.msg.get(zoir)),
                .src_loc = try eb.addSourceLocation(.{
                    .src_path = try eb.addString(src_path),
                    .span_start = err_span.start,
                    .span_main = err_span.main,
                    .span_end = err_span.end,
                    .line = @intCast(err_loc.line),
                    .column = @intCast(err_loc.column),
                    .source_line = try eb.addString(err_loc.source_line),
                }),
                .notes_len = err.note_count,
            });

            const notes_start = try eb.reserveNotes(err.note_count);
            for (notes_start.., err.first_note.., 0..err.note_count) |eb_note_idx, zoir_note_idx, _| {
                const note = zoir.error_notes[zoir_note_idx];
                const note_span: std.zig.Ast.Span = span: {
                    if (note.token == std.zig.Zoir.CompileError.invalid_token) {
                        break :span tree.nodeToSpan(note.node_or_offset);
                    }
                    const token_start = tree.tokens.items(.start)[note.token];
                    const start = token_start + note.node_or_offset;
                    const end = token_start + @as(u32, @intCast(tree.tokenSlice(note.token).len));
                    break :span .{ .start = start, .end = end, .main = start };
                };
                const note_loc = std.zig.findLineColumn(source, note_span.main);

                // This line can cause `wip.extra.items` to be resized.
                const note_index = @intFromEnum(try eb.addErrorMessage(.{
                    .msg = try eb.addString(note.msg.get(zoir)),
                    .src_loc = try eb.addSourceLocation(.{
                        .src_path = try eb.addString(src_path),
                        .span_start = note_span.start,
                        .span_main = note_span.main,
                        .span_end = note_span.end,
                        .line = @intCast(note_loc.line),
                        .column = @intCast(note_loc.column),
                        .source_line = if (note_loc.eql(err_loc))
                            0
                        else
                            try eb.addString(note_loc.source_line),
                    }),
                    .notes_len = 0,
                }));
                eb.extra.items[eb_note_idx] = note_index;
            }
        }
    }

    fn addOtherMessage(wip: *Wip, other: ErrorBundle, msg_index: MessageIndex) !MessageIndex {
        const other_msg = other.getErrorMessage(msg_index);
        const src_loc = try wip.addOtherSourceLocation(other, other_msg.src_loc);
        const msg = try wip.addErrorMessage(.{
            .msg = try wip.addString(other.nullTerminatedString(other_msg.msg)),
            .count = other_msg.count,
            .src_loc = src_loc,
            .notes_len = other_msg.notes_len,
        });
        const notes_start = try wip.reserveNotes(other_msg.notes_len);
        for (notes_start.., other.getNotes(msg_index)) |note, other_note| {
            wip.extra.items[note] = @intFromEnum(try wip.addOtherMessage(other, other_note));
        }
        return msg;
    }

    fn addOtherSourceLocation(
        wip: *Wip,
        other: ErrorBundle,
        index: SourceLocationIndex,
    ) !SourceLocationIndex {
        if (index == .none) return .none;
        const other_sl = other.getSourceLocation(index);

        var ref_traces: std.ArrayListUnmanaged(ReferenceTrace) = .empty;
        defer ref_traces.deinit(wip.gpa);

        if (other_sl.reference_trace_len > 0) {
            var ref_index = other.extraData(SourceLocation, @intFromEnum(index)).end;
            for (0..other_sl.reference_trace_len) |_| {
                const other_ref_trace_ed = other.extraData(ReferenceTrace, ref_index);
                const other_ref_trace = other_ref_trace_ed.data;
                ref_index = other_ref_trace_ed.end;

                const ref_trace: ReferenceTrace = if (other_ref_trace.src_loc == .none) .{
                    // sentinel ReferenceTrace does not store a string index in decl_name
                    .decl_name = other_ref_trace.decl_name,
                    .src_loc = .none,
                } else .{
                    .decl_name = try wip.addString(other.nullTerminatedString(other_ref_trace.decl_name)),
                    .src_loc = try wip.addOtherSourceLocation(other, other_ref_trace.src_loc),
                };
                try ref_traces.append(wip.gpa, ref_trace);
            }
        }

        const src_loc = try wip.addSourceLocation(.{
            .src_path = try wip.addString(other.nullTerminatedString(other_sl.src_path)),
            .line = other_sl.line,
            .column = other_sl.column,
            .span_start = other_sl.span_start,
            .span_main = other_sl.span_main,
            .span_end = other_sl.span_end,
            .source_line = if (other_sl.source_line != 0)
                try wip.addString(other.nullTerminatedString(other_sl.source_line))
            else
                0,
            .reference_trace_len = other_sl.reference_trace_len,
        });

        for (ref_traces.items) |ref_trace| {
            try wip.addReferenceTrace(ref_trace);
        }

        return src_loc;
    }

    fn addExtra(wip: *Wip, extra: anytype) Allocator.Error!u32 {
        const gpa = wip.gpa;
        const fields = @typeInfo(@TypeOf(extra)).@"struct".fields;
        try wip.extra.ensureUnusedCapacity(gpa, fields.len);
        return addExtraAssumeCapacity(wip, extra);
    }

    fn addExtraAssumeCapacity(wip: *Wip, extra: anytype) u32 {
        const fields = @typeInfo(@TypeOf(extra)).@"struct".fields;
        const result: u32 = @intCast(wip.extra.items.len);
        wip.extra.items.len += fields.len;
        setExtra(wip, result, extra);
        return result;
    }

    fn setExtra(wip: *Wip, index: usize, extra: anytype) void {
        const fields = @typeInfo(@TypeOf(extra)).@"struct".fields;
        var i = index;
        inline for (fields) |field| {
            wip.extra.items[i] = switch (field.type) {
                u32 => @field(extra, field.name),
                MessageIndex => @intFromEnum(@field(extra, field.name)),
                SourceLocationIndex => @intFromEnum(@field(extra, field.name)),
                else => @compileError("bad field type"),
            };
            i += 1;
        }
    }

    test addBundleAsRoots {
        var bundle = bundle: {
            var wip: ErrorBundle.Wip = undefined;
            try wip.init(std.testing.allocator);
            errdefer wip.deinit();

            var ref_traces: [3]ReferenceTrace = undefined;
            for (&ref_traces, 0..) |*ref_trace, i| {
                if (i == ref_traces.len - 1) {
                    // sentinel reference trace
                    ref_trace.* = .{
                        .decl_name = 3, // signifies 3 hidden references
                        .src_loc = .none,
                    };
                } else {
                    ref_trace.* = .{
                        .decl_name = try wip.addString("foo"),
                        .src_loc = try wip.addSourceLocation(.{
                            .src_path = try wip.addString("foo"),
                            .line = 1,
                            .column = 2,
                            .span_start = 3,
                            .span_main = 4,
                            .span_end = 5,
                            .source_line = 0,
                        }),
                    };
                }
            }

            const src_loc = try wip.addSourceLocation(.{
                .src_path = try wip.addString("foo"),
                .line = 1,
                .column = 2,
                .span_start = 3,
                .span_main = 4,
                .span_end = 5,
                .source_line = try wip.addString("some source code"),
                .reference_trace_len = ref_traces.len,
            });
            for (&ref_traces) |ref_trace| {
                try wip.addReferenceTrace(ref_trace);
            }

            try wip.addRootErrorMessage(ErrorMessage{
                .msg = try wip.addString("hello world"),
                .src_loc = src_loc,
                .notes_len = 1,
            });
            const i = try wip.reserveNotes(1);
            const note_index = @intFromEnum(wip.addErrorMessageAssumeCapacity(.{
                .msg = try wip.addString("this is a note"),
                .src_loc = try wip.addSourceLocation(.{
                    .src_path = try wip.addString("bar"),
                    .line = 1,
                    .column = 2,
                    .span_start = 3,
                    .span_main = 4,
                    .span_end = 5,
                    .source_line = try wip.addString("another line of source"),
                }),
            }));
            wip.extra.items[i] = note_index;

            break :bundle try wip.toOwnedBundle("");
        };
        defer bundle.deinit(std.testing.allocator);

        const ttyconf: std.io.tty.Config = .no_color;

        var bundle_buf = std.ArrayList(u8).init(std.testing.allocator);
        defer bundle_buf.deinit();
        try bundle.renderToWriter(.{ .ttyconf = ttyconf }, bundle_buf.writer());

        var copy = copy: {
            var wip: ErrorBundle.Wip = undefined;
            try wip.init(std.testing.allocator);
            errdefer wip.deinit();

            try wip.addBundleAsRoots(bundle);

            break :copy try wip.toOwnedBundle("");
        };
        defer copy.deinit(std.testing.allocator);

        var copy_buf = std.ArrayList(u8).init(std.testing.allocator);
        defer copy_buf.deinit();
        try copy.renderToWriter(.{ .ttyconf = ttyconf }, copy_buf.writer());

        try std.testing.expectEqualStrings(bundle_buf.items, copy_buf.items);
    }
}