/*
 * HEIF codec.
 * Copyright (c) 2023 Dirk Farin <dirk.farin@gmail.com>
 *
 * This file is part of libheif.
 *
 * libheif is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as
 * published by the Free Software Foundation, either version 3 of
 * the License, or (at your option) any later version.
 *
 * libheif is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with libheif.  If not, see <http://www.gnu.org/licenses/>.
 */

#ifndef LIBHEIF_UNC_DECODER_LEGACYBASE_H
#define LIBHEIF_UNC_DECODER_LEGACYBASE_H

#include <cstdint>
#include <cstring>
#include <algorithm>
#include <map>
#include <iostream>
#include <cassert>
#include <utility>
#include <vector>
#include <memory>

#include "common_utils.h"
#include "context.h"
#include "compression.h"
#include "error.h"
#include "libheif/heif.h"
#include "unc_types.h"
#include "unc_boxes.h"
#include "unc_codec.h"
#include "unc_decoder.h"


class UncompressedBitReader : public BitReader
{
public:
  UncompressedBitReader(const std::vector<uint8_t>& data) : BitReader(data.data(), data.size()) {}

  void markPixelStart()
  {
    m_pixelStartOffset = get_current_byte_index();
  }

  void markRowStart()
  {
    m_rowStartOffset = get_current_byte_index();
  }

  void markTileStart()
  {
    m_tileStartOffset = get_current_byte_index();
  }

  inline Error handlePixelAlignment(uint32_t pixel_size)
  {
    if (pixel_size != 0) {
      uint32_t bytes_in_pixel = static_cast<uint32_t>(get_current_byte_index() - m_pixelStartOffset);
      if (pixel_size > bytes_in_pixel) {
        uint32_t padding = pixel_size - bytes_in_pixel;
        skip_bytes(padding);
      }
      else if (pixel_size == bytes_in_pixel) {
        // NOP
      }
      else {
        return {
          heif_error_Invalid_input,
          heif_suberror_Unspecified,
          "Uncompressed image: invalid 'pixel_size'"
        };
      }
    }

    return {};
  }

  void handleRowAlignment(uint32_t alignment)
  {
    skip_to_byte_boundary();
    if (alignment != 0) {
      uint32_t bytes_in_row = static_cast<uint32_t>(get_current_byte_index() - m_rowStartOffset);
      uint32_t residual = bytes_in_row % alignment;
      if (residual != 0) {
        uint32_t padding = alignment - residual;
        skip_bytes(padding);
      }
    }
  }

  void handleTileAlignment(uint32_t alignment)
  {
    if (alignment != 0) {
      uint32_t bytes_in_tile = static_cast<uint32_t>(get_current_byte_index() - m_tileStartOffset);
      uint32_t residual = bytes_in_tile % alignment;
      if (residual != 0) {
        uint32_t tile_padding = alignment - residual;
        skip_bytes(tile_padding);
      }
    }
  }

private:
  size_t m_pixelStartOffset = 0;
  size_t m_rowStartOffset = 0;
  size_t m_tileStartOffset = 0;
};


template<typename T> void skip_to_alignment(T& position, uint32_t alignment)
{
  if (alignment == 0) {
    return;
  }

  T residual = position % alignment;
  if (residual == 0) {
    return;
  }

  position += alignment - residual;
}


class unc_decoder_legacybase : public unc_decoder
{
public:
  ~unc_decoder_legacybase() override = default;

  void ensure_channel_list(std::shared_ptr<HeifPixelImage>& img) override { ensureChannelList(img); }

protected:
  unc_decoder_legacybase(uint32_t width, uint32_t height,
                         const std::shared_ptr<const Box_cmpd>& cmpd,
                         const std::shared_ptr<const Box_uncC>& uncC,
                         const std::vector<uint32_t>& uncC_index_to_comp_ids);

  class ChannelListEntry
  {
  public:
    uint32_t get_bytes_per_tile() const
    {
      return bytes_per_tile_row_src * tile_height;
    }

    inline uint64_t getDestinationRowOffset(uint32_t tile_row, uint32_t tile_y) const
    {
      uint64_t dst_row_number = uint64_t{tile_row} * tile_height + tile_y;
      return dst_row_number * dst_plane_stride;
    }

    heif_channel channel = heif_channel_Y;
    uint8_t* dst_plane = nullptr;
    size_t dst_plane_stride;
    uint32_t tile_width;
    uint32_t tile_height;
    uint32_t bytes_per_component_sample;
    // Mixed interleave only: the Cb and Cr planes, indexed in the order the
    // two components are declared in the uncC component list (index 0 is
    // whichever comes first there -- not "Cb" specifically), since that is
    // the order their samples are actually interleaved in the bitstream.
    // Each plane is allocated according to its own component's bit depth, so
    // the two can have different byte widths; using one plane's width for
    // the other's write overruns it (GHSA-x8r2-mggj-j6wr).
    uint8_t* chroma_dst_plane[2] = {nullptr, nullptr};
    size_t chroma_dst_plane_stride[2] = {0, 0};
    uint32_t chroma_bytes_per_component_sample[2] = {0, 0};
    uint16_t bits_per_component_sample;
    uint8_t component_alignment;
    uint32_t bytes_per_tile_row_src;
    bool use_channel;
  };

  std::vector<ChannelListEntry> channelList;

  // TODO: refactor so that the channel list is passed explicitly instead of being lazily initialized as member state.
  void ensureChannelList(std::shared_ptr<HeifPixelImage>& img);

  void processComponentSample(UncompressedBitReader& srcBits, const ChannelListEntry& entry, uint64_t dst_row_offset, uint32_t tile_column, uint32_t tile_x);

  // Handles the case where a row consists of a single component type
  // Not valid for Pixel interleave
  // Not valid for the Cb/Cr channels in Mixed Interleave
  // Not valid for multi-Y pixel interleave
  void processComponentRow(ChannelListEntry& entry, UncompressedBitReader& srcBits, uint64_t dst_row_offset, uint32_t tile_column);

  void processComponentTileSample(UncompressedBitReader& srcBits, const ChannelListEntry& entry, uint64_t dst_offset, uint32_t tile_x);

  // Handles the case where a row consists of a single component type
  // Not valid for Pixel interleave
  // Not valid for the Cb/Cr channels in Mixed Interleave
  // Not valid for multi-Y pixel interleave
  void processComponentTileRow(ChannelListEntry& entry, UncompressedBitReader& srcBits, uint64_t dst_offset);

protected:
  void memcpy_to_native_endian(uint8_t* dst, uint32_t value, uint32_t bytes_per_sample);

private:
  void buildChannelList(std::shared_ptr<HeifPixelImage>& img);
  ChannelListEntry buildChannelListEntry(uint32_t component_idx, Box_uncC::Component component, std::shared_ptr<HeifPixelImage>& img);
};

#endif
