/*
 * Copyright 2021, Haiku, Inc. All rights reserved.
 * Distributed under the terms of the MIT License.
 *
 * Authors:
 *              Emmanuel Gil Peyrot
 */


#include "AVIFTranslator.h"

#include <BufferIO.h>
#include <Catalog.h>
#include <Messenger.h>
#include <TranslatorRoster.h>

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

#include "avif/avif.h"

#include "ConfigView.h"
#include "TranslatorSettings.h"


#undef B_TRANSLATION_CONTEXT
#define B_TRANSLATION_CONTEXT "AVIFTranslator"


class FreeAllocation {
        public:
                FreeAllocation(void* buffer)
                        :
                        fBuffer(buffer)
                {
                }

                ~FreeAllocation()
                {
                        free(fBuffer);
                }

        private:
                void*   fBuffer;
};


// The input formats that this translator knows how to read
static const translation_format sInputFormats[] = {
        {
                AVIF_IMAGE_FORMAT,
                B_TRANSLATOR_BITMAP,
                AVIF_IN_QUALITY,
                AVIF_IN_CAPABILITY,
                "image/avif",
                "AV1 Image File Format"
        },
        {
                B_TRANSLATOR_BITMAP,
                B_TRANSLATOR_BITMAP,
                BITS_IN_QUALITY,
                BITS_IN_CAPABILITY,
                "image/x-be-bitmap",
                "Be Bitmap Format (AVIFTranslator)"
        },
};


// The output formats that this translator knows how to write
static const translation_format sOutputFormats[] = {
        {
                AVIF_IMAGE_FORMAT,
                B_TRANSLATOR_BITMAP,
                AVIF_OUT_QUALITY,
                AVIF_OUT_CAPABILITY,
                "image/avif",
                "AV1 Image File Format"
        },
        {
                B_TRANSLATOR_BITMAP,
                B_TRANSLATOR_BITMAP,
                BITS_OUT_QUALITY,
                BITS_OUT_CAPABILITY,
                "image/x-be-bitmap",
                "Be Bitmap Format (AVIFTranslator)"
        },
};

// Default settings for the Translator
static const TranSetting sDefaultSettings[] = {
        { B_TRANSLATOR_EXT_HEADER_ONLY, TRAN_SETTING_BOOL, false },
        { B_TRANSLATOR_EXT_DATA_ONLY, TRAN_SETTING_BOOL, false },
        { AVIF_SETTING_LOSSLESS, TRAN_SETTING_BOOL, false },
        { AVIF_SETTING_PIXEL_FORMAT, TRAN_SETTING_INT32,
                AVIF_PIXEL_FORMAT_YUV444 },
        { AVIF_SETTING_QUALITY, TRAN_SETTING_INT32, 60 },
        { AVIF_SETTING_SPEED, TRAN_SETTING_INT32, -1 },
        { AVIF_SETTING_TILES_HORIZONTAL, TRAN_SETTING_INT32, 1 },
        { AVIF_SETTING_TILES_VERTICAL, TRAN_SETTING_INT32, 1 },
};

const uint32 kNumInputFormats = sizeof(sInputFormats) /
        sizeof(translation_format);
const uint32 kNumOutputFormats = sizeof(sOutputFormats) /
        sizeof(translation_format);
const uint32 kNumDefaultSettings = sizeof(sDefaultSettings) /
        sizeof(TranSetting);


//      #pragma mark -


AVIFTranslator::AVIFTranslator()
        :
        BaseTranslator(B_TRANSLATE("AVIF images"),
                B_TRANSLATE("AVIF image translator"),
                AVIF_TRANSLATOR_VERSION,
                sInputFormats, kNumInputFormats,
                sOutputFormats, kNumOutputFormats,
                "AVIFTranslator_Settings", sDefaultSettings,
                kNumDefaultSettings, B_TRANSLATOR_BITMAP, AVIF_IMAGE_FORMAT)
{
}


AVIFTranslator::~AVIFTranslator()
{
}


status_t
AVIFTranslator::DerivedIdentify(BPositionIO* stream,
        const translation_format* format, BMessage* settings,
        translator_info* info, uint32 outType)
{
        (void)format;
        (void)settings;
        if (!outType)
                outType = B_TRANSLATOR_BITMAP;
        if (outType != B_TRANSLATOR_BITMAP)
                return B_NO_TRANSLATOR;

        // Read header and first chunck bytes...
        uint32 buf[64];
        ssize_t size = sizeof(buf);
        if (stream->Read(buf, size) != size)
                return B_IO_ERROR;

        // Check it's a valid AVIF format
        avifROData data;
        data.data = reinterpret_cast<const uint8_t*>(buf);
        data.size = static_cast<size_t>(size);
        if (!avifPeekCompatibleFileType(&data))
                return B_ILLEGAL_DATA;

        info->type = AVIF_IMAGE_FORMAT;
        info->group = B_TRANSLATOR_BITMAP;
        info->quality = AVIF_IN_QUALITY;
        info->capability = AVIF_IN_CAPABILITY;
        snprintf(info->name, sizeof(info->name), B_TRANSLATE("AVIF image"));
        strcpy(info->MIME, "image/avif");

        return B_OK;
}


status_t
AVIFTranslator::DerivedTranslate(BPositionIO* stream,
        const translator_info* info, BMessage* ioExtension, uint32 outType,
        BPositionIO* target, int32 baseType)
{
        (void)info;
        if (baseType == 1)
                // if stream is in bits format
                return _TranslateFromBits(stream, ioExtension, outType, target);
        else if (baseType == 0)
                // if stream is NOT in bits format
                return _TranslateFromAVIF(stream, ioExtension, outType, target);
        else
                // if BaseTranslator dit not properly identify the data as
                // bits or not bits
                return B_NO_TRANSLATOR;
}


BView*
AVIFTranslator::NewConfigView(TranslatorSettings* settings)
{
        return new ConfigView(settings);
}


status_t
AVIFTranslator::_TranslateFromBits(BPositionIO* stream, BMessage* ioExtension,
        uint32 outType, BPositionIO* target)
{
        // FIXME: This codepath is completely untested for now, due to libavif
        // being built without any encoder in haikuports.

        (void)ioExtension;
        if (!outType)
                outType = AVIF_IMAGE_FORMAT;
        if (outType != AVIF_IMAGE_FORMAT)
                return B_NO_TRANSLATOR;

        TranslatorBitmap bitsHeader;
        status_t status;

        status = identify_bits_header(stream, NULL, &bitsHeader);
        if (status != B_OK)
                return status;

        avifPixelFormat format = static_cast<avifPixelFormat>(
                fSettings->SetGetInt32(AVIF_SETTING_PIXEL_FORMAT));
        int32 bytesPerPixel;
        avifRGBFormat rgbFormat;
        bool isRGB = true;
        bool ignoreAlpha = false;
        switch (bitsHeader.colors) {
                case B_RGB32:
                        rgbFormat = AVIF_RGB_FORMAT_BGRA;
                        ignoreAlpha = true;
                        bytesPerPixel = 4;
                        break;

                case B_RGB32_BIG:
                        rgbFormat = AVIF_RGB_FORMAT_ARGB;
                        ignoreAlpha = true;
                        bytesPerPixel = 4;
                        break;

                case B_RGBA32:
                        rgbFormat = AVIF_RGB_FORMAT_BGRA;
                        bytesPerPixel = 4;
                        break;

                case B_RGBA32_BIG:
                        rgbFormat = AVIF_RGB_FORMAT_ARGB;
                        bytesPerPixel = 4;
                        break;

                case B_RGB24:
                        rgbFormat = AVIF_RGB_FORMAT_BGR;
                        bytesPerPixel = 3;
                        break;

                case B_RGB24_BIG:
                        rgbFormat = AVIF_RGB_FORMAT_RGB;
                        bytesPerPixel = 3;
                        break;

                case B_YCbCr444:
                        bytesPerPixel = 3;
                        isRGB = false;
                        break;

                case B_GRAY8:
                        bytesPerPixel = 1;
                        isRGB = false;
                        break;

                default:
                        printf("ERROR: Colorspace not supported: %d\n",
                                bitsHeader.colors);
                        return B_NO_TRANSLATOR;
        }

        int width = bitsHeader.bounds.IntegerWidth() + 1;
        int height = bitsHeader.bounds.IntegerHeight() + 1;
        int depth = 8;

        avifImage* image = avifImageCreate(width, height, depth, format);
        image->colorPrimaries = AVIF_COLOR_PRIMARIES_BT709;
        image->transferCharacteristics = AVIF_TRANSFER_CHARACTERISTICS_SRGB;
        image->matrixCoefficients = AVIF_MATRIX_COEFFICIENTS_IDENTITY;

        if (isRGB) {
                image->yuvRange = AVIF_RANGE_FULL;

                avifRGBImage rgb;
                avifRGBImageSetDefaults(&rgb, image);
                rgb.depth = depth;
                rgb.format = rgbFormat;
                rgb.ignoreAlpha = ignoreAlpha;
                int bitsSize = height * bitsHeader.rowBytes;
                rgb.pixels = static_cast<uint8_t*>(malloc(bitsSize));
                if (rgb.pixels == NULL)
                        return B_NO_MEMORY;
                rgb.rowBytes = bitsHeader.rowBytes;

                if (stream->Read(rgb.pixels, bitsSize) != bitsSize) {
                        free(rgb.pixels);
                        return B_IO_ERROR;
                }

                avifResult conversionResult = avifImageRGBToYUV(image, &rgb);
                free(rgb.pixels);
                if (conversionResult != AVIF_RESULT_OK)
                        return B_ERROR;
        } else if (bytesPerPixel == 3) {
                // TODO: Investigate moving that to libavif instead, and do so
                // for other Y'CbCr formats too.
                //
                // See also https://github.com/AOMediaCodec/libavif/pull/235
                assert(bitsHeader.colors == B_YCbCr444);
                int bitsSize = height * bitsHeader.rowBytes;
                uint8_t* pixels = static_cast<uint8_t*>(malloc(bitsSize));
                if (stream->Read(pixels, bitsSize) != bitsSize)
                        return B_IO_ERROR;

                uint8_t* luma = static_cast<uint8_t*>(malloc(bitsSize / 3));
                uint8_t* cb = static_cast<uint8_t*>(malloc(bitsSize / 3));
                uint8_t* cr = static_cast<uint8_t*>(malloc(bitsSize / 3));

                for (int i = 0; i < bitsSize / 3; ++i) {
                        luma[i] = pixels[3 * i + 0];
                        cb[i] = pixels[3 * i + 1];
                        cr[i] = pixels[3 * i + 2];
                }

                image->yuvPlanes[0] = luma;
                image->yuvPlanes[1] = cb;
                image->yuvPlanes[2] = cr;

                image->yuvRowBytes[0] = bitsHeader.rowBytes / 3;
                image->yuvRowBytes[1] = bitsHeader.rowBytes / 3;
                image->yuvRowBytes[2] = bitsHeader.rowBytes / 3;

                image->yuvRange = AVIF_RANGE_LIMITED;
        } else {
                assert(bitsHeader.colors == B_GRAY8);
                int bitsSize = height * bitsHeader.rowBytes;
                uint8_t* luma = static_cast<uint8_t*>(malloc(bitsSize));
                if (stream->Read(luma, bitsSize) != bitsSize)
                        return B_IO_ERROR;

                image->yuvPlanes[0] = luma;
                image->yuvPlanes[1] = nullptr;
                image->yuvPlanes[2] = nullptr;

                image->yuvRowBytes[0] = bitsHeader.rowBytes;
                image->yuvRowBytes[1] = 0;
                image->yuvRowBytes[2] = 0;
        }

        avifRWData output = AVIF_DATA_EMPTY;
        avifEncoder* encoder = avifEncoderCreate();

        system_info info;
        encoder->maxThreads = (get_system_info(&info) == B_OK) ?
                info.cpu_count : 1;

        if (fSettings->SetGetBool(AVIF_SETTING_LOSSLESS)) {
                encoder->minQuantizer = AVIF_QUANTIZER_LOSSLESS;
                encoder->maxQuantizer = AVIF_QUANTIZER_LOSSLESS;
        } else {
                encoder->minQuantizer = encoder->maxQuantizer
                        = fSettings->SetGetInt32(AVIF_SETTING_QUALITY);
        }
        encoder->speed = fSettings->SetGetInt32(AVIF_SETTING_SPEED);
        encoder->tileColsLog2
                = fSettings->SetGetInt32(AVIF_SETTING_TILES_HORIZONTAL);
        encoder->tileRowsLog2
                = fSettings->SetGetInt32(AVIF_SETTING_TILES_VERTICAL);

        avifResult encodeResult = avifEncoderWrite(encoder, image, &output);
        avifImageDestroy(image);
        avifEncoderDestroy(encoder);

        if (encodeResult != AVIF_RESULT_OK) {
                printf("ERROR: Failed to encode: %s\n",
                        avifResultToString(encodeResult));
                avifRWDataFree(&output);
                return B_ERROR;
        }

        // output contains a valid .avif file's contents
        target->Write(output.data, output.size);
        avifRWDataFree(&output);
        return B_OK;
}


status_t
AVIFTranslator::_TranslateFromAVIF(BPositionIO* stream, BMessage* ioExtension,
        uint32 outType, BPositionIO* target)
{
        if (!outType)
                outType = B_TRANSLATOR_BITMAP;
        if (outType != B_TRANSLATOR_BITMAP)
                return B_NO_TRANSLATOR;

        off_t streamLength = 0;
        stream->GetSize(&streamLength);

        off_t streamSize = stream->Seek(0, SEEK_END);
        stream->Seek(0, SEEK_SET);

        void* streamData = malloc(streamSize);
        if (streamData == NULL)
                return B_NO_MEMORY;

        if (stream->Read(streamData, streamSize) != streamSize) {
                free(streamData);
                return B_IO_ERROR;
        }

        avifDecoder* decoder = avifDecoderCreate();
        if (decoder == NULL) {
                free(streamData);
                return B_NO_MEMORY;
        }

        avifResult setIOMemoryResult = avifDecoderSetIOMemory(decoder,
                reinterpret_cast<const uint8_t *>(streamData), streamSize);
        if (setIOMemoryResult != AVIF_RESULT_OK) {
                free(streamData);
                return B_NO_MEMORY;
        }

        avifResult decodeResult = avifDecoderParse(decoder);
        if (decodeResult != AVIF_RESULT_OK) {
                free(streamData);
                return B_ILLEGAL_DATA;
        }

        // We don’t support animations yet.
        if (decoder->imageCount != 1) {
                free(streamData);
                return B_ILLEGAL_DATA;
        }

        avifResult nextImageResult = avifDecoderNextImage(decoder);
        free(streamData);
        if (nextImageResult != AVIF_RESULT_OK)
                return B_ILLEGAL_DATA;

        avifImage* image = decoder->image;
        int width = image->width;
        int height = image->height;
        avifRGBFormat format;
        uint8_t* pixels;
        uint32_t rowBytes;
        color_space colors;

        bool convertToRGB = true;
        if (image->alphaPlane) {
                format = AVIF_RGB_FORMAT_BGRA;
                colors = B_RGBA32;
        } else if (image->yuvFormat == AVIF_PIXEL_FORMAT_YUV400) {
                colors = B_GRAY8;
                convertToRGB = false;
        } else {
                format = AVIF_RGB_FORMAT_BGR;
                colors = B_RGB24;
        }

        if (convertToRGB) {
                avifRGBImage rgb;
                avifRGBImageSetDefaults(&rgb, image);
                rgb.depth = 8;
                rgb.format = format;

                avifRGBImageAllocatePixels(&rgb);
                avifResult conversionResult = avifImageYUVToRGB(image, &rgb);
                if (conversionResult != AVIF_RESULT_OK)
                        return B_ILLEGAL_DATA;

                pixels = rgb.pixels;
                rowBytes = rgb.rowBytes;
        } else {
                // TODO: Add a downsampling (with dithering?) path here, or
                // alternatively add support for higher bit depth to Haiku
                // bitmaps, possibly with HDR too.
                if (image->depth > 8)
                        return B_ILLEGAL_DATA;

                // TODO: Add support for more than just the luma plane.
                pixels = image->yuvPlanes[0];
                rowBytes = image->yuvRowBytes[0];
        }

        uint32 dataSize = rowBytes * height;

        TranslatorBitmap bitmapHeader;
        bitmapHeader.magic = B_TRANSLATOR_BITMAP;
        bitmapHeader.bounds.Set(0, 0, width - 1, height - 1);
        bitmapHeader.rowBytes = rowBytes;
        bitmapHeader.colors = colors;
        bitmapHeader.dataSize = dataSize;

        // write out Be's Bitmap header
        swap_data(B_UINT32_TYPE, &bitmapHeader, sizeof(TranslatorBitmap),
                B_SWAP_HOST_TO_BENDIAN);
        ssize_t bytesWritten = target->Write(&bitmapHeader,
                sizeof(TranslatorBitmap));
        if (bytesWritten < B_OK)
                return bytesWritten;

        if ((size_t)bytesWritten != sizeof(TranslatorBitmap))
                return B_IO_ERROR;

        bool headerOnly = false;
        if (ioExtension != NULL)
                ioExtension->FindBool(B_TRANSLATOR_EXT_HEADER_ONLY,
                        &headerOnly);

        if (headerOnly)
                return B_OK;

        bytesWritten = target->Write(pixels, dataSize);
        if (bytesWritten < B_OK)
                return bytesWritten;
        return B_OK;
}


//      #pragma mark -


BTranslator*
make_nth_translator(int32 n, image_id you, uint32 flags, ...)
{
        (void)you;
        (void)flags;
        if (n != 0)
                return NULL;

        return new AVIFTranslator();
}