/*
 * Copyright (c) 2021, kleines Filmröllchen <filmroellchen@serenityos.org>
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */

#pragma once

#include "FlacTypes.h"
#include "Loader.h"
#include <AK/BitStream.h>
#include <AK/Error.h>
#include <AK/Span.h>
#include <AK/Types.h>

namespace Audio {

ALWAYS_INLINE u8 frame_channel_type_to_channel_count(FlacFrameChannelType channel_type);
// Decodes the sign representation method used in Rice coding.
// Numbers alternate between positive and negative: 0, 1, -1, 2, -2, 3, -3, 4, -4, 5, -5, ...
ALWAYS_INLINE i32 rice_to_signed(u32 x);

// decoders
// read a UTF-8 encoded number, even if it is not a valid codepoint
ALWAYS_INLINE ErrorOr<u64> read_utf8_char(BigEndianInputBitStream& input);
// decode a single number encoded with exponential golomb encoding of the specified order
ALWAYS_INLINE ErrorOr<i32> decode_unsigned_exp_golomb(u8 order, BigEndianInputBitStream& bit_input);

// Loader for the Free Lossless Audio Codec (FLAC)
// This loader supports all audio features of FLAC, although audio from more than two channels is discarded.
// The loader currently supports the STREAMINFO, PADDING, and SEEKTABLE metadata blocks.
// See: https://xiph.org/flac/documentation_format_overview.html
//      https://xiph.org/flac/format.html (identical to IETF draft version 2)
//      https://datatracker.ietf.org/doc/html/draft-ietf-cellar-flac-02 (all section numbers refer to this specification)
//      https://datatracker.ietf.org/doc/html/draft-ietf-cellar-flac-03 (newer IETF draft that uses incompatible numberings and names)
class FlacLoaderPlugin : public LoaderPlugin {
public:
    explicit FlacLoaderPlugin(NonnullOwnPtr<SeekableStream> stream);
    virtual ~FlacLoaderPlugin() override = default;

    static bool sniff(SeekableStream& stream);
    static ErrorOr<NonnullOwnPtr<LoaderPlugin>, LoaderError> create(NonnullOwnPtr<SeekableStream>);

    virtual ErrorOr<Vector<FixedArray<Sample>>, LoaderError> load_chunks(size_t samples_to_read_from_input) override;

    virtual MaybeLoaderError reset() override;
    virtual MaybeLoaderError seek(int sample_index) override;

    virtual int loaded_samples() override { return static_cast<int>(m_loaded_samples); }
    virtual int total_samples() override { return static_cast<int>(m_total_samples); }
    virtual u32 sample_rate() override { return m_sample_rate; }
    virtual u16 num_channels() override { return m_num_channels; }
    virtual ByteString format_name() override { return "FLAC (.flac)"; }
    virtual PcmSampleFormat pcm_format() override { return m_sample_format; }

    bool is_fixed_blocksize_stream() const { return m_min_block_size == m_max_block_size; }
    bool sample_count_unknown() const { return m_total_samples == 0; }

private:
    MaybeLoaderError initialize();
    MaybeLoaderError parse_header();
    // Either returns the metadata block or sets error message.
    // Additionally, increments m_data_start_location past the read meta block.
    ErrorOr<FlacRawMetadataBlock, LoaderError> next_meta_block(BigEndianInputBitStream& bit_input);
    // Fetches and returns the next FLAC frame.
    LoaderSamples next_frame();
    // Helper of next_frame that fetches a sub frame's header
    ErrorOr<FlacSubframeHeader, LoaderError> next_subframe_header(BigEndianInputBitStream& bit_input, u8 channel_index);
    // Helper of next_frame that decompresses a subframe
    ErrorOr<void, LoaderError> parse_subframe(Vector<i64>& samples, FlacSubframeHeader& subframe_header, BigEndianInputBitStream& bit_input);
    // Subframe-internal data decoders (heavy lifting)
    ErrorOr<Vector<i64>, LoaderError> decode_fixed_lpc(FlacSubframeHeader& subframe, BigEndianInputBitStream& bit_input);
    ErrorOr<Vector<i64>, LoaderError> decode_verbatim(FlacSubframeHeader& subframe, BigEndianInputBitStream& bit_input);
    ErrorOr<void, LoaderError> decode_custom_lpc(Vector<i64>& decoded, FlacSubframeHeader& subframe, BigEndianInputBitStream& bit_input);
    MaybeLoaderError decode_residual(Vector<i64>& decoded, FlacSubframeHeader& subframe, BigEndianInputBitStream& bit_input);
    // decode a single rice partition that has its own rice parameter
    ALWAYS_INLINE ErrorOr<Vector<i64>, LoaderError> decode_rice_partition(u8 partition_type, u32 partitions, u32 partition_index, FlacSubframeHeader& subframe, BigEndianInputBitStream& bit_input);
    MaybeLoaderError load_seektable(FlacRawMetadataBlock&);
    // Note that failing to read a Vorbis comment block is not treated as an error of the FLAC loader, since metadata is optional.
    void load_vorbis_comment(FlacRawMetadataBlock&);
    MaybeLoaderError load_picture(FlacRawMetadataBlock&);

    // Converters for special coding used in frame headers
    ALWAYS_INLINE ErrorOr<u32, LoaderError> convert_sample_count_code(u8 sample_count_code);
    ALWAYS_INLINE ErrorOr<u32, LoaderError> convert_sample_rate_code(u8 sample_rate_code);
    ALWAYS_INLINE ErrorOr<u8, LoaderError> convert_bit_depth_code(u8 bit_depth_code);

    bool should_insert_seekpoint_at(u64 sample_index) const;

    // Data obtained directly from the FLAC metadata: many values have specific bit counts
    u32 m_sample_rate { 0 };    // 20 bit
    u8 m_num_channels { 0 };    // 3 bit
    u8 m_bits_per_sample { 0 }; // 5 bits for the integer bit depth
    // Externally visible format; the smallest integer format that's larger than the precise bit depth.
    PcmSampleFormat m_sample_format;
    // Blocks are units of decoded audio data
    u16 m_min_block_size { 0 };
    u16 m_max_block_size { 0 };
    // Frames are units of encoded audio data, both of these are 24-bit
    u32 m_min_frame_size { 0 }; // 24 bit
    u32 m_max_frame_size { 0 }; // 24 bit
    u64 m_total_samples { 0 };  // 36 bit
    u8 m_md5_checksum[128 / 8]; // 128 bit (!)
    size_t m_loaded_samples { 0 };

    // keep track of the start of the data in the FLAC stream to seek back more easily
    u64 m_data_start_location { 0 };
    Optional<FlacFrameHeader> m_current_frame;
    u64 m_current_sample_or_frame { 0 };
    SeekTable m_seektable;

    // Keep around a few temporary buffers whose allocated space can be reused.
    // This is an empirical optimization since allocations and deallocations take a lot of time in the decoder.
    mutable Vector<Vector<i64>, 2> m_subframe_buffers;
};

}