/**********************************************************************************************
    Copyright (C) 2018 Oliver Eichler <oliver.eichler@gmx.de>

    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, either version 3 of the License, or
    (at your option) any later version.

    This program 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 General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.

**********************************************************************************************/

#include "CApp.h"

#include <gdal_alg.h>
#include <gdal_priv.h>
#include <iostream>

const GDALColorEntry CApp::noColor = {255,255,255,0};

void printStdoutQString(const QString& str)
{
    QByteArray array = str.toUtf8();
    printf("%s", array.data());
    printf("\n");
}

void printStderrQString(const QString& str)
{
    QByteArray array = str.toUtf8();
    fprintf(stderr, "%s", array.data());
    fprintf(stderr, "\n");
}



CApp::CApp(qint32 ncolors, const QString& pctFilename, const QString &sctFilename, const QString &srcFilename, const QString &tarFilename)
    : ncolors(ncolors)
    , pctFilename(pctFilename)
    , sctFilename(sctFilename)
    , srcFilename(srcFilename)
    , tarFilename(tarFilename)
{
    GDALAllRegister();
}

qint32 CApp::exec()
{
    qint32 res = 0;
    GDALColorTable * ct = nullptr;
    GDALDataset * dsSrc = nullptr;
    try
    {
        dsSrc = (GDALDataset*)GDALOpenShared(srcFilename.toUtf8(),GA_ReadOnly);
        if(dsSrc == nullptr)
        {
            throw tr("Failed to open source file.");
        }

        if(dsSrc->GetRasterCount() < 3 || dsSrc->GetRasterCount() > 4)
        {
            throw tr("Raster band count of source file must be either 3 or 4.");
        }

        if(QFile(tarFilename).exists())
        {
            QFile::remove(tarFilename);
        }

        ct = createColorTable(ncolors, pctFilename, dsSrc);
        saveColorTable(ct, sctFilename);
        ditherMap(dsSrc, tarFilename, ct);
    }
    catch(const QString& msg)
    {
        printStderrQString(msg);
        res = -1;
    }


    GDALClose(dsSrc);
    delete ct;
    return res;
}

GDALColorTable * CApp::createColorTable(qint32 ncolors, const QString& pctFilename, GDALDataset * dataset)
{
    GDALColorTable * ct = nullptr;
    try
    {
        if(pctFilename.isEmpty())
        {
            ct = (GDALColorTable*)GDALCreateColorTable(GPI_RGB);

            printStdoutQString(tr("Calculate optimal color table from source file"));

            int ok = GDALComputeMedianCutPCT(dataset->GetRasterBand(1),
                                             dataset->GetRasterBand(2),
                                             dataset->GetRasterBand(3),
                                             nullptr,
                                             ncolors,
                                             ct,
                                             GDALTermProgress,
                                             0
                                             );

            if(ok != CE_None)
            {
                throw tr("Failed to create color table.");
            }
        }
        else
        {
            GDALDataset * dsPct = (GDALDataset*)GDALOpenShared(pctFilename.toUtf8(),GA_ReadOnly);
            if(dsPct == nullptr)
            {
                throw tr("Failed to open file with palette.");
            }

            GDALRasterBand * band = (GDALRasterBand*)dsPct->GetRasterBand(1);

            if((dsPct->GetRasterCount() != 1) || (band->GetColorInterpretation() != GCI_PaletteIndex))
            {
                GDALClose(dsPct);
                throw tr("Palette file does not have a single band with a color table");
            }

            int ok = 0;
            band->GetNoDataValue(&ok);

            if(ok || band->GetColorTable()->GetColorEntryCount() > 255)
            {
                GDALClose(dsPct);
                throw tr("The color table must not contain a \"no data\" value and it's size must not exceed 255 colors.");
            }

            ct = dsPct->GetRasterBand(1)->GetColorTable()->Clone();
        }
    }
    catch(const QString& msg)
    {
        delete ct;
        throw msg;
    }
    return ct;
}

void CApp::saveColorTable(GDALColorTable * ct, QString& sctFilename)
{
    if(sctFilename.isEmpty())
    {
        return;
    }

    if(!sctFilename.endsWith(".vrt"))
    {
        sctFilename += ".vrt";
    }

    QByteArray buf = sctFilename.toUtf8();
    printStdoutQString(tr("Save color table to: %1").arg(buf.data()));

    GDALDriverManager * drvman  = GetGDALDriverManager();
    GDALDriver * driver         = drvman->GetDriverByName("VRT");
    GDALDataset * dataset       = driver->Create(sctFilename.toUtf8(), 1, 1, 1, GDT_Byte, {});

    dataset->GetRasterBand(1)->SetColorInterpretation(GCI_PaletteIndex);
    dataset->GetRasterBand(1)->SetColorTable(ct);

    dataset->FlushCache();
    GDALClose(dataset);
}

void CApp::ditherMap(GDALDataset * dsSrc, const QString& tarFilename, GDALColorTable *ct)
{
    if(tarFilename.isEmpty())
    {
        return;
    }

    qint32 xsize = dsSrc->GetRasterBand(1)->GetXSize();
    qint32 ysize = dsSrc->GetRasterBand(1)->GetYSize();

    GDALDriverManager * drvman  = nullptr;
    GDALDriver * driver         = nullptr;
    GDALDataset * dataset       = nullptr;

    try
    {
        const char * cargs[] = {"TILED=YES","COMPRESS=LZW", 0};
        drvman  = GetGDALDriverManager();
        driver  = drvman->GetDriverByName("GTiff");
        dataset = driver->Create(tarFilename.toUtf8(), xsize, ysize, 1, GDT_Byte, (char**)cargs);

        if(dataset == nullptr)
        {
            throw tr("Failed to create target file.");
        }

        dataset->GetRasterBand(1)->SetColorTable(ct);
        dataset->GetRasterBand(1)->SetNoDataValue(ct->GetColorEntryCount());
        dataset->SetProjection(dsSrc->GetProjectionRef());

        double adfGeoTransform[6] = {0};
        dsSrc->GetGeoTransform(adfGeoTransform);
        dataset->SetGeoTransform(adfGeoTransform);

        printStdoutQString(tr("Dither source file to target file"));
        int res = GDALDitherRGB2PCT(dsSrc->GetRasterBand(1),
                                    dsSrc->GetRasterBand(2),
                                    dsSrc->GetRasterBand(3),
                                    dataset->GetRasterBand(1),
                                    ct,
                                    GDALTermProgress,
                                    0
                                    );
        if(res != CE_None)
        {
            throw tr("Failed to dither file.");
        }

        if(dsSrc->GetRasterCount() == 3)
        {
            return;
        }

        GDALRasterBand * alpha = dsSrc->GetRasterBand(4);
        GDALRasterBand * band  = dataset->GetRasterBand(1);

        QByteArray buffer1(xsize, 0);
        QByteArray buffer2(xsize, 0);

        quint8 nodata = band->GetNoDataValue();
        printStdoutQString(tr("Apply alpha channel as no data value to target file"));
        for(int y = 0; y < ysize; y++)
        {
            GDALTermProgress(double(xsize * y)/(xsize*ysize),0,0);
            res = alpha->RasterIO(GF_Read, 0, y, xsize, 1, buffer1.data(), xsize, 1, GDT_Byte, 0, 0);
            if(res != CE_None)
            {
                throw tr("Failed to read from alpha channel.");
            }

            res = band->RasterIO(GF_Read, 0, y, xsize, 1, buffer2.data(), xsize, 1, GDT_Byte, 0, 0);
            if(res != CE_None)
            {
                throw tr("Failed to read from target file.");
            }

            for(int x = 0; x < xsize; x++)
            {
                if(buffer1[x] != char(0xFF))
                {
                    buffer2[x] = nodata;
                }
            }

            res = band->RasterIO(GF_Write, 0, y, xsize, 1, buffer2.data(), xsize, 1, GDT_Byte, 0, 0);
            if(res != CE_None)
            {
                throw tr("Failed to write to target file.");
            }
        }
        GDALTermProgress(1.0,0,0);
    }
    catch(const QString& msg)
    {
        GDALClose(dataset);
        throw msg;
    }

    dataset->FlushCache();
    GDALClose(dataset);
}
