// Copyright 2024 The Wuffs Authors.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.
//
// SPDX-License-Identifier: Apache-2.0 OR MIT

pub status "#bad QUIRK_SOURCE_LENGTH"
pub status "#bad QUIRK_WIDTH_AND_HEIGHT"
pub status "#bad header"
pub status "#inconsistent width and height"
pub status "#missing QUIRK_SOURCE_LENGTH"
pub status "#truncated input"
pub status "#unsupported QUIRK_SOURCE_LENGTH"
pub status "#unsupported VP8 file"

pub const DECODER_WORKBUF_LEN_MAX_INCL_WORST_CASE : base.u64 = 0x30FF_8000

pub struct decoder? implements base.image_decoder(
        width  : base.u32[..= 0x3FFF],
        height : base.u32[..= 0x3FFF],

        // Width and height again, but this time measured in 16×16 macroblocks.
        mbw : base.u32[..= 0x400],
        mbh : base.u32[..= 0x400],

        // The call sequence state machine is discussed in
        // (/doc/std/image-decoders-call-sequence.md).
        call_sequence : base.u8,

        is_key_frame : base.bool,

        quirk_has_alpha_side_channel : base.bool,

        frame_config_io_position : base.u64,

        // This uses the same representation as an wuffs_base__optional_u63:
        // ((value_as_u63 << 1) | has_value_as_u1).
        quirk_source_length : base.u64,

        quirk_width_and_height : base.u64,

        // Just under 144 MiB, for 9 partitions each up to 0xFF_FFFF bytes.
        partitioned_data_length : base.u32[..= 0x08FF_FFF7],

        // args.workbuf strides and offsets.
        //
        // The foobar_end offsets are measured after any "prefix trimmed work
        // buffer" adjustment (if QUIRK_HAS_ALPHA_SIDE_CHANNEL is enabled).
        //
        // The workbuf holds the YUV buffers (after decoding from the source
        // bytes but before converting YUV to RGBX) and then also a copy of
        // most of the source bytes themselves - the "at least 2, up to 9"
        // partitions after the frame header. Decoding successive rows of
        // macroblocks alternates between the other (not-first) partitions.
        // Source bytes are processed with random (not sequential) access.
        workbuf_yuv_y_stride  : base.u32[..= 0x4000],
        workbuf_yuv_uv_stride : base.u32[..= 0x2000],
        workbuf_yuv_y_end     : base.u64[..= 0x1000_0000],
        workbuf_yuv_u_end     : base.u64[..= 0x1400_0000],
        workbuf_yuv_v_end     : base.u64[..= 0x1800_0000],

        // Partition lengths and per-partition entropy decoder state.
        //
        // On naming: "foobar_m1" is "foobar minus 1" and "foobar_m1s" is the
        // plural, "an array of (foobar - 1) values".
        part_workbuf_ris : array[9] base.u64,
        part_range_m1s   : array[9] base.u32[..= 0xFF],
        part_bits        : array[9] base.u32,
        part_n_bits      : array[9] base.u32,
        part_lens        : array[9] base.u32[..= 0xFF_FFFF],

        // The number of other partitions (other than the first one), minus 1.
        num_other_partitions_m1 : base.u32[..= 7],

        // Segmentation parameters.
        seg_enabled    : base.bool,
        seg_update_map : base.bool,
        seg_absolute   : base.bool,
        seg_quants     : array[4] base.u8,
        seg_strengths  : array[4] base.u8,
        seg_probs      : array[3] base.u8,

        // Loop filter parameters.
        filt_simple         : base.bool,
        filt_level          : base.u8[..= 63],
        filt_sharpness      : base.u8[..= 7],
        filt_deltas_enabled : base.bool,
        filt_deltas         : array[2] array[4] base.u8,

        // Dequantization parameters.
        quant_yac_qi     : base.u32[..= 127],
        quant_ydc_delta  : base.u32,
        quant_y2dc_delta : base.u32,
        quant_y2ac_delta : base.u32,
        quant_uvdc_delta : base.u32,
        quant_uvac_delta : base.u32,

        // Look-up tables derived from the loop filter parameters.
        //
        // The array is indexed by ((2 * segment) + inner), where segment is in
        // [0 ..= 3] and inner is in [0 ..= 1]. The inner bit is per any of the
        // three equivalent conditions below (C is libwebp, Go is
        // golang.org/x/image/vp8, Wuffs is this code):
        //
        //  - what C     calls   "is_i4x4    == false" / "is_i4x4    == true".
        //  - what Go    calls   "usePredY16 == true"  / "usePredY16 == false".
        //  - what Wuffs calls   "luma_mode  <  4"     / "luma_mode  >= 4".
        //
        // The u32 bits are packed like this (unless level == 0, in which case
        // the high 24 bits are unused and also all-zero):
        //
        //  - 31         is "is_i4x4", also known as "inner".
        //  - 24 ..= 30 are unused.
        //  - 16 ..= 23 are hlevel (high edge variance level).
        //  -  8 ..= 15 are ilevel (inner level).
        //  -  0 ..=  7 are level.
        loop_filters : array[8] base.u32,

        // Look-up tables derived from the dequantization parameters.
        //
        // The first index is the segment. The second index is Y/Y2/UV. The
        // third index is DC/AC.
        dequants : array[4] array[3] array[2] base.u16,

        // These 9 bits combine what RFC 6386 calls mb_no_skip_coeff (which we
        // hold in the 0x100 bit) and prob_skip_false (in the low 8 bits).
        prob_skip : base.u32[..= 0x1FF],

        // 4×4 luma subblock prediction modes.
        mb_subblock_modes : array[16] base.u8[..= 9],

        // Two u32 values whose bits are packed like this:
        //
        //  - 24 ..= 31 are unused.
        //  - 20 ..= 23 are chroma (V) HZNC bits.
        //  - 16 ..= 19 are chroma (U) HZNC bits.
        //  -  0 ..= 15 are luma   (Y) HNZC bits.
        //
        // HNZC stands for has-non-zero-coefficients.
        mb_dc_nz : base.u32,
        mb_ac_nz : base.u32,

        // The workspace for the reconstruction-in-progress macroblock,
        // containing the predicted and then fully decoded (prediction plus
        // post-IDCT residuals) YUV samples to copy out to the workbuf. Its
        // edges also hold samples from previous (left and top) macroblocks
        // that are used in calculating this macroblock's predictions.
        //
        // See the comments near decoder.reconstruct for the layout.
        yuv_cache : array[26] array[32] base.u8,

        swizzler : base.pixel_swizzler,
        util     : base.utility,
) + (
        swizzle_ycck_scratch_buffer_2k : array[2048] base.u8,

        // 1056 is (4 planes, 8 bands, 3 contexts, 11 tokens) flattened.
        coeff_probs : array[1056] base.u8,

        // Macroblock coefficients. Per this.decode_one_macroblock:
        //
        // this.decode_coefficients decodes these macroblock coefficients from
        // the bitstream and dequantizes them, writing them to this.mb_coeffs.
        // At this point, these are pre-IDCT residuals.
        //
        // this.reconstruct calculates the predictions (writing them to
        // this.yuv_cache, not this.mb_coeffs) and then applies the inverse WHT
        // (if applicable) and the inverse DCT to the pre-IDCT residuals, by
        // calling this.inverse_dct_etc.
        //
        // this.inverse_dct_etc reads the pre-IDCT residuals from
        // this.mb_coeffs and writes fully decoded samples (predictions plus
        // post-IDCT residuals) to this.yuv_cache.
        mb_coeffs : array[25] array[16] base.u16,

        // Macroblock states are a u32 whose bits are packed like this:
        //
        //  - 28 ..= 31 are Chroma (UV)  HNZC bits.
        //  - 24 ..= 27 are Luma4  (Y4)  HNZC bits.
        //  - 23        is  Luma16 (Y16) HNZC bit.
        //  - 16 ..= 22 are unused.
        //  - 12 ..= 15 are the 3rd predictor mode.
        //  -  8 ..= 12 are the 2nd predictor mode.
        //  -  4 ..=  7 are the 1st predictor mode.
        //  -  0 ..=  3 are the 0th predictor mode.
        //
        // HNZC stands for has-non-zero-coefficients.
        //
        // The mb_states_top array is indexed by mbx.
        mb_states_left : base.u32,
        mb_states_top  : array[0x400] base.u32,

        // Macroblock filter parameters.
        //
        // The first index is (mby & 1). As we decode rows (and columns) of
        // macroblocks, the filter parameters are decoded at the time of the
        // R'th row but are applied (to row R) at the time of the (R+1)'th row.
        //
        // The second index is mbx.
        //
        // The u8 bits are packed like this:
        //
        //  -  4 ..=  7 are unused
        //  -  3        is  decode_one_macroblock's (not skip).
        //  -  0 ..=  2 are an index into loop_filters.
        //
        mb_filters : array[2] array[0x400] base.u8,
)

pub func decoder.get_quirk(key: base.u32) base.u64 {
    if args.key == base.QUIRK_SOURCE_LENGTH {
        return this.quirk_source_length
    } else if args.key == QUIRK_HAS_ALPHA_SIDE_CHANNEL {
        if this.quirk_has_alpha_side_channel {
            return 1
        }
    } else if args.key == QUIRK_WIDTH_AND_HEIGHT {
        return this.quirk_width_and_height
    }
    return 0
}

pub func decoder.set_quirk!(key: base.u32, value: base.u64) base.status {
    var w : base.u32
    var h : base.u32

    if args.key == base.QUIRK_SOURCE_LENGTH {
        if (args.value & 1) == 0 {
            if args.value <> 0 {
                return "#bad QUIRK_SOURCE_LENGTH"
            }
            this.quirk_source_length = 0
            this.partitioned_data_length = 0
        } else if (args.value >> 1) < 10 {
            return "#bad QUIRK_SOURCE_LENGTH"
        } else if (args.value >> 1) > (0x08FF_FFF7 + 10) {
            return "#unsupported QUIRK_SOURCE_LENGTH"
        } else {
            this.quirk_source_length = args.value
            this.partitioned_data_length = ((args.value >> 1) as base.u32) - 10
        }
        return ok

    } else if args.key == QUIRK_HAS_ALPHA_SIDE_CHANNEL {
        this.quirk_has_alpha_side_channel = args.value <> 0
        return ok

    } else if args.key == QUIRK_WIDTH_AND_HEIGHT {
        w = (args.value & 0xFFFF_FFFF) as base.u32
        h = (args.value >> 32) as base.u32
        if (w < 1) or (0x3FFF < w) or (h < 1) or (0x3FFF < h) {
            return "#bad QUIRK_WIDTH_AND_HEIGHT"
        }
        this.quirk_width_and_height = args.value
        this.width = w
        this.height = h
        this.calculate_mbw_mbh!()
        return ok
    }

    return base."#unsupported option"
}

pub func decoder.decode_image_config?(dst: nptr base.image_config, src: base.io_reader) {
    var status : base.status

    while true {
        status =? this.do_decode_image_config?(dst: args.dst, src: args.src)
        if (status == base."$short read") and args.src.is_closed() {
            return "#truncated input"
        }
        yield? status
    }
}

pri func decoder.do_decode_image_config?(dst: nptr base.image_config, src: base.io_reader) {
    var c32 : base.u32

    if this.call_sequence <> 0x00 {
        return base."#bad call sequence"
    }

    if this.quirk_source_length == 0 {
        return "#missing QUIRK_SOURCE_LENGTH"
    }

    // RFC 6386 Section 9. Frame Header, up to and including 9.1. Uncompressed
    // Data Chunk. decode_header_partition handles the rest of section 9.

    c32 = args.src.read_u24le_as_u32?()
    this.is_key_frame = ((c32 & 0x01) == 0)
    if not this.is_key_frame {
        // TODO: support non-key frames.
        return "#unsupported VP8 file"
    }
    this.part_lens[0] = c32 >> 5
    if this.part_lens[0] > this.partitioned_data_length {
        return "#bad header"
    }

    c32 = args.src.read_u24le_as_u32?()
    if c32 <> '\x9D\x01\x2A'le {
        return "#bad header"
    }

    c32 = args.src.read_u32le?()
    this.width = 0x3FFF & (c32 >> 0)
    this.height = 0x3FFF & (c32 >> 16)
    if (this.width == 0) or (this.height == 0) {
        return "#bad header"
    } else if (this.quirk_width_and_height <> 0) and
            (this.quirk_width_and_height <> ((this.width as base.u64) | ((this.height as base.u64) << 32))) {
        return "#inconsistent width and height"
    }
    this.calculate_mbw_mbh!()

    this.frame_config_io_position = args.src.position()

    if args.dst <> nullptr {
        args.dst.set!(
                pixfmt: base.PIXEL_FORMAT__BGRX,
                pixsub: 0,
                width: this.width,
                height: this.height,
                first_frame_io_position: this.frame_config_io_position,
                first_frame_is_opaque: false)
    }

    this.call_sequence = 0x20
}

pri func decoder.calculate_mbw_mbh!() {
    this.mbw = (this.width + 15) / 16
    this.mbh = (this.height + 15) / 16

    this.workbuf_yuv_y_stride = this.mbw * 16
    this.workbuf_yuv_uv_stride = this.mbw * 8
    this.workbuf_yuv_y_end =
            ((this.workbuf_yuv_y_stride as base.u64) * ((this.mbh * 16) as base.u64))
    this.workbuf_yuv_u_end = this.workbuf_yuv_y_end +
            ((this.workbuf_yuv_uv_stride as base.u64) * ((this.mbh * 8) as base.u64))
    this.workbuf_yuv_v_end = this.workbuf_yuv_u_end +
            ((this.workbuf_yuv_uv_stride as base.u64) * ((this.mbh * 8) as base.u64))
}

pub func decoder.decode_frame_config?(dst: nptr base.frame_config, src: base.io_reader) {
    var status : base.status

    while true {
        status =? this.do_decode_frame_config?(dst: args.dst, src: args.src)
        if (status == base."$short read") and args.src.is_closed() {
            return "#truncated input"
        }
        yield? status
    }
}

pri func decoder.do_decode_frame_config?(dst: nptr base.frame_config, src: base.io_reader) {
    if this.call_sequence == 0x20 {
        // No-op.
    } else if this.call_sequence < 0x20 {
        this.do_decode_image_config?(dst: nullptr, src: args.src)
    } else if this.call_sequence == 0x28 {
        if this.frame_config_io_position <> args.src.position() {
            return base."#bad restart"
        }
    } else if this.call_sequence == 0x40 {
        this.call_sequence = 0x60
        return base."@end of data"
    } else {
        return base."@end of data"
    }

    if args.dst <> nullptr {
        args.dst.set!(bounds: this.util.make_rect_ie_u32(
                min_incl_x: 0,
                min_incl_y: 0,
                max_excl_x: this.width,
                max_excl_y: this.height),
                duration: 0,
                index: 0,
                io_position: this.frame_config_io_position,
                disposal: 0,
                opaque_within_bounds: false,
                overwrite_instead_of_blend: false,
                background_color: 0xFF00_0000)
    }

    this.call_sequence = 0x40
}

pub func decoder.decode_frame?(dst: ptr base.pixel_buffer, src: base.io_reader, blend: base.pixel_blend, workbuf: slice base.u8, opts: nptr base.decode_frame_options) {
    var wh     : base.u64[..= 0xFFF_8001]
    var wb_len : base.u64[..= DECODER_WORKBUF_LEN_MAX_INCL_WORST_CASE]
    var status : base.status

    // The args.workbuf, with a "prefix trimmed work buffer" adjustment, if
    // QUIRK_HAS_ALPHA_SIDE_CHANNEL is enabled, in which case the first (width
    // × height) bytes of args.workbuf hold alpha samples that are e.g. part of
    // WebP, not part of VP8 per se.
    //
    // The prefix isn't used when decoding VP8 to YCbCr, but the final swizzle
    // to the destination pixel buffer should combine YCbCr and Alpha.
    var trimmed : slice base.u8

    while true {
        trimmed = args.workbuf
        if this.quirk_has_alpha_side_channel {
            if this.quirk_width_and_height == 0 {
                return "#inconsistent width and height"
            }
            wh = (this.width * this.height) as base.u64
            if wh > args.workbuf.length() {
                return base."#bad workbuf length"
            }
            trimmed = args.workbuf[wh ..]
        }

        wb_len = this.workbuf_yuv_v_end + (this.partitioned_data_length as base.u64) + 8
        if wb_len > trimmed.length() {
            return base."#bad workbuf length"
        } else {
            trimmed = trimmed[.. wb_len]
        }

        status =? this.do_decode_frame?(dst: args.dst, src: args.src, blend: args.blend, workbuf: trimmed, opts: args.opts)
        if status.is_ok() {
            break
        } else if (status == base."$short read") and args.src.is_closed() {
            return "#truncated input"
        }
        yield? status
    }

    status = this.swizzle!(dst: args.dst, blend: args.blend, workbuf: args.workbuf)
    if not status.is_ok() {
        return status
    }
}

pri func decoder.do_decode_frame?(dst: ptr base.pixel_buffer, src: base.io_reader, blend: base.pixel_blend, workbuf: slice base.u8, opts: nptr base.decode_frame_options) {
    var i      : base.u32
    var status : base.status

    if this.call_sequence == 0x40 {
        // No-op.
    } else if this.call_sequence < 0x40 {
        this.do_decode_frame_config?(dst: nullptr, src: args.src)
    } else {
        return base."@end of data"
    }

    this.copy_partitions_to_workbuf?(src: args.src, workbuf: args.workbuf)

    // Initialize per-partition bitstreams.
    i = 0
    while i < 9 {
        this.part_workbuf_ris[i] = this.workbuf_yuv_v_end
        this.part_range_m1s[i] = 0xFE
        this.part_bits[i] = 0
        this.part_n_bits[i] = 0
        i += 1
    }

    status = this.decode_header_partition!(workbuf: args.workbuf)
    if not status.is_ok() {
        return status
    }

    status = this.decode_other_partition_lengths!(workbuf: args.workbuf)
    if not status.is_ok() {
        return status
    }

    status = this.swizzler.prepare!(
            dst_pixfmt: args.dst.pixel_format(),
            dst_palette: args.dst.palette(),
            src_pixfmt: this.util.make_pixel_format(repr: base.PIXEL_FORMAT__BGRX),
            src_palette: this.util.empty_slice_u8(),
            blend: args.blend)
    if not status.is_ok() {
        return status
    }

    status = this.decode_macroblocks!(workbuf: args.workbuf)
    if not status.is_ok() {
        return status
    }

    this.call_sequence = 0x60
}

pub func decoder.frame_dirty_rect() base.rect_ie_u32 {
    return this.util.make_rect_ie_u32(
            min_incl_x: 0,
            min_incl_y: 0,
            max_excl_x: this.width,
            max_excl_y: this.height)
}

pub func decoder.num_animation_loops() base.u32 {
    return 0
}

pub func decoder.num_decoded_frame_configs() base.u64 {
    if this.call_sequence > 0x20 {
        return 1
    }
    return 0
}

pub func decoder.num_decoded_frames() base.u64 {
    if this.call_sequence > 0x40 {
        return 1
    }
    return 0
}

pub func decoder.restart_frame!(index: base.u64, io_position: base.u64) base.status {
    if this.call_sequence < 0x20 {
        return base."#bad call sequence"
    }
    if (args.index <> 0) or (args.io_position <> this.frame_config_io_position) {
        return base."#bad argument"
    }
    this.call_sequence = 0x28
    return ok
}

pub func decoder.set_report_metadata!(fourcc: base.u32, report: base.bool) {
    // No-op. VP8 doesn't support metadata (but WebP does).
}

pub func decoder.tell_me_more?(dst: base.io_writer, minfo: nptr base.more_information, src: base.io_reader) {
    return base."#no more information"
}

pub func decoder.workbuf_len() base.range_ii_u64 {
    var wh     : base.u64[..= 0xFFF_8001]
    var wb_len : base.u64[..= DECODER_WORKBUF_LEN_MAX_INCL_WORST_CASE - 0xFFF_8001]

    if this.quirk_source_length == 0 {
        // (min_incl > max_incl) means an empty range_ii_u64. The workbuf
        // length isn't known yet, as we don't know partitioned_data_length.
        return this.util.make_range_ii_u64(min_incl: 1, max_incl: 0)
    }

    if this.quirk_has_alpha_side_channel {
        wh = (this.width * this.height) as base.u64
    }

    wb_len = this.workbuf_yuv_v_end + (this.partitioned_data_length as base.u64) + 8
    return this.util.make_range_ii_u64(
            min_incl: wh + wb_len,
            max_incl: wh + wb_len)
}