// jpegWriter.C
// Copyright (c) 2009 The Foundry Visionmongers Ltd.  All Rights Reserved.

// Write jpeg JFIF files. This is an example of a file writer that uses
// a separate library. It also reuses the FILE opened by FileWriter but
// none of it's read/write methods. The reason it uses a FileWriter is
// to handle the complexity of Windoze driver letters and the use of a
// temporary output file.

#include <stdio.h>
extern "C" {
#include "jpeg.h"
}
#include "DDImage/FileWriter.h"
#include "DDImage/Row.h"
#include "DDImage/ARRAY.h"
#include "DDImage/Knobs.h"

#include "exifWriter.h"

static struct jpeg_error_mgr jerr;

static const char* kSubSamplingOptions[] = {
  "4:1:1", "4:2:2", "4:4:4", nullptr
};

using namespace DD::Image;

class jpegWriter : public ExifWriter
{
  struct jpeg_compress_struct cinfo;
  float quality;
  int   _subSampling;
public:
  jpegWriter(Write* iop) : ExifWriter(iop)
  {
    cinfo.err = jpeg_std_error(&jerr);
    jpeg_create_compress(&cinfo);
    quality = .75f;
    _subSampling = 0;
  }
  ~jpegWriter() override
  {
    jpeg_destroy_compress(&cinfo);
  }

  void execute() override;
  void knobs(Knob_Callback) override;
  static const Writer::Description d;

  const char* help() override { return "jpeg"; }
};

static Writer* build(Write* iop) { return new jpegWriter(iop); }
const Writer::Description jpegWriter::d("jpeg\0jpg\0", build);

void jpegWriter::execute()
{
  if (!open())
    return;

  bool mono = num_channels() <= 1;
  ChannelSet channels = channel_mask(mono ? 1 : 3);
  input0().request(0, 0, width(), height(), channels, 1);

  jpeg_stdio_dest(&cinfo, (FILE*)file);
  cinfo.image_width = width();
  cinfo.image_height = height();
  if (!mono) {
    cinfo.input_components = 3;
    cinfo.in_color_space = JCS_RGB;
  }
  else {
    cinfo.input_components = 1;
    cinfo.in_color_space = JCS_GRAYSCALE;
  }
  jpeg_set_defaults(&cinfo);
  
  // note: jpeg_set_defaults() sets the other fields ([1] and [2]) to 1  
  switch (_subSampling) {
    case 0:
    default: {
      cinfo.comp_info[0].h_samp_factor = 2;
      cinfo.comp_info[0].v_samp_factor = 2;
      break;
    }
    case 1: {
      cinfo.comp_info[0].h_samp_factor = 2;
      cinfo.comp_info[0].v_samp_factor = 1;
      break;
    }
    case 2: {
      cinfo.comp_info[0].h_samp_factor = 1;
      cinfo.comp_info[0].v_samp_factor = 1;
      break;
    }
  }
  
  cinfo.Y_density = 1200;
  cinfo.X_density = (UINT16)(iop->format().pixel_aspect() * cinfo.Y_density + .5);
  jpeg_set_quality(&cinfo, fast_rint(quality * 100), FALSE);
  jpeg_start_compress(&cinfo, TRUE);

  ExifData* exif = makeExifData();
  unsigned char* data;
  unsigned int data_size;
  exif_data_save_data(exif, &data, &data_size);

  jpeg_write_marker(&cinfo, JPEG_APP0 + 1, data, data_size);

  exif_data_free(exif);

  ARRAY(uchar, buffer, cinfo.input_components * cinfo.image_width);
  unsigned char* arrayp = &buffer[0];

  Row row(0, width());
  for (int y = 0; y < height(); y++) {
    iop->status(double(y) / height());
    get(height() - y - 1, 0, width(), channels, row);
    if (aborted()) {
      jpeg_destroy_compress(&cinfo);
      jpeg_create_compress(&cinfo);
      return;
    }
    for (int i = 0; i < cinfo.input_components; i++)
      to_byte(i, buffer + i, row[channel(i)], nullptr, width(), cinfo.input_components);
    jpeg_write_scanlines(&cinfo, &arrayp, 1);
  }

  jpeg_finish_compress(&cinfo);
  close();
}

void jpegWriter::knobs(Knob_Callback f)
{
  Float_knob(f, &quality, "_jpeg_quality", "quality");
  Enumeration_knob(f, &_subSampling, kSubSamplingOptions, "_jpeg_sub_sampling", "sub-sampling");
  Tooltip(f, "Chroma sub-sampling to use:\n\n"
          "<b>4:1:1</b>\n"
          "<b>4:2:2</b>\n"
          "<b>4:4:4</b>\n"
          "\n<b>Note</b>: 4:4:4 provides the best quality output.");
}