const std = @import("std");

const Format = @import("Format.zig");

pub const format: Format = .{
    .magic_length = magic.len,
    //  4B - chunk length
    //  4B - chunk type
    // 13B - IHDR data
    //  4B - checksum
    .info_length = magic.len + 25,
    .extension = "png",
    .media_type = "image/png",
    .isFormat = isPng,
};

const magic = "\x89PNG\r\n\x1A\n";

// --- IHDR --------------------------------------------------------------------

pub const Header = struct {
    width: u32,
    height: u32,
    bit_depth: BitDepth,
    color_type: ColorType,
    compression_method: CompressionMethod,
    filter_method: FilterMethod,
    interlace_method: InterlaceMethod,

    pub fn decode(chunk: Chunk) !Header {
        std.debug.assert(chunk.chunk_type == .IHDR);

        if (chunk.data.len != 13) return error.InvalidPng;
        const width = std.mem.readInt(u32, chunk.data[0..4], .big);
        const height = std.mem.readInt(u32, chunk.data[4..8], .big);
        const bit_depth = chunk.data[8];
        const color_type = chunk.data[9];
        const compression_method = chunk.data[10];
        const filter_method = chunk.data[11];
        const interlace_method = chunk.data[12];

        if (width == 0 or width > 0x7FFF_FFFF or
            height == 0 or height > 0x7FFF_FFFF or
            bit_depth == 0 or bit_depth > 16 or !std.math.isPowerOfTwo(bit_depth) or
            color_type == 1 or color_type == 5 or color_type > 6 or
            compression_method != 0 or
            filter_method != 0 or
            interlace_method > 1)
        {
            return error.InvalidPng;
        }

        switch (color_type) {
            0 => {
                // all bit depths allowed
            },
            2 => {
                if (bit_depth < 8) return error.InvalidPng;
            },
            3 => {
                if (bit_depth > 8) return error.InvalidPng;
            },
            4 => {
                if (bit_depth < 8) return error.InvalidPng;
            },
            6 => {
                if (bit_depth < 8) return error.InvalidPng;
            },
            else => unreachable,
        }

        return .{
            .width = width,
            .height = height,
            .bit_depth = @enumFromInt(bit_depth),
            .color_type = @enumFromInt(color_type),
            .compression_method = @enumFromInt(compression_method),
            .filter_method = @enumFromInt(filter_method),
            .interlace_method = @enumFromInt(interlace_method),
        };
    }

    test "refAllDecls" {
        std.testing.refAllDecls(@This());
    }
};

pub const BitDepth = enum(u8) {
    @"1" = 1,
    @"2" = 2,
    @"4" = 4,
    @"8" = 8,
    @"16" = 16,

    pub fn range(self: BitDepth) usize {
        return @as(usize, 2) << @intCast(@intFromEnum(self));
    }

    test "refAllDecls" {
        std.testing.refAllDecls(@This());
    }
};

pub const ColorType = enum(u8) {
    grayscale = 0,
    rgb = 2,
    palette = 3,
    grayscale_alpha = 4,
    rgba = 6,

    pub fn paletteUsed(self: ColorType) bool {
        return @intFromEnum(self) & 0b0000_0001 != 0;
    }

    pub fn colorUsed(self: ColorType) bool {
        return @intFromEnum(self) & 0b0000_0010 != 0;
    }

    pub fn alphaChannelUsed(self: ColorType) bool {
        return @intFromEnum(self) & 0b0000_0100 != 0;
    }

    test "refAllDecls" {
        std.testing.refAllDecls(@This());
    }
};

pub const CompressionMethod = enum(u8) {
    flate = 0,
};

pub const FilterMethod = enum(u8) {
    adaptive = 0,
};

pub const InterlaceMethod = enum(u8) {
    none = 0,
    adam7 = 1,
};

// --- PLTE --------------------------------------------------------------------

pub const Palette = struct {
    len: u16,
    entries: [256][3]u8,

    pub fn asSlice(self: *const Palette) []const [3]u8 {
        return self.entries[0..self.len];
    }

    pub fn decode(chunk: Chunk, ctx: Header) !Palette {
        std.debug.assert(chunk.chunk_type == .PLTE);

        const len = std.math.divExact(usize, chunk.data.len, 3) catch return error.InvalidPng;

        switch (ctx.color_type) {
            .grayscale => return error.InvalidPng,
            .rgb => if (len < 1 or len > 256) return error.InvalidPng,
            .palette => if (len < 1 or len > ctx.bit_depth.range()) return error.InvalidPng,
            .grayscale_alpha => return error.InvalidPng,
            .rgba => if (len < 1 or len > 256) return error.InvalidPng,
        }

        var entries: [256][3]u8 = undefined;
        @memcpy(entries[0..len], @as([]const [3]u8, @ptrCast(chunk.data)));
        return .{
            .len = @intCast(len),
            .entries = entries,
        };
    }

    test "refAllDecls" {
        std.testing.refAllDecls(@This());
    }
};

// --- tRNS --------------------------------------------------------------------

pub const Transparency = union(enum) {
    grayscale: u16,
    rgb: [3]u16,
    palette: TransparencyPalette,

    pub fn initGrayscale(grayscale: u16) Transparency {
        return .{ .grayscale = grayscale };
    }

    pub fn initRgb(rgb: [3]u16) Transparency {
        return .{ .rgb = rgb };
    }

    pub fn initPalette(palette: TransparencyPalette) Transparency {
        return .{ .palette = palette };
    }

    pub fn decode(chunk: Chunk, header: Header, maybe_palette: ?Palette) !Transparency {
        std.debug.assert(chunk.chunk_type == .tRNS);

        return switch (header.color_type) {
            .grayscale => {
                if (chunk.data.len != 2) return error.InvalidPng;
                const grayscale = std.mem.readInt(u16, chunk.data[0..2], .big);
                return .initGrayscale(grayscale);
            },
            .rgb => {
                if (chunk.data.len != 6) return error.InvalidPng;
                const r = std.mem.readInt(u16, chunk.data[0..2], .big);
                const g = std.mem.readInt(u16, chunk.data[2..4], .big);
                const b = std.mem.readInt(u16, chunk.data[4..6], .big);
                return .initRgb(.{ r, g, b });
            },
            .palette => if (maybe_palette) |palette| {
                const len = chunk.data.len;
                if (len < 1 or len > palette.len) return error.InvalidPng;
                var entries: [256]u8 = undefined;
                @memcpy(entries[0..len], chunk.data[0..len]);
                return .initPalette(.{
                    .len = @intCast(len),
                    .entries = entries,
                });
            } else {
                return error.InvalidPng;
            },
            .grayscale_alpha => error.InvalidPng,
            .rgba => error.InvalidPng,
        };
    }

    test "refAllDecls" {
        std.testing.refAllDecls(@This());
    }
};

const TransparencyPalette = struct {
    len: u16,
    entries: [256]u8,

    pub fn asSlice(self: *const TransparencyPalette) []const u8 {
        return self.entries[0..self.len];
    }

    test "refAllDecls" {
        std.testing.refAllDecls(@This());
    }
};

// --- gAMA --------------------------------------------------------------------

const Gamma = struct {
    /// Gamma times `scale`, i.e. 100000, For example, a gamma of 1/2.2 would be
    /// stored as 45455.
    gamma: u32,

    const scale = 100000;

    pub fn initExponent(exponent: f32) Gamma {
        const gamma: u32 = @intFromFloat(@round(exponent * scale));
        return .{ .gamma = gamma };
    }

    pub fn asExponent(self: Gamma) f32 {
        return @as(f32, @floatFromInt(self.gamma)) / scale;
    }

    pub fn decode(chunk: Chunk) Gamma {
        std.debug.assert(chunk.chunk_type == .gAMA);

        if (chunk.data.len != 4) return error.InvalidPng;
        const gamma = std.mem.readInt(u32, chunk.data[0..4], .big);
        return .{ .gamma = gamma };
    }

    test "refAllDecls" {
        std.testing.refAllDecls(@This());
    }
};

// --- cHRM --------------------------------------------------------------------

const Chromaticities = struct {
    white_point_x: u32,
    white_point_y: u32,
    red_x: u32,
    red_y: u32,
    green_x: u32,
    green_y: u32,
    blue_x: u32,
    blue_y: u32,

    pub fn decode(chunk: Chunk) !Chunk {
        std.debug.assert(chunk.chunk_type == .cHRM);

        if (chunk.data.len != 32) return error.InvalidPng;
        return .{
            .white_point_x = std.mem.readInt(u32, chunk.data[0..4], .big),
            .white_point_y = std.mem.readInt(u32, chunk.data[4..8], .big),
            .red_x = std.mem.readInt(u32, chunk.data[8..12], .big),
            .red_y = std.mem.readInt(u32, chunk.data[12..16], .big),
            .green_x = std.mem.readInt(u32, chunk.data[16..20], .big),
            .green_y = std.mem.readInt(u32, chunk.data[20..24], .big),
            .blue_x = std.mem.readInt(u32, chunk.data[24..28], .big),
            .blue_y = std.mem.readInt(u32, chunk.data[28..32], .big),
        };
    }

    test "refAllDecls" {
        std.testing.refAllDecls(@This());
    }
};

// --- sRGB --------------------------------------------------------------------

const RenderingIntent = enum(u8) {
    /// Perceptual intent is for images preferring good adaptation to the output
    /// device gamut at the expense of colorimetric accuracy, like photographs.
    perceptual = 0,
    /// Relative colorimetric intent is for images requiring color appearance
    /// matching (relative to the output device white point), like logos.
    relative_colorimetric = 1,
    /// Saturation intent is for images preferring preservation of saturation at
    /// the expense of hue and lightness, like charts and graphs.
    saturation = 2,
    /// Absolute colorimetric intent is for images requiring preservation of
    /// absolute colorimetry, like proofs (previews of images destined for a
    /// different output device).
    absolute_colorimetric = 3,

    pub fn decode(chunk: Chunk) RenderingIntent {
        std.debug.assert(chunk.chunk_type == .sRGB);

        if (chunk.data.len != 1) return error.InvalidPng;
        const rendering_intent = chunk.data[0];
        if (rendering_intent > 3) return error.InvalidPng;
        return @enumFromInt(rendering_intent);
    }

    test "refAllDecls" {
        std.testing.refAllDecls(@This());
    }
};

// --- pHYs --------------------------------------------------------------------

pub const PhysicalPixelDimensions = struct {
    pixels_per_unit_x: u32,
    pixels_per_unit_y: u32,
    unit: PhysicalUnit,

    pub fn decode(chunk: Chunk) !PhysicalPixelDimensions {
        std.debug.assert(chunk.chunk_type == .pHYs);

        if (chunk.data.len != 9) return error.InvalidPng;
        const pixels_per_unit_x = std.mem.readInt(u32, chunk.data[0..4], .big);
        const pixels_per_unit_y = std.mem.readInt(u32, chunk.data[4..8], .big);
        const unit = chunk.data[8];

        if (pixels_per_unit_x == 0 or pixels_per_unit_y == 0 or unit > 1) {
            return error.InvalidPng;
        }

        return .{
            .pixels_per_unit_x = pixels_per_unit_x,
            .pixels_per_unit_y = pixels_per_unit_y,
            .unit = @enumFromInt(unit),
        };
    }

    test "refAllDecls" {
        std.testing.refAllDecls(@This());
    }
};

pub const PhysicalUnit = enum(u8) {
    unknown = 0,
    meter = 1,
};

// --- tIME --------------------------------------------------------------------

pub const LastModificationTime = struct {
    year: u16,
    /// 1 - 12
    month: u8,
    /// 1 - 31
    day: u8,
    /// 0 - 23
    hour: u8,
    /// 0 - 59
    minute: u8,
    /// 0 - 60 (may include leap seconds)
    second: u8,

    pub fn decode(chunk: Chunk) !LastModificationTime {
        std.debug.assert(chunk.chunk_type == .tIME);

        if (chunk.data.len == 7) return error.InvalidPng;
        const year = std.mem.readInt(u16, chunk.data[0..2], .big);
        const month = chunk.data[2];
        const day = chunk.data[3];
        const hour = chunk.data[4];
        const minute = chunk.data[5];
        const second = chunk.data[6];

        // NOTE This implementation is intentionally trying not to be too clever
        // with validation.
        if (month < 1 or month > 12 or
            day < 1 or day > 31 or
            hour > 23 or
            minute > 59 or
            second > 60)
        {
            return error.InvalidPng;
        }

        return .{
            .year = year,
            .month = month,
            .day = day,
            .hour = hour,
            .minute = minute,
            .second = second,
        };
    }

    test "refAllDecls" {
        std.testing.refAllDecls(@This());
    }
};

// -----------------------------------------------------------------------------

pub const ChunkType = enum(u32) {
    IHDR = fromName("IHDR"),
    PLTE = fromName("PLTE"),
    IDAT = fromName("IDAT"),
    IEND = fromName("IEND"),
    tRNS = fromName("tRNS"),
    gAMA = fromName("gAMA"),
    cHRM = fromName("cHRM"),
    sRGB = fromName("sRGB"),
    iCCP = fromName("iCCP"),
    tEXt = fromName("tEXt"),
    zTXt = fromName("zTXt"),
    iTXt = fromName("iTXt"),
    bKGD = fromName("bKGD"),
    pHYs = fromName("pHYs"),
    sBIT = fromName("sBIT"),
    sPLT = fromName("sPLT"),
    hIST = fromName("hIST"),
    tIME = fromName("tIME"),
    _,

    fn fromName(name: *const [4]u8) u32 {
        return std.mem.bytesToValue(u32, name);
    }

    pub fn fromBytes(bytes: *const [4]u8) ChunkType {
        return @enumFromInt(fromName(bytes));
    }

    pub fn toBytes(self: ChunkType) [4]u8 {
        return std.mem.toBytes(self);
    }

    pub fn isAncillary(self: ChunkType) bool {
        return self.toBytes()[0] & 0b0010_0000 != 0;
    }

    pub fn isPrivate(self: ChunkType) bool {
        return self.toBytes()[1] & 0b0010_0000 != 0;
    }

    pub fn isSafeToCopy(self: ChunkType) bool {
        return self.toBytes()[3] & 0b0010_0000 != 0;
    }

    test "refAllDecls" {
        std.testing.refAllDecls(@This());
    }
};

pub const Chunk = struct {
    chunk_type: ChunkType,
    data: []const u8,
};

pub const StandardKeyword = enum {
    /// Short (one line) title or caption for image
    Title,
    /// Name of image’s creator
    Author,
    /// Description of image (possibly long)
    Description,
    /// Copyright notice
    Copyright,
    /// Time of original image creation
    Creation,
    /// Software used to create the image
    Software,
    /// Legal disclaimer
    Disclaimer,
    /// Warning of nature of content
    Warning,
    /// Device used to create the image
    Source,
    /// Miscellaneous comment; conversion from GIF comment
    Comment,

    pub const map: std.StaticStringMap(StandardKeyword) = blk: {
        const fields = @typeInfo(StandardKeyword).@"enum".fields;

        var kvs_list: [fields.len]struct { []const u8, StandardKeyword } = undefined;
        for (fields, 0..) |field, i| {
            kvs_list[i] = .{ field.name, @field(StandardKeyword, field.name) };
        }

        break :blk .initComptime(kvs_list);
    };

    pub fn isStandardKeyword(keyword: []const u8) ?StandardKeyword {
        return map.get(keyword);
    }

    test "refAllDecls" {
        std.testing.refAllDecls(@This());
    }
};

/// The caller asserts that the buffer is at least `format.magic_length` bytes
/// long.
pub fn isPng(buffer: []const u8) bool {
    return std.mem.eql(u8, buffer[0..format.magic_length], magic);
}

/// The caller asserts that the buffer is at least `format.info_length` bytes
/// long.
pub fn info(buffer: []const u8) ?Header {
    std.debug.assert(buffer.len >= format.info_length);

    if (!isPng(buffer)) {
        return null;
    }

    const chunk, _ = (decodeChunk(buffer[format.magic_length..]) catch return null) orelse return null;
    if (chunk.chunk_type != .IHDR) {
        return null;
    }

    const header = Header.decode(chunk) catch return null;
    return header;
}

pub fn decodeChunk(buffer: []const u8) !?struct { Chunk, []const u8 } {
    var rest: []const u8 = buffer;

    if (rest.len < 8) {
        // Not enough data
        return null;
    }

    const length = std.mem.readInt(u32, rest[0..4], .big);
    const chunk_type: ChunkType = .fromBytes(rest[4..8]);
    rest = rest[8..];

    if (length > 0x7FFF_FFFF) {
        return error.InvalidPng;
    }

    if (rest.len < length) {
        // Not enough data
        return null;
    }

    const data = rest[0..length];
    rest = rest[length..];

    if (rest.len < 4) {
        // Not enough data
        return null;
    }

    const crc = std.mem.readInt(u32, rest[0..4], .big);
    rest = rest[4..];

    if (std.hash.crc.Crc32IsoHdlc.hash(data) != crc) {
        return error.InvalidPng;
    }

    return .{
        .{ .chunk_type = chunk_type, .data = data },
        rest,
    };
}

pub fn decodeChunks(buffer: []const u8, chunks: ?[]Chunk) !usize {
    if (buffer.len < format.magic_length or !isPng(buffer)) {
        return error.InvalidPng;
    }

    var index: usize = 0;
    var rest: []const u8 = buffer[format.magic_length..];

    while (rest.len > 0) : (index += 1) {
        const chunk, rest = try decodeChunk(rest) orelse return error.InvalidPng;

        if (chunks) |ck| {
            if (index < ck.len) {
                ck[index] = chunk;
            }
        }
    }

    return index;
}

pub fn decodeChunksAlloc(buffer: []const u8, allocator: std.mem.Allocator) ![]Chunk {
    const n = try decodeChunks(buffer, null);

    const chunks = try allocator.alloc(Chunk, n);
    errdefer allocator.free(chunks);

    _ = try decodeChunks(buffer, chunks);
    return chunks;
}

pub fn encodeChunks(chunks: []const Chunk, writer: *std.Io.Writer) !void {
    try writer.writeAll(magic);

    for (chunks) |chunk| {
        var crc: std.hash.crc.Crc32IsoHdlc = .init();
        var data: []const u8 = chunk.data;

        while (data.len > 0) {
            const bytes_written = try writer.write(data);
            crc.update(data[0..bytes_written]);
            data = data[bytes_written..];
        }

        try writer.writeInt(u32, crc.final(), .big);
    }
}

test "refAllDecls" {
    std.testing.refAllDecls(@This());
}