const std = @import("std");
const Allocator = std.mem.Allocator;

/// A streaming text buffer used by content blocks.
/// Thin alias over ArrayList(u8) — amortized O(1) appends,
/// no O(n²) re-copying.
pub const TextualBlock = std.ArrayList(u8);

/// Create a TextualBlock with initial content (copies the slice).
pub fn textualBlockFromSlice(alloc: Allocator, slice: []const u8) !TextualBlock {
    var buf: TextualBlock = .empty;
    try buf.appendSlice(alloc, slice);
    return buf;
}

/// A reasoning/thinking block from the assistant.
///
/// `text` is the streamed reasoning content. `signature` is the opaque
/// integrity token Anthropic emits with extended-thinking responses and
/// requires echoed back verbatim on follow-up turns. It is optional because
/// other providers (OpenAI-compatible APIs) do not produce it.
pub const ThinkingBlock = struct {
    text: TextualBlock = .empty,
    signature: ?[]const u8 = null,

    pub fn deinit(self: *ThinkingBlock, alloc: Allocator) void {
        self.text.deinit(alloc);
        if (self.signature) |sig| alloc.free(sig);
    }
};

pub const ToolUseBlock = struct {
    id: []const u8,
    name: []const u8,
    input: TextualBlock = .empty,

    pub fn deinit(self: *ToolUseBlock, alloc: Allocator) void {
        alloc.free(self.id);
        alloc.free(self.name);
        self.input.deinit(alloc);
    }
};

pub const ToolResultBlock = struct {
    tool_use_id: []const u8,
    content: TextualBlock = .empty,

    pub fn deinit(self: *ToolResultBlock, alloc: Allocator) void {
        alloc.free(self.tool_use_id);
        self.content.deinit(alloc);
    }
};

/// How a `.system` content block combines with the system text collected
/// before it. `append` adds to the running effective prompt; `replace`
/// discards everything collected so far and starts fresh from this block.
pub const SystemMode = enum { append, replace };

/// A system-prompt content block. System prompts remain `.system`-role
/// messages; this block records the mode that governs how its text folds
/// into the effective system prompt (see `effectiveSystemBlocks`).
pub const SystemBlock = struct {
    text: TextualBlock = .empty,
    mode: SystemMode = .append,

    pub fn deinit(self: *SystemBlock, alloc: Allocator) void {
        self.text.deinit(alloc);
    }
};

pub const ContentBlock = union(enum) {
    Text: TextualBlock,
    Thinking: ThinkingBlock,
    ToolUse: ToolUseBlock,
    ToolResult: ToolResultBlock,
    System: SystemBlock,

    pub fn deinit(self: *ContentBlock, alloc: Allocator) void {
        switch (self.*) {
            .Text => |*b| b.deinit(alloc),
            .Thinking => |*b| b.deinit(alloc),
            .ToolUse => |*b| b.deinit(alloc),
            .ToolResult => |*b| b.deinit(alloc),
            .System => |*b| b.deinit(alloc),
        }
    }
};

pub const MessageRole = enum {
    system,
    user,
    assistant,
};

pub const Message = struct {
    role: MessageRole,
    content: std.ArrayList(ContentBlock) = .empty,

    pub fn deinit(self: *Message, alloc: Allocator) void {
        for (self.content.items) |*block| {
            block.deinit(alloc);
        }
        self.content.deinit(alloc);
    }
};

pub const Conversation = struct {
    messages: std.ArrayList(Message) = .empty,
    allocator: Allocator,

    pub fn init(allocator: Allocator) Conversation {
        return .{
            .allocator = allocator,
        };
    }

    /// Append a system message in `append` mode. Adds to the effective
    /// system prompt. (Back-compatible: same external behavior as before
    /// the `.System` block existed.)
    pub fn addSystemMessage(self: *Conversation, text: []const u8) !void {
        return self.appendSystemBlock(text, .append);
    }

    /// Append a system message in `replace` mode. When the effective
    /// prompt is rebuilt (see `effectiveSystemBlocks`), this discards all
    /// prior system text and starts fresh.
    pub fn replaceSystemMessage(self: *Conversation, text: []const u8) !void {
        return self.appendSystemBlock(text, .replace);
    }

    /// Append a `.system`-role message whose single content block is a
    /// `.System` block carrying `mode`.
    fn appendSystemBlock(self: *Conversation, text: []const u8, mode: SystemMode) !void {
        const tb = try textualBlockFromSlice(self.allocator, text);
        var content: std.ArrayList(ContentBlock) = .empty;
        errdefer {
            for (content.items) |*b| b.deinit(self.allocator);
            content.deinit(self.allocator);
        }
        try content.append(self.allocator, .{ .System = .{ .text = tb, .mode = mode } });
        try self.messages.append(self.allocator, .{
            .role = .system,
            .content = content,
        });
    }

    pub fn addUserMessage(self: *Conversation, text: []const u8) !void {
        const tb = try textualBlockFromSlice(self.allocator, text);
        var content: std.ArrayList(ContentBlock) = .empty;
        try content.append(self.allocator, .{ .Text = tb });
        try self.messages.append(self.allocator, .{
            .role = .user,
            .content = content,
        });
    }

    /// Append an assistant message. Ownership of the blocks is transferred
    /// to the conversation; the caller must not deinit them after this call.
    pub fn addAssistantMessage(self: *Conversation, blocks: []const ContentBlock) !void {
        var content: std.ArrayList(ContentBlock) = .empty;
        try content.ensureTotalCapacity(self.allocator, blocks.len);
        for (blocks) |block| {
            content.appendAssumeCapacity(block);
        }
        try self.messages.append(self.allocator, .{
            .role = .assistant,
            .content = content,
        });
    }

    pub fn deinit(self: *Conversation) void {
        for (self.messages.items) |*msg| {
            msg.deinit(self.allocator);
        }
        self.messages.deinit(self.allocator);
    }
};

/// Derive the effective ordered list of system-text blocks from a slice of
/// messages. This is the single shared rule that governs both provider
/// serialization and session rebuild.
///
/// Walk the messages in order; for each `.system` message's `.System`
/// block:
///   - `append`: add the block's text to the running list.
///   - `replace`: clear the running list, then add this block's text.
///
/// The returned slices are **borrowed** from `messages` — valid only as
/// long as the underlying conversation is unmodified. The caller owns the
/// returned `ArrayList` itself and must `deinit` it (this frees the slice
/// storage, not the borrowed text).
///
/// Running this walk over a *prefix* of the messages reconstructs the
/// effective prompt as of that point — the `/tree` faithfulness property.
pub fn effectiveSystemBlocks(
    alloc: Allocator,
    messages: []const Message,
) !std.ArrayList([]const u8) {
    var out: std.ArrayList([]const u8) = .empty;
    errdefer out.deinit(alloc);
    for (messages) |msg| {
        if (msg.role != .system) continue;
        for (msg.content.items) |block| {
            switch (block) {
                .System => |sb| {
                    switch (sb.mode) {
                        .append => {},
                        .replace => out.clearRetainingCapacity(),
                    }
                    try out.append(alloc, sb.text.items);
                },
                // Be tolerant of plain `.Text` blocks on a system message
                // (e.g. hand-built test conversations): treat them as
                // append-mode text.
                .Text => |tb| try out.append(alloc, tb.items),
                else => {},
            }
        }
    }
    return out;
}

test "Conversation - add messages and verify content" {
    const allocator = std.testing.allocator;

    var conv = Conversation.init(allocator);
    defer conv.deinit();

    try conv.addSystemMessage("You are a helpful assistant.");
    try conv.addUserMessage("Hello!");
    try conv.addAssistantMessage(&.{
        .{ .Text = try textualBlockFromSlice(allocator, "Hi there!") },
    });

    try std.testing.expectEqual(@as(usize, 3), conv.messages.items.len);

    try std.testing.expectEqual(MessageRole.system, conv.messages.items[0].role);
    try std.testing.expectEqualStrings(
        "You are a helpful assistant.",
        conv.messages.items[0].content.items[0].System.text.items,
    );
    try std.testing.expectEqual(
        SystemMode.append,
        conv.messages.items[0].content.items[0].System.mode,
    );

    try std.testing.expectEqual(MessageRole.user, conv.messages.items[1].role);
    try std.testing.expectEqualStrings("Hello!", conv.messages.items[1].content.items[0].Text.items);

    try std.testing.expectEqual(MessageRole.assistant, conv.messages.items[2].role);
    try std.testing.expectEqualStrings("Hi there!", conv.messages.items[2].content.items[0].Text.items);
}

test "TextualBlock - incremental append" {
    const allocator = std.testing.allocator;

    var tb = TextualBlock.empty;
    defer tb.deinit(allocator);

    try tb.appendSlice(allocator, "Hello");
    try tb.appendSlice(allocator, " world");
    try std.testing.expectEqualStrings("Hello world", tb.items);
}

test "Conversation - deinit frees without leaks" {
    const allocator = std.testing.allocator;

    var conv = Conversation.init(allocator);
    try conv.addSystemMessage("system");
    try conv.addUserMessage("user message");
    try conv.addAssistantMessage(&.{
        .{ .Text = try textualBlockFromSlice(allocator, "response") },
    });
    conv.deinit();
}

test "ContentBlock - Thinking variant" {
    const allocator = std.testing.allocator;

    var conv = Conversation.init(allocator);
    defer conv.deinit();

    try conv.addAssistantMessage(&.{
        .{ .Thinking = .{ .text = try textualBlockFromSlice(allocator, "hmm...") } },
        .{ .Text = try textualBlockFromSlice(allocator, "answer") },
    });

    try std.testing.expectEqual(@as(usize, 2), conv.messages.items[0].content.items.len);
    try std.testing.expectEqualStrings("hmm...", conv.messages.items[0].content.items[0].Thinking.text.items);
    try std.testing.expectEqualStrings("answer", conv.messages.items[0].content.items[1].Text.items);
}

test "System block - addSystemMessage records append mode, replaceSystemMessage records replace mode" {
    const allocator = std.testing.allocator;

    var conv = Conversation.init(allocator);
    defer conv.deinit();

    try conv.addSystemMessage("base");
    try conv.replaceSystemMessage("fresh");

    try std.testing.expectEqual(SystemMode.append, conv.messages.items[0].content.items[0].System.mode);
    try std.testing.expectEqualStrings("base", conv.messages.items[0].content.items[0].System.text.items);
    try std.testing.expectEqual(SystemMode.replace, conv.messages.items[1].content.items[0].System.mode);
    try std.testing.expectEqualStrings("fresh", conv.messages.items[1].content.items[0].System.text.items);
}

test "effectiveSystemBlocks - append accumulates in order" {
    const allocator = std.testing.allocator;

    var conv = Conversation.init(allocator);
    defer conv.deinit();

    try conv.addSystemMessage("a");
    try conv.addSystemMessage("b");
    try conv.addUserMessage("hi");
    try conv.addSystemMessage("c");

    var blocks = try effectiveSystemBlocks(allocator, conv.messages.items);
    defer blocks.deinit(allocator);

    try std.testing.expectEqual(@as(usize, 3), blocks.items.len);
    try std.testing.expectEqualStrings("a", blocks.items[0]);
    try std.testing.expectEqualStrings("b", blocks.items[1]);
    try std.testing.expectEqualStrings("c", blocks.items[2]);
}

test "effectiveSystemBlocks - replace wipes everything collected so far" {
    const allocator = std.testing.allocator;

    var conv = Conversation.init(allocator);
    defer conv.deinit();

    try conv.addSystemMessage("a");
    try conv.addSystemMessage("b");
    try conv.replaceSystemMessage("fresh");
    try conv.addSystemMessage("after");

    var blocks = try effectiveSystemBlocks(allocator, conv.messages.items);
    defer blocks.deinit(allocator);

    try std.testing.expectEqual(@as(usize, 2), blocks.items.len);
    try std.testing.expectEqualStrings("fresh", blocks.items[0]);
    try std.testing.expectEqualStrings("after", blocks.items[1]);
}

test "effectiveSystemBlocks - prefix reconstructs prompt as of that point" {
    const allocator = std.testing.allocator;

    var conv = Conversation.init(allocator);
    defer conv.deinit();

    try conv.addSystemMessage("a");
    try conv.replaceSystemMessage("fresh");
    try conv.addSystemMessage("after");

    // Truncate at position 1 (only the first `addSystemMessage`).
    var blocks = try effectiveSystemBlocks(allocator, conv.messages.items[0..1]);
    defer blocks.deinit(allocator);
    try std.testing.expectEqual(@as(usize, 1), blocks.items.len);
    try std.testing.expectEqualStrings("a", blocks.items[0]);
}