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);
    }
};

/// A binary attachment stored in a tool result. `media_type` is owned
/// (e.g. "image/png"); `data` is the streaming text buffer holding
/// base64-encoded bytes.
pub const StoredMediaPart = struct {
    media_type: []const u8,
    data: TextualBlock = .empty,

    pub fn deinit(self: *StoredMediaPart, alloc: Allocator) void {
        alloc.free(self.media_type);
        self.data.deinit(alloc);
    }
};

/// One part of a tool result as stored in the conversation. Text uses the
/// streaming `TextualBlock` form to preserve incremental-append semantics.
pub const ResultPartStored = union(enum) {
    text: TextualBlock,
    media: StoredMediaPart,

    pub fn deinit(self: *ResultPartStored, alloc: Allocator) void {
        switch (self.*) {
            .text => |*t| t.deinit(alloc),
            .media => |*m| m.deinit(alloc),
        }
    }
};

pub const ToolResultBlock = struct {
    tool_use_id: []const u8,
    parts: std.ArrayList(ResultPartStored) = .empty,
    /// True when this result reports a tool failure rather than success.
    /// Serialized to providers that support it (Anthropic's `is_error`);
    /// recorded but unserialized for OpenAI Chat (the error text in `parts`
    /// carries the signal there). Defaults to false for backward
    /// compatibility with sessions written before this field existed.
    is_error: bool = false,

    pub fn deinit(self: *ToolResultBlock, alloc: Allocator) void {
        alloc.free(self.tool_use_id);
        for (self.parts.items) |*p| p.deinit(alloc);
        self.parts.deinit(alloc);
    }

    /// Concatenate all text parts into `out`. Media parts are skipped.
    /// Used by callers (compaction, OpenAI fan-out) that want only the
    /// textual portion.
    pub fn appendTextInto(self: ToolResultBlock, alloc: Allocator, out: *TextualBlock) !void {
        for (self.parts.items) |p| {
            if (p == .text) try out.appendSlice(alloc, p.text.items);
        }
    }

    /// True if any part is a media attachment.
    pub fn hasMedia(self: ToolResultBlock) bool {
        for (self.parts.items) |p| {
            if (p == .media) return true;
        }
        return false;
    }
};

/// 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);
    }
};

/// A compaction summary block. Carries the synthetic seed text that
/// stands in for a compacted conversation prefix. Like a `replace`-mode
/// `.System` block, it changes replay semantics: when the effective
/// conversation is rebuilt, a `.CompactionSummary` block resets all prior
/// *conversation* turns (user/assistant), and only the latest summary
/// plus the messages after it contribute to active context. System blocks
/// are unaffected — they are derived separately by `effectiveSystemBlocks`
/// and survive compaction untouched.
///
/// A `.CompactionSummary` block sits alone in a `user`-role message.
pub const CompactionSummaryBlock = struct {
    text: TextualBlock = .empty,

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

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

    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),
            .CompactionSummary => |*b| b.deinit(alloc),
        }
    }
};

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

/// Token usage reported by a provider for a single assistant turn.
///
/// All input categories sum to the total prompt tokens billed for the turn:
///   `input + cache_read + cache_write` = total prompt tokens.
///
/// Crucially, `input` is *cumulative*: it is the entire prompt sent for the
/// turn (the whole prior conversation), not the size of the turn's own
/// message. Compaction relies on this to back out per-turn sizes from
/// neighbor deltas (see `compaction.computeSplit`).
///
/// `reasoning` is a **subset** of `output`, not additive — the portion of
/// the output spent on internal reasoning (OpenAI o-series, Anthropic
/// extended thinking). Cost is computed from `output` only; `reasoning` is
/// tracked for display.
///
/// All fields are `u64` and default to 0. Providers that don't report a
/// given category leave it at 0.
///
/// Defined here (rather than in `session.zig`) so in-memory `Message`s can
/// carry usage without a module cycle; `session.zig` re-exports it.
pub const Usage = struct {
    /// Fresh input tokens billed at the model's base input rate.
    input: u64 = 0,
    /// Output tokens billed at the model's base output rate.
    output: u64 = 0,
    /// Input tokens served from cache (typ. 0.1× base input rate on
    /// Anthropic, 0.5× on OpenAI).
    cache_read: u64 = 0,
    /// Input tokens written to a new cache entry (Anthropic only at
    /// time of writing; typ. 1.25× base input rate). OpenAI's cache
    /// hits don't bill a write premium, so this stays 0 there.
    cache_write: u64 = 0,
    /// Subset of `output` spent on reasoning. For OpenAI o-series
    /// models and Anthropic extended thinking. Display-only — cost is
    /// already accounted for via `output`.
    reasoning: u64 = 0,
};

pub const Message = struct {
    role: MessageRole,
    content: std.ArrayList(ContentBlock) = .empty,
    /// Provider-reported token usage for this message's turn. Set on
    /// assistant messages (live, via the provider; replayed, from disk);
    /// null on user/system messages and when the provider emitted no usage.
    /// Used by compaction to size the retention window.
    usage: ?Usage = null,
    /// Opaque per-message metadata bag. `libpanto` never interprets these
    /// bytes; the documented contract is that, when present, they are valid
    /// JSON (so a store may keep them as a JSON column and tools may
    /// deserialize them). Round-trips through persistence: set before a turn
    /// commits, read back off the `Message` after `load`. Borrowed; owned by
    /// whoever set it (the conversation's allocator on the replay path).
    metadata: ?[]const u8 = null,

    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 a `user`-role message carrying a single `.CompactionSummary`
    /// block holding `text` (copied). This is the seed that stands in for
    /// a compacted conversation prefix; see `CompactionSummaryBlock`.
    pub fn addCompactionSummary(self: *Conversation, text: []const u8) !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, .{ .CompactionSummary = .{ .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 {
        return self.addAssistantMessageWithUsage(blocks, null);
    }

    /// Append an assistant message tagged with its provider-reported
    /// `usage`. Ownership of the blocks is transferred to the conversation.
    pub fn addAssistantMessageWithUsage(
        self: *Conversation,
        blocks: []const ContentBlock,
        usage: ?Usage,
    ) !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,
            .usage = usage,
        });
    }

    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;
}

/// Index of the message carrying the latest `.CompactionSummary` block,
/// or null if the conversation has never been compacted.
///
/// Compaction replay is reset-like: only the latest compaction summary and
/// the messages after it contribute to active conversation context. This
/// returns the anchor message (the summary itself); active conversation is
/// `messages[anchor..]`. System messages are unaffected and are derived
/// independently via `effectiveSystemBlocks`.
pub fn latestCompactionIndex(messages: []const Message) ?usize {
    var anchor: ?usize = null;
    for (messages, 0..) |msg, i| {
        for (msg.content.items) |block| {
            if (block == .CompactionSummary) {
                anchor = i;
                break;
            }
        }
    }
    return anchor;
}

/// The active (post-compaction) conversation message window that should be
/// sent to a provider. If the conversation has been compacted, this is the
/// latest compaction summary message plus everything after it; otherwise
/// it is the whole message list. System messages within the window are
/// still emitted/handled by the caller's own role filtering — this only
/// trims the *prefix* superseded by compaction.
///
/// The returned slice borrows from `messages`.
pub fn activeMessageWindow(messages: []const Message) []const Message {
    if (latestCompactionIndex(messages)) |anchor| {
        return messages[anchor..];
    }
    return messages;
}

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

test "addCompactionSummary - sits alone in a user message" {
    const allocator = std.testing.allocator;

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

    try conv.addUserMessage("hi");
    try conv.addCompactionSummary("summary of earlier history");

    const m = conv.messages.items[1];
    try std.testing.expectEqual(MessageRole.user, m.role);
    try std.testing.expectEqual(@as(usize, 1), m.content.items.len);
    try std.testing.expectEqualStrings(
        "summary of earlier history",
        m.content.items[0].CompactionSummary.text.items,
    );
}

test "latestCompactionIndex - null when never compacted" {
    const allocator = std.testing.allocator;

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

    try conv.addSystemMessage("sys");
    try conv.addUserMessage("hi");
    try conv.addAssistantMessage(&.{
        .{ .Text = try textualBlockFromSlice(allocator, "hello") },
    });

    try std.testing.expect(latestCompactionIndex(conv.messages.items) == null);
}

test "latestCompactionIndex - returns the latest summary anchor" {
    const allocator = std.testing.allocator;

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

    try conv.addSystemMessage("sys");
    try conv.addUserMessage("old");
    try conv.addCompactionSummary("S1");
    try conv.addUserMessage("mid");
    try conv.addCompactionSummary("S2"); // index 4
    try conv.addUserMessage("recent");

    try std.testing.expectEqual(@as(?usize, 4), latestCompactionIndex(conv.messages.items));
}