// Copyright (c) the JPEG XL Project Authors. All rights reserved.

//

// Use of this source code is governed by a BSD-style

// license that can be found in the LICENSE file.

#include "lib/extras/dec/pnm.h"

#include <jxl/codestream_header.h>

#include <jxl/encode.h>

#include <jxl/types.h>

#include <algorithm>

#include <cmath>

#include <cstddef>

#include <cstdint>

#include <cstdlib>

#include <cstring>

#include <utility>

#include <vector>

#include "lib/extras/dec/color_hints.h"

#include "lib/extras/mmap.h"

#include "lib/extras/packed_image.h"

#include "lib/extras/size_constraints.h"

#include "lib/jxl/base/bits.h"

#include "lib/jxl/base/c_callback_support.h"

#include "lib/jxl/base/common.h"

#include "lib/jxl/base/span.h"

#include "lib/jxl/base/status.h"

namespace jxl {

namespace extras {

namespace {

class Parser {

 public:

  explicit Parser(const Span<const uint8_t> bytes)

      : pos_(bytes.data()), end_(pos_ + bytes.size()) {}

  // Sets "pos" to the first non-header byte/pixel on success.

  Status ParseHeader(HeaderPNM* header, const uint8_t** pos) {

    // codec.cc ensures we have at least two bytes => no range check here.

    if (pos_[0] != 'P') return false;

    const uint8_t type = pos_[1];

    pos_ += 2;

    switch (type) {

      case '1':

        return JXL_FAILURE("ascii pbm not supported");

      case '2':

        return JXL_FAILURE("ascii pgm not supported");

      case '3':

        return JXL_FAILURE("ascii ppm not supported");

      case '4':

        return JXL_FAILURE("pbm not supported");

      case '5':

        header->is_gray = true;

        return ParseHeaderPNM(header, pos);

      case '6':

        header->is_gray = false;

        return ParseHeaderPNM(header, pos);

      case '7':

        return ParseHeaderPAM(header, pos);

      case 'F':

        header->is_gray = false;

        return ParseHeaderPFM(header, pos);

      case 'f':

        header->is_gray = true;

        return ParseHeaderPFM(header, pos);

      default:

        return false;

    }

  }

  // Exposed for testing

  Status ParseUnsigned(size_t* number) {

    if (pos_ == end_) return JXL_FAILURE("PNM: reached end before number");

    if (!IsDigit(*pos_)) return JXL_FAILURE("PNM: expected unsigned number");

    *number = 0;

    while (pos_ < end_ && *pos_ >= '0' && *pos_ <= '9') {

      const size_t digit = *pos_ - '0';

      if (!SafeMul(*number, static_cast<size_t>(10), *number) ||

          !SafeAdd(*number, digit, *number)) {

        return JXL_FAILURE("PNM: unsigned number too large");

      }

      ++pos_;

    }

    return true;

  }

  Status ParseSigned(double* number) {

    if (pos_ == end_) return JXL_FAILURE("PNM: reached end before signed");

    if (*pos_ != '-' && *pos_ != '+' && !IsDigit(*pos_)) {

      return JXL_FAILURE("PNM: expected signed number");

    }

    // Skip sign

    const bool is_neg = *pos_ == '-';

    if (is_neg || *pos_ == '+') {

      ++pos_;

      if (pos_ == end_) return JXL_FAILURE("PNM: reached end before digits");

    }

    // Leading digits

    *number = 0.0;

    while (pos_ < end_ && *pos_ >= '0' && *pos_ <= '9') {

      *number *= 10;

      *number += *pos_ - '0';

      ++pos_;

    }

    // Decimal places?

    if (pos_ < end_ && *pos_ == '.') {

      ++pos_;

      double place = 0.1;

      while (pos_ < end_ && *pos_ >= '0' && *pos_ <= '9') {

        *number += (*pos_ - '0') * place;

        place *= 0.1;

        ++pos_;

      }

    }

    if (is_neg) *number = -*number;

    return true;

  }

 private:

  static bool IsDigit(const uint8_t c) { return '0' <= c && c <= '9'; }

  static bool IsLineBreak(const uint8_t c) { return c == '\r' || c == '\n'; }

  static bool IsWhitespace(const uint8_t c) {

    return IsLineBreak(c) || c == '\t' || c == ' ';

  }

  Status SkipBlank() {

    if (pos_ == end_) return JXL_FAILURE("PNM: reached end before blank");

    const uint8_t c = *pos_;

    if (c != ' ' && c != '\n') return JXL_FAILURE("PNM: expected blank");

    ++pos_;

    return true;

  }

  Status SkipSingleWhitespace() {

    if (pos_ == end_) return JXL_FAILURE("PNM: reached end before whitespace");

    if (!IsWhitespace(*pos_)) return JXL_FAILURE("PNM: expected whitespace");

    ++pos_;

    return true;

  }

  Status SkipWhitespace() {

    if (pos_ == end_) return JXL_FAILURE("PNM: reached end before whitespace");

    if (!IsWhitespace(*pos_) && *pos_ != '#') {

      return JXL_FAILURE("PNM: expected whitespace/comment");

    }

    while (pos_ < end_ && IsWhitespace(*pos_)) {

      ++pos_;

    }

    // Comment(s)

    while (pos_ != end_ && *pos_ == '#') {

      while (pos_ != end_ && !IsLineBreak(*pos_)) {

        ++pos_;

      }

      // Newline(s)

      while (pos_ != end_ && IsLineBreak(*pos_)) pos_++;

    }

    while (pos_ < end_ && IsWhitespace(*pos_)) {

      ++pos_;

    }

    return true;

  }

  Status MatchString(const char* keyword, bool skipws = true) {

    const uint8_t* ppos = pos_;

    const uint8_t* kw = reinterpret_cast<const uint8_t*>(keyword);

    while (*kw) {

      if (ppos >= end_) return JXL_FAILURE("PAM: unexpected end of input");

      if (*kw != *ppos) return false;

      ppos++;

      kw++;

    }

    pos_ = ppos;

    if (skipws) {

      JXL_RETURN_IF_ERROR(SkipWhitespace());

    } else {

      JXL_RETURN_IF_ERROR(SkipSingleWhitespace());

    }

    return true;

  }

  Status ParseHeaderPAM(HeaderPNM* header, const uint8_t** pos) {

    size_t depth = 3;

    size_t max_val = 255;

    JXL_RETURN_IF_ERROR(SkipWhitespace());

    while (!MatchString("ENDHDR", /*skipws=*/false)) {

      if (MatchString("WIDTH")) {

        JXL_RETURN_IF_ERROR(ParseUnsigned(&header->xsize));

        JXL_RETURN_IF_ERROR(SkipWhitespace());

      } else if (MatchString("HEIGHT")) {

        JXL_RETURN_IF_ERROR(ParseUnsigned(&header->ysize));

        JXL_RETURN_IF_ERROR(SkipWhitespace());

      } else if (MatchString("DEPTH")) {

        JXL_RETURN_IF_ERROR(ParseUnsigned(&depth));

        JXL_RETURN_IF_ERROR(SkipWhitespace());

      } else if (MatchString("MAXVAL")) {

        JXL_RETURN_IF_ERROR(ParseUnsigned(&max_val));

        JXL_RETURN_IF_ERROR(SkipWhitespace());

      } else if (MatchString("TUPLTYPE")) {

        if (MatchString("RGB_ALPHA")) {

          header->has_alpha = true;

        } else if (MatchString("RGB")) {

        } else if (MatchString("GRAYSCALE_ALPHA")) {

          header->has_alpha = true;

          header->is_gray = true;

        } else if (MatchString("GRAYSCALE")) {

          header->is_gray = true;

        } else if (MatchString("BLACKANDWHITE_ALPHA")) {

          header->has_alpha = true;

          header->is_gray = true;

          max_val = 1;

        } else if (MatchString("BLACKANDWHITE")) {

          header->is_gray = true;

          max_val = 1;

        } else if (MatchString("Alpha")) {

          header->ec_types.push_back(JXL_CHANNEL_ALPHA);

        } else if (MatchString("Depth")) {

          header->ec_types.push_back(JXL_CHANNEL_DEPTH);

        } else if (MatchString("SpotColor")) {

          header->ec_types.push_back(JXL_CHANNEL_SPOT_COLOR);

        } else if (MatchString("SelectionMask")) {

          header->ec_types.push_back(JXL_CHANNEL_SELECTION_MASK);

        } else if (MatchString("Black")) {

          header->ec_types.push_back(JXL_CHANNEL_BLACK);

        } else if (MatchString("CFA")) {

          header->ec_types.push_back(JXL_CHANNEL_CFA);

        } else if (MatchString("Thermal")) {

          header->ec_types.push_back(JXL_CHANNEL_THERMAL);

        } else if (MatchString("Unknown")) {

          header->ec_types.push_back(JXL_CHANNEL_UNKNOWN);

        } else if (MatchString("Optional")) {

          header->ec_types.push_back(JXL_CHANNEL_OPTIONAL);

        } else {

          return JXL_FAILURE("PAM: unknown TUPLTYPE");

        }

      } else {

        constexpr size_t kMaxHeaderLength = 20;

        char unknown_header[kMaxHeaderLength + 1];

        size_t len = std::min<size_t>(kMaxHeaderLength, end_ - pos_);

        strncpy(unknown_header, reinterpret_cast<const char*>(pos_), len);

        unknown_header[len] = 0;

        return JXL_FAILURE("PAM: unknown header keyword: %s", unknown_header);

      }

    }

    size_t num_channels = header->is_gray ? 1 : 3;

    if (header->has_alpha) num_channels++;

    if (num_channels + header->ec_types.size() != depth) {

      return JXL_FAILURE("PAM: bad DEPTH");

    }

    if (max_val == 0 || max_val >= 65536) {

      return JXL_FAILURE("PAM: bad MAXVAL");

    }

    // e.g. When `max_val` is 1 , we want 1 bit:

    header->bits_per_sample = FloorLog2Nonzero(max_val) + 1;

    if ((1u << header->bits_per_sample) - 1 != max_val)

      return JXL_FAILURE("PNM: unsupported MaxVal (expected 2^n - 1)");

    // PAM does not pack bits as in PBM.

    header->floating_point = false;

    header->big_endian = true;

    *pos = pos_;

    return true;

  }

  Status ParseHeaderPNM(HeaderPNM* header, const uint8_t** pos) {

    JXL_RETURN_IF_ERROR(SkipWhitespace());

    JXL_RETURN_IF_ERROR(ParseUnsigned(&header->xsize));

    JXL_RETURN_IF_ERROR(SkipWhitespace());

    JXL_RETURN_IF_ERROR(ParseUnsigned(&header->ysize));

    JXL_RETURN_IF_ERROR(SkipWhitespace());

    size_t max_val;

    JXL_RETURN_IF_ERROR(ParseUnsigned(&max_val));

    if (max_val == 0 || max_val >= 65536) {

      return JXL_FAILURE("PNM: bad MaxVal");

    }

    header->bits_per_sample = FloorLog2Nonzero(max_val) + 1;

    if ((1u << header->bits_per_sample) - 1 != max_val)

      return JXL_FAILURE("PNM: unsupported MaxVal (expected 2^n - 1)");

    header->floating_point = false;

    header->big_endian = true;

    JXL_RETURN_IF_ERROR(SkipSingleWhitespace());

    *pos = pos_;

    return true;

  }

  Status ParseHeaderPFM(HeaderPNM* header, const uint8_t** pos) {

    JXL_RETURN_IF_ERROR(SkipSingleWhitespace());

    JXL_RETURN_IF_ERROR(ParseUnsigned(&header->xsize));

    JXL_RETURN_IF_ERROR(SkipBlank());

    JXL_RETURN_IF_ERROR(ParseUnsigned(&header->ysize));

    JXL_RETURN_IF_ERROR(SkipSingleWhitespace());

    // The scale has no meaning as multiplier, only its sign is used to

    // indicate endianness. All software expects nominal range 0..1.

    double scale;

    JXL_RETURN_IF_ERROR(ParseSigned(&scale));

    if (scale == 0.0) {

      return JXL_FAILURE("PFM: bad scale factor value.");

    } else if (std::abs(scale) != 1.0) {

      JXL_WARNING("PFM: Discarding non-unit scale factor");

    }

    header->big_endian = scale > 0.0;

    header->bits_per_sample = 32;

    header->floating_point = true;

    JXL_RETURN_IF_ERROR(SkipSingleWhitespace());

    *pos = pos_;

    return true;

  }

  const uint8_t* pos_;

  const uint8_t* const end_;

};

}  // namespace

struct PNMChunkedInputFrame {

  JxlChunkedFrameInputSource operator()() {

    return JxlChunkedFrameInputSource{

        this,

        METHOD_TO_C_CALLBACK(

            &PNMChunkedInputFrame::GetColorChannelsPixelFormat),

        METHOD_TO_C_CALLBACK(&PNMChunkedInputFrame::GetColorChannelDataAt),

        METHOD_TO_C_CALLBACK(&PNMChunkedInputFrame::GetExtraChannelPixelFormat),

        METHOD_TO_C_CALLBACK(&PNMChunkedInputFrame::GetExtraChannelDataAt),

        METHOD_TO_C_CALLBACK(&PNMChunkedInputFrame::ReleaseCurrentData)};

  }

  void /* NOLINT */ GetColorChannelsPixelFormat(JxlPixelFormat* pixel_format) {

    *pixel_format = format;

  }

  const void* GetColorChannelDataAt(size_t xpos, size_t ypos, size_t xsize,

                                    size_t ysize, size_t* row_offset) {

    const size_t bytes_per_channel =

        DivCeil(dec->header_.bits_per_sample, jxl::kBitsPerByte);

    const size_t num_channels = dec->header_.is_gray ? 1 : 3;

    const size_t bytes_per_pixel = num_channels * bytes_per_channel;

    *row_offset = dec->header_.xsize * bytes_per_pixel;

    const size_t offset = ypos * *row_offset + xpos * bytes_per_pixel;

    return dec->pnm_.data() + offset + dec->data_start_;

  }

  void GetExtraChannelPixelFormat(size_t ec_index,

                                  JxlPixelFormat* pixel_format) {

    (void)this;

    *pixel_format = {};

    JXL_DEBUG_ABORT("Not implemented");

  }

  const void* GetExtraChannelDataAt(size_t ec_index, size_t xpos, size_t ypos,

                                    size_t xsize, size_t ysize,

                                    size_t* row_offset) {

    (void)this;

    *row_offset = 0;

    JXL_DEBUG_ABORT("Not implemented");

    return nullptr;

  }

  void ReleaseCurrentData(const void* buffer) {}

  JxlPixelFormat format;

  const ChunkedPNMDecoder* dec;

};

StatusOr<ChunkedPNMDecoder> ChunkedPNMDecoder::Init(

    const char* path, const SizeConstraints* constraints) {

  ChunkedPNMDecoder dec;

  JXL_ASSIGN_OR_RETURN(dec.pnm_, MemoryMappedFile::Init(path));

  size_t size = dec.pnm_.size();

  if (size < 2) return JXL_FAILURE("Invalid ppm");

  size_t hdr_buf = std::min<size_t>(size, 10 * 1024);

  Span<const uint8_t> span(dec.pnm_.data(), hdr_buf);

  Parser parser(span);

  HeaderPNM& header = dec.header_;

  const uint8_t* pos = nullptr;

  if (!parser.ParseHeader(&header, &pos)) {

    return StatusCode::kGenericError;

  }

  dec.data_start_ = pos - span.data();

  if (header.bits_per_sample == 0 || header.bits_per_sample > 16) {

    return JXL_FAILURE("Invalid bits_per_sample");

  }

  if (header.has_alpha || !header.ec_types.empty() || header.floating_point) {

    return JXL_FAILURE("Only PGM and PPM inputs are supported");

  }

  JXL_RETURN_IF_ERROR(

      VerifyDimensions(constraints, header.xsize, header.ysize));

  const size_t bytes_per_channel =

      DivCeil(dec.header_.bits_per_sample, jxl::kBitsPerByte);

  const size_t num_channels = dec.header_.is_gray ? 1 : 3;

  const size_t bytes_per_pixel = num_channels * bytes_per_channel;

  size_t row_size;

  if (!SafeMul(dec.header_.xsize, bytes_per_pixel, row_size)) {

    return JXL_FAILURE("PNM image dimensions are too large");

  }

  size_t required_size;

  if (!SafeMul(header.ysize, row_size, required_size) ||

      !SafeAdd(required_size, dec.data_start_, required_size)) {

    return JXL_FAILURE("PNM image dimensions are too large");

  }

  if (size < required_size) {

    return JXL_FAILURE("PNM file too small");

  }

  return dec;

}

jxl::Status ChunkedPNMDecoder::InitializePPF(const ColorHints& color_hints,

                                             PackedPixelFile* ppf) {

  // PPM specifies that in the raster, the sample values are "nonlinear"

  // (BP.709, with gamma number of 2.2). Deviate from the specification and

  // assume `sRGB` in our implementation.

  JXL_RETURN_IF_ERROR(ApplyColorHints(color_hints, /*color_already_set=*/false,

                                      header_.is_gray, ppf));

  ppf->info.xsize = header_.xsize;

  ppf->info.ysize = header_.ysize;

  ppf->info.bits_per_sample = header_.bits_per_sample;

  ppf->info.exponent_bits_per_sample = 0;

  ppf->info.orientation = JXL_ORIENT_IDENTITY;

  ppf->info.alpha_bits = 0;

  ppf->info.alpha_exponent_bits = 0;

  ppf->info.num_color_channels = (header_.is_gray ? 1 : 3);

  ppf->info.num_extra_channels = 0;

  const JxlDataType data_type =

      header_.bits_per_sample > 8 ? JXL_TYPE_UINT16 : JXL_TYPE_UINT8;

  const JxlPixelFormat format{

      /*num_channels=*/ppf->info.num_color_channels,

      /*data_type=*/data_type,

      /*endianness=*/header_.big_endian ? JXL_BIG_ENDIAN : JXL_LITTLE_ENDIAN,

      /*align=*/0,

  };

  PNMChunkedInputFrame frame;

  frame.format = format;

  frame.dec = this;

  ppf->chunked_frames.emplace_back(header_.xsize, header_.ysize, frame);

  return true;

}

Status DecodeImagePNM(const Span<const uint8_t> bytes,

                      const ColorHints& color_hints, PackedPixelFile* ppf,

                      const SizeConstraints* constraints) {

  Parser parser(bytes);

  HeaderPNM header = {};

  const uint8_t* pos = nullptr;

  if (!parser.ParseHeader(&header, &pos)) return false;

  JXL_RETURN_IF_ERROR(

      VerifyDimensions(constraints, header.xsize, header.ysize));

  if (header.bits_per_sample == 0 || header.bits_per_sample > 32) {

    return JXL_FAILURE("PNM: bits_per_sample invalid");

  }

  // PPM specifies that in the raster, the sample values are "nonlinear"

  // (BP.709, with gamma number of 2.2). Deviate from the specification and

  // assume `sRGB` in our implementation.

  JXL_RETURN_IF_ERROR(ApplyColorHints(color_hints, /*color_already_set=*/false,

                                      header.is_gray, ppf));

  ppf->info.xsize = header.xsize;

  ppf->info.ysize = header.ysize;

  if (header.floating_point) {

    ppf->info.bits_per_sample = 32;

    ppf->info.exponent_bits_per_sample = 8;

  } else {

    ppf->info.bits_per_sample = header.bits_per_sample;

    ppf->info.exponent_bits_per_sample = 0;

  }

  ppf->info.orientation = JXL_ORIENT_IDENTITY;

  // No alpha in PNM and PFM

  ppf->info.alpha_bits = (header.has_alpha ? ppf->info.bits_per_sample : 0);

  ppf->info.alpha_exponent_bits = 0;

  ppf->info.num_color_channels = (header.is_gray ? 1 : 3);

  uint32_t num_alpha_channels = (header.has_alpha ? 1 : 0);

  uint32_t num_interleaved_channels =

      ppf->info.num_color_channels + num_alpha_channels;

  ppf->info.num_extra_channels = num_alpha_channels + header.ec_types.size();

  for (auto type : header.ec_types) {

    PackedExtraChannel pec = {};

    pec.ec_info.bits_per_sample = ppf->info.bits_per_sample;

    pec.ec_info.type = type;

    ppf->extra_channels_info.emplace_back(std::move(pec));

  }

  JxlDataType data_type;

  if (header.floating_point) {

    // There's no float16 pnm version.

    data_type = JXL_TYPE_FLOAT;

  } else {

    if (header.bits_per_sample > 8) {

      data_type = JXL_TYPE_UINT16;

    } else {

      data_type = JXL_TYPE_UINT8;

    }

  }

  // No align - pixels are tightly packed.

  constexpr size_t kAlign = 0;

  size_t twidth = PackedImage::BitsPerChannel(data_type) / 8;

  const JxlPixelFormat format{

      /*num_channels=*/num_interleaved_channels,

      /*data_type=*/data_type,

      /*endianness=*/header.big_endian ? JXL_BIG_ENDIAN : JXL_LITTLE_ENDIAN,

      kAlign,

  };

  // EC format is same as color, but 1-channel.

  JxlPixelFormat ec_format = format;

  ec_format.num_channels = 1;

  // Compute required pixel-data size with overflow checks. Without these,

  // a crafted header (large xsize/ysize) wraps the multiplication and lets

  // the size check below pass while the actual memcpy below reads OOB.

  size_t total_channels = num_interleaved_channels + header.ec_types.size();

  size_t required_pnm_size;

  if (!SafeMul(header.xsize, total_channels, required_pnm_size) ||

      !SafeMul(required_pnm_size, twidth, required_pnm_size) ||

      !SafeMul(required_pnm_size, header.ysize, required_pnm_size)) {

    return JXL_FAILURE("PNM image dimensions are too large");

  }

  size_t pnm_remaining_size = bytes.data() + bytes.size() - pos;

  if (pnm_remaining_size < required_pnm_size) {

    return JXL_FAILURE("PNM file too small");

  }

  ppf->frames.clear();

  {

    JXL_ASSIGN_OR_RETURN(

        PackedFrame frame,

        PackedFrame::Create(header.xsize, header.ysize, format));

    ppf->frames.emplace_back(std::move(frame));

  }

  auto* frame = &ppf->frames.back();

  uint8_t* out = reinterpret_cast<uint8_t*>(frame->color.pixels());

  std::vector<uint8_t*> ec_out;

  for (size_t i = 0; i < header.ec_types.size(); ++i) {

    JXL_ASSIGN_OR_RETURN(

        PackedImage ec,

        PackedImage::Create(header.xsize, header.ysize, ec_format));

    frame->extra_channels.emplace_back(std::move(ec));

    ec_out.emplace_back(

        reinterpret_cast<uint8_t*>(frame->extra_channels.back().pixels()));

    JXL_DASSERT(frame->extra_channels.back().stride == header.xsize * twidth);

  }

  JXL_DASSERT(frame->color.stride ==

              header.xsize * num_interleaved_channels * twidth);

  if (ec_out.empty()) {

    const bool flipped_y = (header.bits_per_sample == 32);  // PFMs are flipped

    if (!flipped_y) {

    // When there are no EC and input is not flipped we can copy the whole

    // image at once.

      memcpy(out, pos, header.ysize * frame->color.stride);

    } else {

      // Otherwise copy row-by-row.

      for (size_t y = 0; y < header.ysize; ++y) {

        size_t y_out = header.ysize - 1 - y;

        const uint8_t* row_in = pos + y * frame->color.stride;

        uint8_t* row_out = out + y_out * frame->color.stride;

        memcpy(row_out, row_in, frame->color.stride);

      }

    }

  } else {

    // In case there are EC, we have to deinterleave data pixel-wise.

    JXL_RETURN_IF_ERROR(PackedImage::ValidateDataType(data_type));

    size_t color_stride = twidth * num_interleaved_channels;

    for (size_t y = 0; y < header.ysize; ++y) {

      for (size_t x = 0; x < header.xsize; ++x) {

        memcpy(out, pos, frame->color.pixel_stride());

        out += color_stride;

        pos += color_stride;

        for (auto& p : ec_out) {

          memcpy(p, pos, twidth);

          pos += twidth;

          p += twidth;

        }

      }

    }

  }

  if (ppf->info.exponent_bits_per_sample == 0) {

    ppf->input_bitdepth.type = JXL_BIT_DEPTH_FROM_CODESTREAM;

  }

  return true;

}

// Exposed for testing.

Status PnmParseSigned(Bytes str, double* v) {

  return Parser(str).ParseSigned(v);

}

Status PnmParseUnsigned(Bytes str, size_t* v) {

  return Parser(str).ParseUnsigned(v);

}

}  // namespace extras

}  // namespace jxl