const std = @import("std");

// pub fn serialize(writer: *BitWriter, allocator: std.mem.Allocator, comptime T: type, value: T) !void {
//     const info = @typeInfo(T);

//     switch (info) {
//         .int, .bool => {
//             try writer.write(allocator, T, value);
//         },

//         .@"struct" => |s| {
//             if (!s.is_packed)
//                 @compileError("serialize only supports packed structs (got " ++ @typeName(T) ++ ")");

//             inline for (s.fields) |field| {
//                 const field_value = @field(value, field.name);
//                 try serialize(writer, allocator, field_value);
//             }
//         },

//         .@"union" => |u| {
//             if (u.layout != .@"packed")
//                 @compileError("serialize only supports packed structs (got " ++ @typeName(T) ++ ")");

//             inline for (u.fields) |field| {
//                 const field_value = @field(value, field.name);
//                 try serialize(writer, allocator, field_value);
//             }
//         },

//         else => @compileError("Unsupported type in serialize: " ++ @typeName(T)),
//     }
// }


pub const BitWriter = struct {
    // We store a pointer to the ArrayList so we can grow it
    bytes: *std.ArrayList(u8),
    bit_offset: usize = 0,

    pub fn init(bytes: *std.ArrayList(u8)) BitWriter {
        return .{ .bytes = bytes };
    }

    pub fn serialize(self: *BitWriter, allocator: std.mem.Allocator, value: anytype) !void {
        const T = @TypeOf(value);
        const info = @typeInfo(T);

        switch (info) {
            .@"struct" => |s| {
                inline for (s.fields) |f|
                    try self.serialize(allocator, @field(value, f.name));
            },
            .@"union" => |u| {
                const E = u.tag_type orelse @compileError(std.fmt.comptimePrint("Union {s} has to tag type.", .{@typeName(T)}));
                try self.serialize(allocator, @as(E, value));
                switch (value) {
                    inline else => |f| try self.serialize(allocator, f),
                }
            },
            .@"enum" => try self.serialize(allocator, @intFromEnum(value)),
            .pointer => |p| switch (p.size) {
                .one => try self.serialize(allocator, value.*),
                .slice => {
                    try self.serialize(allocator, value.len);
                    for (value) |e|
                        try self.serialize(allocator, e);
                },
                .many, .c => @compileError(std.fmt.comptimePrint("Unsupported type: {}", .{T})),
            },
            .int, .bool, .float, => try self.write(allocator, value),
            .void => {},
            .array => |a| if (@sizeOf(a.child) == 1) {
                try self.write(allocator, @ptrCast(&value));
            } else {
                for (value) |e|
                    try self.serialize(allocator, e);
            },
            .optional => if (value) |_t| {
                try self.serialize(allocator, true);
                try self.serialize(allocator, _t);
            } else {
                try self.serialize(allocator, false);
            },
            .vector => |v| try self.serialize(allocator, @as([v.len]v.child, value)),
            else => @compileError(std.fmt.comptimePrint("Unsupported type: {}", .{T})),
        }
    }

    fn write(self: *BitWriter, allocator: std.mem.Allocator, value: anytype) !void {
        const T = @TypeOf(value);
        const bit_width = @bitSizeOf(T);

        const UnsignedT = std.meta.Int(.unsigned, bit_width);
        const bits_to_write = @as(UnsignedT, @bitCast(value));

        var remaining : usize = bit_width;
        while (remaining > 0) {
            const byte_index = self.bit_offset / 8;
            const bit_in_byte = self.bit_offset % 8;

            // If we are moving into a new byte, grow the ArrayList
            if (byte_index >= self.bytes.items.len) {
                try self.bytes.append(allocator, 0);
            }

            const space_in_byte = 8 - bit_in_byte;
            const chunk_size = @min(remaining, space_in_byte);

            const mask : u8 = std.math.shr(u8, 0xFF, 8 - chunk_size);

            const bits : u8 = @truncate(std.math.shr(UnsignedT, bits_to_write, remaining - chunk_size) & mask);

            self.bytes.items[byte_index] |= std.math.shl(u8, bits, space_in_byte - chunk_size);

            self.bit_offset += chunk_size;
            remaining -= chunk_size;
        }
    }

    pub fn written(self: *BitWriter) []const u8 {
        const bytes_used = (self.bit_offset + 7) / 8; // round up bits → bytes
        return self.bytes.items[0..bytes_used];
    }
};

pub const BitReader = struct {
    bytes: []const u8,
    bit_offset: usize = 0,

    pub fn init(bytes: []const u8) BitReader {
        return .{ .bytes = bytes };
    }

    pub fn deserialize(self: *BitReader, allocator: std.mem.Allocator, T: type) !T {
        return switch (@typeInfo(T)) {
            .@"struct" => |s| b: {
                var data: T = undefined;
                inline for (s.fields) |f|
                    @field(data, f.name) = try self.deserialize(allocator, f.type);
                break :b data;
            },
            .@"union" => |u| switch (try self.deserialize(allocator, u.tag_type.?)) {
                inline else => |t| b: {
                    var data: T = @unionInit(T, @tagName(t), undefined);
                    const field = &@field(data, @tagName(t));
                    field.* = try self.deserialize(allocator, @TypeOf(field.*));
                    break :b data;
                },
            },
            .@"enum" => |e| @enumFromInt(try self.deserialize(allocator, e.tag_type)),
            .pointer => |p| switch (p.size) {
                .one => b: {
                    const ptr = try allocator.create(p.child);
                    errdefer allocator.destroy(ptr);
                    ptr.* = try self.deserialize(allocator, p.child);
                    break :b ptr;
                },
                .slice => b: {
                    const slice = try allocator.alloc(p.child, try self.deserialize(allocator, usize));
                    errdefer allocator.free(slice);
                    for (slice) |*e|
                        e.* = try self.deserialize(allocator, p.child);
                    break :b slice;
                },
                .many, .c => @compileError(std.fmt.comptimePrint("Unsupported type: {}", .{T})),
            },
            .int, .bool, .float => self.read(T),
            .void => {},
            .array => |a| b: {
                var array: T = undefined;
                if (@sizeOf(a.child) == 1) {
                    try self.read(&array);
                } else {
                    for (&array) |*e|
                        e.* = try self.deserialize(allocator, a.child);
                }
                break :b array;
            },
            .optional => |o| if (try self.deserialize(allocator, bool))
                try self.deserialize(o.child, allocator)
            else
                null,
            .vector => |v| try self.deserialize(allocator, [v.len]v.child),
            else => @compileError(std.fmt.comptimePrint("Unsupported type: {}", .{T})),
        };
    }

    pub fn read(self: *BitReader, comptime T: type) T {
        // const info = @typeInfo(T);

        const bit_width = @bitSizeOf(T);

        const UnsignedT = std.meta.Int(.unsigned, bit_width);

        var remaining: usize = bit_width;
        var result: UnsignedT = 0;

        while (remaining > 0) {
            const byte_index = self.bit_offset / 8;
            const bit_in_byte = self.bit_offset % 8;

            const space_in_byte = 8 - bit_in_byte;
            const chunk_size = @min(remaining, space_in_byte);

            const mask: u8 = std.math.shr(u8, 0xFF, 8 - chunk_size);

            // Extract chunk MSB-first from current byte
            const shift_from_byte = space_in_byte - chunk_size;

            const bits: u8 = std.math.shr(u8, self.bytes[byte_index], shift_from_byte) & mask;

            // Shift into result (from MSB side)
            result |= std.math.shl(UnsignedT, @intCast(bits), remaining - chunk_size);

            self.bit_offset += chunk_size;
            remaining -= chunk_size;
        }

        return @bitCast(result);
    }
};

test "Bit Stream" {
    std.debug.print("Bit Stream\n", .{});
    const allocator = std.testing.allocator;

    var aw = try std.ArrayList(u8).initCapacity(allocator, 128);
    defer aw.deinit(allocator);

    var writer = BitWriter.init(&aw);

    const _a : u3   = 0x5;
    const _b : bool = true;
    const _c : u3   = 0x1;
    const _d : u4   = 0xE;
    const _e : u32  = 0xFFFFFFFA;
    const _f : i32  = -3;
    const _g : f32  = 6.667;

    const ts = struct {
        a: i32,
        n: []const u8,

        pub fn format(self: @This(), w: *std.io.Writer) !void {

            try w.print(".{{ .a = {}, .n = \"{s}\" }}", .{ self.a, self.n });
        }
    };

    const _h : ts = .{
        .a = 44,
        .n = "hello",
    };

    try writer.serialize(allocator, _a);
    try writer.serialize(allocator, _b);
    try writer.serialize(allocator, _c);
    try writer.serialize(allocator, _d);
    try writer.serialize(allocator, _e);
    try writer.serialize(allocator, _f);
    try writer.serialize(allocator, _g);
    try writer.serialize(allocator, _h);

    var reader = BitReader.init(writer.written());

    const _1 : u3   = try reader.deserialize(allocator, u3);
    const _2 : bool = try reader.deserialize(allocator, bool);
    const _3 : u3   = try reader.deserialize(allocator, u3);
    const _4 : u4   = try reader.deserialize(allocator, u4);
    const _5 : u32  = try reader.deserialize(allocator, u32);
    const _6 : i32  = try reader.deserialize(allocator, i32);
    const _7        = try reader.deserialize(allocator, f32);
    const _8        = try reader.deserialize(allocator, ts);
    defer allocator.free(_8.n);

    try std.testing.expectEqual(_a,   _1);
    try std.testing.expectEqual(_b,   _2);
    try std.testing.expectEqual(_c,   _3);
    try std.testing.expectEqual(_d,   _4);
    try std.testing.expectEqual(_e,   _5);
    try std.testing.expectEqual(_f,   _6);
    try std.testing.expectEqual(_g,   _7);

    try std.testing.expectEqualDeep(_h, _8);
}
