const std = @import("std");
const Compress = std.compress.flate.Compress;
const Dir = std.Io.Dir;
const SIZE_ALIGN: usize = 0x200;
pub fn main(init: std.process.Init) !void {
    const args = try init.minimal.args.toSlice(init.arena.allocator());

    if (args.len != 4) {
        std.debug.print("Usage: {s} <kernel_bin> <output_file> <assembly_file>\n", .{args[0]});
        std.process.exit(1);
    }
    const kernel_bin_path = args[1];
    const output_path = args[2];
    const assembly_path = args[3];
    std.debug.print("Compressing kernel: {s} -> {s}\n", .{kernel_bin_path, output_path});

    const kernel_file = try Dir.cwd().openFile(init.io, kernel_bin_path, .{});
    defer kernel_file.close(init.io);
    
    const compressed_file = try Dir.cwd().createFile(init.io, output_path, .{ .truncate = true });
    defer compressed_file.close(init.io);

    var read_buffer: [4096]u8 = undefined;
    var write_buffer: [4096]u8 = undefined;
    var work_buffer: [std.compress.flate.max_window_len]u8 = undefined;
    
    var writer = compressed_file.writer(init.io, &write_buffer);
    var reader = kernel_file.reader(init.io, &[0]u8{});
    var compressor = try Compress.init(&writer.interface, &work_buffer, .gzip, .best);
    var compressor_writer = &compressor.writer;
    var total_read: usize = 0;
    while (true) {
        const read_bytes = try reader.interface.readSliceShort(&read_buffer);
        if (read_bytes == 0) break;
        total_read += read_bytes;
        try compressor_writer.writeAll(read_buffer[0..read_bytes]);
    }
    try compressor_writer.flush();
    const aligned_size = (total_read + SIZE_ALIGN - 1) & ~(SIZE_ALIGN - 1);
    var assembly_file = try Dir.cwd().createFile(init.io, assembly_path, .{ .truncate = true });
    defer assembly_file.close(init.io);
    var assembly_writer = assembly_file.writer(init.io, &[0]u8{});
    try assembly_writer.interface.print(
        \\/* Auto-generated compressed kernel data */
        \\    .section .rodata.compressed,"a"
        \\    .global compressed_kernel_data
        \\compressed_kernel_data:
        \\    .incbin "{s}"
        \\compressed_end:
        \\    .set compressed_kernel_size, compressed_end - compressed_kernel_data
        \\    .set uncompressed_kernel_size, {d}
        \\    .global compressed_kernel_size
        \\    .global uncompressed_kernel_size
    , .{output_path, aligned_size});
    try assembly_writer.interface.flush();
    std.debug.print("Uncompressed aligned size: 0x{x} bytes\n", .{aligned_size});
}