/*
 * Copyright (C) 2020 - 2026 René Rebe, ExactCODE GmbH
 * 
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; version 2. A copy of the GNU General
 * Public License can be found in the file LICENSE.
 * 
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANT-
 * ABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General
 * Public License for more details.
 *
 * Alternatively, commercial licensing options are available from the
 * copyright holder ExactCODE GmbH Germany.
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <iostream>
#include <sstream>
#include <vector>

#include <jxl/decode.h>
#include <jxl/decode_cxx.h>
#include <jxl/resizable_parallel_runner.h>
#include <jxl/resizable_parallel_runner_cxx.h>

#include <jxl/encode.h>
#include <jxl/encode_cxx.h>
#include <jxl/thread_parallel_runner.h>
#include <jxl/thread_parallel_runner_cxx.h>

#include "jpegxl.hh"

/* TODO:
 * streaming API
 * implement < 8bit support
 * implement alpha channel support
 * quick magic check
 * modular e.g. lossless mode and other quality options
 * ICC color profiles
 * thumbnail decoding
 * multi frame images / animations
 * re-compressing jpeg coefficients
 * lossless rotation
 */

int JPEGXLCodec::readImage(std::istream* stream, Image& image, const std::string& decompress)
{
  if (false) {
    // quick magic check
    char buf [6];
    stream->read (buf, sizeof (buf));
    stream->seekg (0);
    // TODO: naked bitstream, or BMFF
    if (buf[4] != 'j' || buf[5] != 'P')
      return false;
  }
  
  // Multi-threaded parallel runner.
  auto runner = JxlResizableParallelRunnerMake(nullptr);

  auto dec = JxlDecoderMake(nullptr);
  if (JXL_DEC_SUCCESS !=
      JxlDecoderSubscribeEvents(dec.get(),
				JXL_DEC_BASIC_INFO | JXL_DEC_FULL_IMAGE)) {
    std::cerr << "JxlDecoderSubscribeEvents failed" << std::endl;
    return false;
  }
  
  if (JXL_DEC_SUCCESS !=
      JxlDecoderSetParallelRunner(dec.get(),
				  JxlResizableParallelRunner,
				  runner.get())) {
    std::cerr << "JxlDecoderSetParallelRunner failed" << std::endl;
    return false;
  }

  JxlBasicInfo info;
  JxlPixelFormat format = {4, JXL_TYPE_UINT8, JXL_NATIVE_ENDIAN, 0};
  
  // read file
  stream->seekg(0, std::ios::end);
  size_t size = stream->tellg();
  stream->seekg(0);
  std::vector<uint8_t> jxl;
  jxl.resize(size);
  stream->read((char*)&jxl[0], size); // TODO: return
  
  JxlDecoderSetInput(dec.get(), &jxl[0], size);
  
  for (;;) {
    JxlDecoderStatus status = JxlDecoderProcessInput(dec.get());
    if (status == JXL_DEC_ERROR) {
      std::cerr << "Decoder error" << std::endl;
      return false;
    } else if (status == JXL_DEC_NEED_MORE_INPUT) {
      std::cerr << "Error, already provided all input" << std::endl;
      return false;
    } else if (status == JXL_DEC_BASIC_INFO) {
      if (JXL_DEC_SUCCESS != JxlDecoderGetBasicInfo(dec.get(), &info)) {
        std::cerr << "JxlDecoderGetBasicInfo failed" << std::endl;
        return false;
      }
      image.bps = info.bits_per_sample;
      image.spp = format.num_channels = info.num_color_channels; // TODO: Alpha
      format.data_type = image.bps > 8 ? JXL_TYPE_UINT16 : JXL_TYPE_UINT8;
      image.resize(info.xsize, info.ysize);
      
      JxlResizableParallelRunnerSetThreads(runner.get(),
					   JxlResizableParallelRunnerSuggestThreads(info.xsize, info.ysize));
    } else if (status == JXL_DEC_NEED_IMAGE_OUT_BUFFER) {
      size_t buffer_size;
      if (JXL_DEC_SUCCESS !=
          JxlDecoderImageOutBufferSize(dec.get(), &format, &buffer_size)) {
        std::cerr << "JxlDecoderImageOutBufferSize failed" << std::endl;
        return false;
      }
      if (buffer_size != image.stride() * image.h) {
        std::cerr << "Invalid out buffer size " << buffer_size << std::endl;
        return false;
      }
      
      if (JXL_DEC_SUCCESS != JxlDecoderSetImageOutBuffer(dec.get(), &format,
							 image.getRawData(),
                                                         buffer_size)) {
        std::cerr << "JxlDecoderSetImageOutBuffer failed" << std::endl;
        return false;
      }
    } else if (status == JXL_DEC_FULL_IMAGE) {
      // Nothing to do. Do not yet return. If the image is an animation, more
      // full frames may be decoded. This example only keeps the last one.
    } else if (status == JXL_DEC_SUCCESS) {
      // All decoding successfully finished.
      // It's not required to call JxlDecoderReleaseInput(dec.get()) here since
      // the decoder will be destroyed.
      return true;
    } else {
      std::cerr << "Unknown decoder status" << std::endl;
      return false;
    }
  }

  return false;
}


bool JPEGXLCodec::writeImage(std::ostream* stream, Image& image, int quality,
			     const std::string& compress)
{
  auto enc = JxlEncoderMake(nullptr/*memory_manager*/);
  auto runner = JxlThreadParallelRunnerMake(nullptr /*memory_manager*/,
					    JxlThreadParallelRunnerDefaultNumWorkerThreads());
  if (JXL_ENC_SUCCESS != JxlEncoderSetParallelRunner(enc.get(),
                                                     JxlThreadParallelRunner,
                                                     runner.get())) {
    std::cerr << "JxlEncoderSetParallelRunner failed" << std::endl;
    return false;
  }
  
  // TODO: stride alignment!
  JxlPixelFormat pixel_format = {
    image.spp, image.bps <= 8 ? JXL_TYPE_UINT8 : JXL_TYPE_UINT16, JXL_NATIVE_ENDIAN, 0
  };

  JxlBasicInfo basic_info = {};
  JxlEncoderInitBasicInfo(&basic_info);
  basic_info.xsize = image.w;
  basic_info.ysize = image.h;
  basic_info.bits_per_sample = image.bps;
  
  if (JxlEncoderSetBasicInfo(enc.get(), &basic_info) != JXL_ENC_SUCCESS) {
    std::cerr << "JxlEncoderSetBasicInfo failed" << std::endl;
    return false;
  }
  
  JxlColorEncoding color_encoding = {};
  JxlColorEncodingSetToSRGB(&color_encoding,
                            pixel_format.num_channels < 3/*is_gray*/);
  if (JxlEncoderSetColorEncoding(enc.get(), &color_encoding) != JXL_ENC_SUCCESS) {
    std::cerr << "JxlEncoderSetColorEncoding failed" << std::endl;
    return false;
  }

  JxlEncoderFrameSettings* frame_settings =
    JxlEncoderFrameSettingsCreate(enc.get(), NULL);

  if (!frame_settings) {
    std::cerr << "JxlEncoderFrameSettingsCreate failed" << std::endl;
    return false;
  }

  /* use BMFF container, not raw codestream */
  JxlEncoderUseContainer(enc.get(), true);

  if (compress == "lossless") {
    JxlEncoderSetFrameLossless(frame_settings, true);
  } else {
    JxlEncoderSetFrameDistance(frame_settings,
			       JxlEncoderDistanceFromQuality(quality));
  }
  
  if (JXL_ENC_SUCCESS !=
      JxlEncoderAddImageFrame(frame_settings,
                              &pixel_format, image.getRawData(),
                              image.stride() * image.h)) {
    std::cerr << "JxlEncoderAddImageFrame failed" << std::endl;
    return false;
  }

  JxlEncoderCloseFrames(enc.get());
  
  std::vector<uint8_t> compressed;
  compressed.resize(4096);
  
  JxlEncoderStatus process_result;
  for (process_result = JXL_ENC_NEED_MORE_OUTPUT; process_result == JXL_ENC_NEED_MORE_OUTPUT;) {
    uint8_t* next_out = compressed.data();
    size_t avail_out = compressed.size();
    process_result = JxlEncoderProcessOutput(enc.get(), &next_out, &avail_out);
    
    size_t write_out = compressed.size() - avail_out;
    if (write_out) {
      stream->write((char*)&compressed[0], write_out);
    }
  }
  
  if (process_result != JXL_ENC_SUCCESS) {
    std::cerr << "JxlEncoderProcessOutput failed" << std::endl;
    return false;
  }
  
  return true;
}

JPEGXLCodec jpegxl_loader;