const std = @import("std");
const jpg = @import("root.zig");

pub const Segment = union(enum) {
    temporary: void,
    reset: void,
    start_of_image: void,
    end_of_image: void,
    jfif: jpg.Jfif,
    frame: jpg.Frame,
    quantization_tables: [4]?jpg.QuantizationTable,
    huffman_tables: []jpg.HuffmanTable,
    scan: jpg.Scan,

    const Self = @This();

    pub fn read(gpa: std.mem.Allocator, reader: *std.Io.Reader) !Self {
        if (try reader.peekByte() != 0xff) {
            return error.InvalidMarkerHeader;
        }
        reader.toss(1);

        const marker  = jpg.Marker.from_u8(try reader.peekByte())
            catch return error.InvalidMarkerByte;
        reader.toss(1);

        return switch (marker) {
            .temporary => .{ .temporary = void{} },
            .reset => .{ .reset = void{} },
            .start_of_image => .{ .start_of_image = void{} },
            .end_of_image => .{ .end_of_image = void{} },
            .application_segment => |n| read_application_segment(reader, n),
            .define_quantization_table => dqt: {
                var size = try reader.takeInt(u16, .big) - 2;
                var tables: [4]?jpg.QuantizationTable = .{ null, null, null, null };

                while (size > 0) {
                    const table = try jpg.QuantizationTable.read(reader);
                    tables[table.identifier] = table;
                    size -= @truncate(table.size_in_file());
                }

                break :dqt .{ .quantization_tables = tables };
            },
            .start_of_frame => |n| sof: {
                reader.toss(2);
                break :sof .{ .frame = try .read(gpa, n, reader) };
            },
            .define_huffman_table => dhf: {
                var size = try reader.takeInt(u16, .big) - 2;
                var tables: std.ArrayList(jpg.HuffmanTable) = .empty;
                errdefer tables.deinit(gpa);

                while (size > 0) {
                    const table = try tables.addOne(gpa);
                    table.* = try .read(gpa, reader);

                    size -= @truncate(table.size_in_file());
                }

                break :dhf .{ .huffman_tables = try tables.toOwnedSlice(gpa) };
            },
            .start_of_scan => sos: {
                reader.toss(2);
                break :sos .{ .scan = try .read(gpa, reader) };
            },
            else => {
                std.debug.panic("unimplemented tag '{s}'\n", .{@tagName(marker)});
            },
        };
    }

    fn read_application_segment(reader: *std.Io.Reader, n: usize) !Self {
        return switch (n) {
            0 => jfif: {
                reader.toss(2);
                break :jfif .{
                    .jfif = try .read(reader)
                };
            },
            else => std.debug.panic("unimplemented app{}\n", .{n}),
        };
    }

    pub fn deinit(self: *Self, gpa: std.mem.Allocator) void {
        switch (self.*) {
            .frame => |*frame| frame.deinit(gpa),
            .huffman_tables => |tables| {
                for (tables) |*table| table.deinit(gpa);
                gpa.free(tables);
            },
            .scan => |*s| s.deinit(gpa),
            else => {},
        }
    }
};


test "rgb" {
    const buffer = @embedFile("../testing/rgb.jpg");
    var reader: std.Io.Reader = .fixed(buffer);

    {
        var segment: Segment = try .read(std.testing.allocator, &reader);
        defer segment.deinit(std.testing.allocator);
        try std.testing.expectEqual(.start_of_image, std.meta.activeTag(segment));
    }

    {
        var segment: Segment = try .read(std.testing.allocator, &reader);
        defer segment.deinit(std.testing.allocator);
        try std.testing.expectEqual(.jfif, std.meta.activeTag(segment));
    }

    {
        var segment: Segment = try .read(std.testing.allocator, &reader);
        defer segment.deinit(std.testing.allocator);
        const table = segment.quantization_tables[0].?;
        try std.testing.expectEqual(.b8, table.precision);
        try std.testing.expectEqual(0, table.identifier);
    }

    {
        var segment: Segment = try .read(std.testing.allocator, &reader);
        defer segment.deinit(std.testing.allocator);
        const table = segment.quantization_tables[1].?;
        try std.testing.expectEqual(.b8, table.precision);
        try std.testing.expectEqual(1, table.identifier);
    }

    {
        var segment: Segment = try .read(std.testing.allocator, &reader);
        defer segment.deinit(std.testing.allocator);
        try std.testing.expectEqual(0, segment.frame.index);
        try std.testing.expectEqual(8, segment.frame.sample_precision);
        try std.testing.expectEqual(450, segment.frame.number_of_lines);
        try std.testing.expectEqual(600, segment.frame.samples_per_line);
        try std.testing.expectEqual(3, segment.frame.components.len);

        for (segment.frame.components, 0..) |component, index| {
            try std.testing.expectEqual(index + 1, component.identifier);
        }
    }

    for (0..4) |_| {
        var segment: Segment = try .read(std.testing.allocator, &reader);
        defer segment.deinit(std.testing.allocator);
        try std.testing.expectEqual(.huffman_tables, std.meta.activeTag(segment));
    }

    {
        var segment: Segment = try .read(std.testing.allocator, &reader);
        defer segment.deinit(std.testing.allocator);
        try std.testing.expectEqual(.scan, std.meta.activeTag(segment));
    }

    {
        var segment: Segment = try .read(std.testing.allocator, &reader);
        defer segment.deinit(std.testing.allocator);
        try std.testing.expectEqual(.end_of_image, std.meta.activeTag(segment));
    }
}