const std = @import("std");

// A single message in the conversation history. We own all string memory.
const Message = struct {
    role: []const u8,
    content: ?[]const u8, // null for assistant messages that only have tool_calls
    // For tool result messages
    tool_call_id: ?[]const u8 = null,
    // For assistant messages that contain tool calls (raw JSON value kept so
    // we can re-serialize it verbatim)
    tool_calls_raw: ?std.json.Value = null,

    allocator: std.mem.Allocator,

    fn deinit(self: *Message) void {
        self.allocator.free(self.role);
        if (self.content) |c| self.allocator.free(c);
        if (self.tool_call_id) |id| self.allocator.free(id);
    }
};

// Serialize the conversation history into the JSON messages array.
fn writeMessages(jw: *std.json.Stringify, messages: []const Message) !void {
    try jw.beginArray();
    for (messages) |msg| {
        try jw.beginObject();
        try jw.objectField("role");
        try jw.write(msg.role);

        if (msg.tool_call_id) |id| {
            // tool result message
            try jw.objectField("tool_call_id");
            try jw.write(id);
        }

        if (msg.tool_calls_raw) |tc| {
            // assistant message with tool_calls: write content as null
            try jw.objectField("content");
            try jw.write(null);
            try jw.objectField("tool_calls");
            try jw.write(tc);
        } else {
            try jw.objectField("content");
            if (msg.content) |c| {
                try jw.write(c);
            } else {
                try jw.write(null);
            }
        }

        try jw.endObject();
    }
    try jw.endArray();
}

pub fn main() !void {
    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    defer _ = gpa.deinit();
    const allocator = gpa.allocator();

    // Parse -p flag
    const args = try std.process.argsAlloc(allocator);
    defer std.process.argsFree(allocator, args);

    var prompt: ?[]const u8 = null;
    var i: usize = 1;
    while (i < args.len) : (i += 1) {
        if (std.mem.eql(u8, args[i], "-p") and i + 1 < args.len) {
            i += 1;
            prompt = args[i];
        }
    }

    const prompt_str = prompt orelse @panic("Prompt must not be empty");

    const api_key = std.posix.getenv("OPENROUTER_API_KEY") orelse @panic("OPENROUTER_API_KEY is not set");
    const base_url = std.posix.getenv("OPENROUTER_BASE_URL") orelse "https://openrouter.ai/api/v1";

    const url_str = try std.fmt.allocPrint(allocator, "{s}/chat/completions", .{base_url});
    defer allocator.free(url_str);

    const auth_value = try std.fmt.allocPrint(allocator, "Bearer {s}", .{api_key});
    defer allocator.free(auth_value);

    // Conversation history — seed with the user prompt.
    // In Zig 0.15 ArrayList is unmanaged: init via .{}, pass allocator to each op.
    var messages: std.ArrayList(Message) = .{};
    defer {
        for (messages.items) |*m| m.deinit();
        messages.deinit(allocator);
    }
    try messages.append(allocator, .{
        .role = try allocator.dupe(u8, "user"),
        .content = try allocator.dupe(u8, prompt_str),
        .allocator = allocator,
    });

    var client: std.http.Client = .{ .allocator = allocator };
    defer client.deinit();

    // Keep parsed responses alive so tool_calls_raw values remain valid.
    var parsed_responses: std.ArrayList(std.json.Parsed(std.json.Value)) = .{};
    defer {
        for (parsed_responses.items) |*p| p.deinit();
        parsed_responses.deinit(allocator);
    }

    // ── Agent loop ────────────────────────────────────────────────────────────
    while (true) {
        // Build request body
        var body_out: std.io.Writer.Allocating = .init(allocator);
        defer body_out.deinit();
        var jw: std.json.Stringify = .{ .writer = &body_out.writer };

        try jw.beginObject();

        try jw.objectField("model");
        try jw.write("anthropic/claude-haiku-4.5");

        try jw.objectField("messages");
        try writeMessages(&jw, messages.items);

        try jw.objectField("tools");
        try jw.beginArray();
        // Read tool
        try jw.write(.{
            .type = "function",
            .function = .{
                .name = "Read",
                .description = "Read and return the contents of a file",
                .parameters = .{
                    .type = "object",
                    .properties = .{
                        .file_path = .{
                            .type = "string",
                            .description = "The path to the file to read",
                        },
                    },
                    .required = &[_][]const u8{"file_path"},
                },
            },
        });
        // Write tool
        try jw.write(.{
            .type = "function",
            .function = .{
                .name = "Write",
                .description = "Write content to a file",
                .parameters = .{
                    .type = "object",
                    .properties = .{
                        .file_path = .{
                            .type = "string",
                            .description = "The path of the file to write to",
                        },
                        .content = .{
                            .type = "string",
                            .description = "The content to write to the file",
                        },
                    },
                    .required = &[_][]const u8{ "file_path", "content" },
                },
            },
        });
        // Bash tool
        try jw.write(.{
            .type = "function",
            .function = .{
                .name = "Bash",
                .description = "Execute a shell command",
                .parameters = .{
                    .type = "object",
                    .properties = .{
                        .command = .{
                            .type = "string",
                            .description = "The command to execute",
                        },
                    },
                    .required = &[_][]const u8{"command"},
                },
            },
        });
        try jw.endArray();

        try jw.endObject();
        const body = body_out.written();

        // Send request
        var response_out: std.io.Writer.Allocating = .init(allocator);
        defer response_out.deinit();

        _ = try client.fetch(.{
            .location = .{ .url = url_str },
            .method = .POST,
            .payload = body,
            .extra_headers = &.{
                .{ .name = "content-type", .value = "application/json" },
                .{ .name = "authorization", .value = auth_value },
            },
            .response_writer = &response_out.writer,
        });
        const response_body = response_out.written();
        std.debug.print("Response: {s}\n", .{response_body});

        // Parse and keep alive for tool_calls_raw references
        const parsed = try std.json.parseFromSlice(std.json.Value, allocator, response_body, .{});
        try parsed_responses.append(allocator, parsed);
        const parsed_ptr = &parsed_responses.items[parsed_responses.items.len - 1];

        const choices = parsed_ptr.value.object.get("choices") orelse @panic("No choices in response");
        if (choices.array.items.len == 0) @panic("Empty choices array");

        const choice = choices.array.items[0];
        const finish_reason = choice.object.get("finish_reason").?.string;
        const message = choice.object.get("message").?;

        // Check whether this turn has tool calls
        const has_tool_calls = blk: {
            const tc = message.object.get("tool_calls") orelse break :blk false;
            break :blk tc != .null and tc.array.items.len > 0;
        };

        if (has_tool_calls) {
            const tc_val = message.object.get("tool_calls").?;
            try messages.append(allocator, .{
                .role = try allocator.dupe(u8, "assistant"),
                .content = null,
                .tool_calls_raw = tc_val,
                .allocator = allocator,
            });

            // Execute every tool call and append results
            for (tc_val.array.items) |tool_call| {
                const call_id = tool_call.object.get("id").?.string;
                const func = tool_call.object.get("function").?;
                const func_name = func.object.get("name").?.string;
                const arguments_str = func.object.get("arguments").?.string;

                std.debug.print("Tool call: {s}\n", .{func_name});

                var result_owned: []u8 = &.{};

                if (std.mem.eql(u8, func_name, "Read")) {
                    const args_parsed = try std.json.parseFromSlice(std.json.Value, allocator, arguments_str, .{});
                    defer args_parsed.deinit();
                    const file_path = args_parsed.value.object.get("file_path").?.string;
                    std.debug.print("Reading file: {s}\n", .{file_path});
                    result_owned = try std.fs.cwd().readFileAlloc(allocator, file_path, std.math.maxInt(usize));
                } else if (std.mem.eql(u8, func_name, "Write")) {
                    const args_parsed = try std.json.parseFromSlice(std.json.Value, allocator, arguments_str, .{});
                    defer args_parsed.deinit();
                    const file_path = args_parsed.value.object.get("file_path").?.string;
                    const content = args_parsed.value.object.get("content").?.string;
                    std.debug.print("Writing file: {s}\n", .{file_path});
                    // Create parent directories if needed
                    if (std.fs.path.dirname(file_path)) |dir| {
                        try std.fs.cwd().makePath(dir);
                    }
                    try std.fs.cwd().writeFile(.{ .sub_path = file_path, .data = content });
                    result_owned = try allocator.dupe(u8, "OK");
                } else if (std.mem.eql(u8, func_name, "Bash")) {
                    const args_parsed = try std.json.parseFromSlice(std.json.Value, allocator, arguments_str, .{});
                    defer args_parsed.deinit();
                    const command = args_parsed.value.object.get("command").?.string;
                    std.debug.print("Running command: {s}\n", .{command});
                    const result = try std.process.Child.run(.{
                        .allocator = allocator,
                        .argv = &.{ "sh", "-c", command },
                    });
                    defer allocator.free(result.stdout);
                    defer allocator.free(result.stderr);
                    // Combine stdout + stderr into the tool result
                    result_owned = try std.fmt.allocPrint(allocator, "{s}{s}", .{ result.stdout, result.stderr });
                }

                try messages.append(allocator, .{
                    .role = try allocator.dupe(u8, "tool"),
                    .content = result_owned,
                    .tool_call_id = try allocator.dupe(u8, call_id),
                    .allocator = allocator,
                });
            }
        } else {
            // No tool calls — record the assistant reply and stop
            const content_val = message.object.get("content") orelse std.json.Value{ .null = {} };
            const content_str = if (content_val == .string) content_val.string else "";
            try messages.append(allocator, .{
                .role = try allocator.dupe(u8, "assistant"),
                .content = try allocator.dupe(u8, content_str),
                .allocator = allocator,
            });
        }

        // Stop when the model signals it's done
        if (std.mem.eql(u8, finish_reason, "stop") or !has_tool_calls) {
            const last = messages.items[messages.items.len - 1];
            try std.fs.File.stdout().writeAll(last.content orelse "");
            break;
        }
    }
}