//
// SPDX-License-Identifier: BSD-3-Clause
// Copyright (c) Contributors to the OpenEXR Project.
//

//-----------------------------------------------------------------------------
//
//	Utility program to combine or separate multipart image files
//
//-----------------------------------------------------------------------------

#include <ImfChannelList.h>
#include <ImfDeepScanLineInputPart.h>
#include <ImfDeepScanLineOutputPart.h>
#include <ImfDeepTiledInputPart.h>
#include <ImfDeepTiledOutputPart.h>
#include <ImfInputPart.h>
#include <ImfMultiPartInputFile.h>
#include <ImfMultiPartOutputFile.h>
#include <ImfOutputPart.h>
#include <ImfPartHelper.h>
#include <ImfPartType.h>
#include <ImfStringAttribute.h>
#include <ImfTiledInputPart.h>
#include <ImfTiledOutputPart.h>
#include <ImfMisc.h>
#include <OpenEXRConfig.h>

#include <Iex.h>
#include <OpenEXRConfig.h>

#include <algorithm>
#include <assert.h>
#include <cctype>
#include <iostream>
#include <sstream>
#include <stdlib.h>
#include <utility> // pair
#include <vector>

using IMATH_NAMESPACE::Box2i;
using namespace std;

using namespace OPENEXR_IMF_NAMESPACE;

#if defined(ANDROID) || defined(__ANDROID_API__)
#    define IMF_PATH_SEPARATOR "/"
#elif defined(_WIN32) || defined(_WIN64)
#    define IMF_PATH_SEPARATOR "\\"
#else
#    define IMF_PATH_SEPARATOR "/"
#endif

void
copy_tile (
    MultiPartInputFile&  input,
    MultiPartOutputFile& output,
    int                  inPart,
    int                  outPart)
{
    TiledInputPart  in (input, inPart);
    TiledOutputPart out (output, outPart);

    out.copyPixels (in);
}

void
copy_tiledeep (
    MultiPartInputFile&  input,
    MultiPartOutputFile& output,
    int                  inPart,
    int                  outPart)
{
    DeepTiledInputPart  in (input, inPart);
    DeepTiledOutputPart out (output, outPart);

    out.copyPixels (in);
}

void
copy_scanline (
    MultiPartInputFile&  input,
    MultiPartOutputFile& output,
    int                  inPart,
    int                  outPart)
{
    InputPart  in (input, inPart);
    OutputPart out (output, outPart);

    out.copyPixels (in);
}

void
copy_scanlinedeep (
    MultiPartInputFile&  input,
    MultiPartOutputFile& output,
    int                  inPart,
    int                  outPart)
{
    DeepScanLineInputPart  in (input, inPart);
    DeepScanLineOutputPart out (output, outPart);

    out.copyPixels (in);
}

bool
is_number (const std::string& s)
{
    std::string::const_iterator it = s.begin ();
    while (it != s.end () && std::isdigit (*it))
        ++it;
    return !s.empty () && it == s.end ();
}

void
parse_partname (string& part_name)
{
    // strip off a path delimitation
    size_t posSlash = part_name.rfind (IMF_PATH_SEPARATOR);
    if (posSlash != string::npos) part_name = part_name.substr (posSlash + 1);

    // strip the exr ext
    size_t pos = part_name.rfind (".exr");
    if (pos != string::npos)
    {
        part_name = part_name.substr (0, pos);

        // strip off the frame number
        size_t pos2 = part_name.rfind (".");
        if (pos2 != string::npos)
        {
            string frame = part_name.substr (pos2 + 1, pos);
            if (is_number (frame)) part_name = part_name.substr (0, pos2);
        }
    }
}

///
/// If input is <...>::<partname] then extract part name
/// Else, use the stripped version of the filename for partname
/// If input is <file>:<partnum> then extract part number,
/// Else, use all parts
///
void
parse_filename (
    string& file_name, string& part_name, bool& force_part_name, int& part_num)
{
    force_part_name    = false;
    part_name          = file_name; // default is the file_name
    size_t doublecolon = file_name.rfind ("::");
    if (doublecolon != string::npos)
    {
        part_name       = file_name.substr (doublecolon + 2);
        file_name       = file_name.substr (0, doublecolon);
        force_part_name = true;
    }
    else
    {
        parse_partname (part_name);
    }

    size_t colon = file_name.rfind (':');
    if (colon == string::npos)
    {
        part_num = -1; // use all parts
    }
    else
    {
        string num = file_name.substr (colon + 1);
        if (is_number (num))
        {
            part_num  = atoi (num.c_str ());
            file_name = file_name.substr (0, colon);
        }
        else
        {
            throw invalid_argument("part number must be a number");
    }
    }
}

void
make_unique_names (vector<Header>& headers)
{
    set<string> names;
    for (size_t i = 0; i < headers.size (); i++)
    {
        Header&     h = headers[i];
        std::string base_name;

        // if no name at all, set it to <type><partnum> (first part is part 1)
        if (!h.hasName ())
        {
            // We should not be here, we have populated all headers with names.
            throw invalid_argument("Software Error: header does not have a valid name");
        }
        else
        {
            base_name = h.name ();
        }

        // check name has already been used, if so add a _<number> to it
        if (names.find (base_name) != names.end ())
        {
            ostringstream s;
            size_t        backup = 1;
            do
            {
                s.clear ();
                s << h.name () << "_" << i << "_" << backup;
                backup++;
            } while (names.find (s.str ()) != names.end ());
            h.setName (s.str ());
        }

        names.insert (h.name ());
    }
}

void
filename_check (vector<string> names, const char* aname)
{
    string bname (aname);
    for (size_t i = 0; i < names.size (); i++)
        if (bname.compare (names[i]) == 0)
            throw invalid_argument("input and output file names cannot be the same");
}

void
convert (
    vector<const char*> in,
    vector<const char*> views,
    const char*         outname,
    bool                override)
{
    if (in.size () != 1)
        throw invalid_argument("can only convert one file at once - use 'combine' mode for multiple files");

    MultiPartInputFile infile (in[0]);
    
    if (infile.parts () != 1)
        throw invalid_argument("can only convert single part EXRs to multipart EXR-2.0 files: use 'split' mode instead");


    vector<MultiViewChannelName> input_channels;

    string hero;
    if (hasMultiView (infile.header (0)))
    {
        StringVector h = multiView (infile.header (0));
        if (h.size () > 0) { hero = h[0]; }
    }

    // retrieve channel names from input file in view-friendly format
    GetChannelsInMultiPartFile (infile, input_channels);

    vector<MultiViewChannelName> output_channels = input_channels;
    // remap channels to multiple output parts
    int parts = SplitChannels (
                               output_channels.begin (), output_channels.end (), true, hero);

    vector<Header>      output_headers (parts);
    vector<FrameBuffer> output_framebuffers (parts);
    FrameBuffer         input_framebuffer;

    //
    // make all output headers the same as the input header but
    // with no channels
    //
    for (int i = 0; i < parts; i++)
    {
        Header& h = output_headers[i];
        h         = infile.header (0);
        if (hasMultiView (h)) h.erase ("multiView");

        string fn, pname;
        int    pnum;
        bool   pforce;
        parse_filename (fn, pname, pforce, pnum);
        if (!h.hasName () || pforce) h.setName (pname);

        h.channels () = ChannelList ();
    }

    make_unique_names (output_headers);

    const ChannelList& in_chanlist = infile.header (0).channels ();

    int channel_count = 0;
    for (ChannelList::ConstIterator i = in_chanlist.begin ();
         i != in_chanlist.end ();
         ++i)
    {
        ++channel_count;
    }

    Box2i dataWindow = infile.header (0).dataWindow ();
    int   pixel_count =
        (dataWindow.size ().y + 1) * (dataWindow.size ().x + 1);
    int pixel_width = dataWindow.size ().x + 1;

    // offset in pixels between base of array and 0,0
    int pixel_base = dataWindow.min.y * pixel_width + dataWindow.min.x;

    vector<vector<char>> channelstore (channel_count);

    //
    // insert channels into correct header and framebuffers
    //
    for (size_t i = 0; i < input_channels.size (); i++)
    {
        // read the part we should be writing channel into, insert into header
        int                        part = output_channels[i].part_number;
        ChannelList::ConstIterator chan =
            in_chanlist.find (input_channels[i].internal_name);
        Header& h = output_headers[part];
        h.channels ().insert (output_channels[i].name, chan.channel ());

        if (output_channels[i].view != "")
        {
            h.setView (output_channels[i].view);
        }

        // compute size of channel
        size_t samplesize = sizeof (float);
        if (chan.channel ().type == HALF) { samplesize = sizeof (half); }
        channelstore[i].resize (samplesize * pixel_count);

        output_framebuffers[part].insert (output_channels[i].name,
                                          Slice (chan.channel ().type,
                                                 &channelstore[i][0] - pixel_base * samplesize,
                                                 samplesize,
                                                 pixel_width * samplesize));

        input_framebuffer.insert (input_channels[i].internal_name,
                                  Slice (chan.channel ().type,
                                         &channelstore[i][0] - pixel_base * samplesize,
                                         samplesize,
                                         pixel_width * samplesize));
    }

    //
    // create output file
    //
    MultiPartOutputFile outfile (outname, &output_headers[0], output_headers.size ());
                                 
    InputPart inpart (infile, 0);

    //
    // read file
    //
    inpart.setFrameBuffer (input_framebuffer);
    inpart.readPixels (dataWindow.min.y, dataWindow.max.y);

    //
    // write each part
    //

    for (size_t i = 0; i < output_framebuffers.size (); i++)
    {
        OutputPart outpart (outfile, i);
        outpart.setFrameBuffer (output_framebuffers[i]);
        outpart.writePixels (dataWindow.max.y + 1 - dataWindow.min.y);
    }
}

void
combine (
    vector<const char*> in,
    vector<const char*> views,
    const char*         outname,
    bool                override)
{
    size_t                      numInputs = in.size ();
    int                         numparts;
    vector<int>                 partnums;
    vector<MultiPartInputFile*> inputs;
    vector<MultiPartInputFile*> fordelete;
    MultiPartInputFile*         infile;
    vector<Header>              headers;
    vector<string>              fornamecheck;

    //
    // parse all inputs
    //

    // Input format :
    // We support the following syntax for each input
    // <file>[:<partnum>][::<newpartname>]

    for (size_t i = 0; i < numInputs; i++)
    {
        string filename (in[i]);
        string partname;
        bool   forcepartname;
        int    partnum;
        parse_filename (filename, partname, forcepartname, partnum);

        if (partnum == -1)
        {
            fornamecheck.push_back (filename);

            infile = new MultiPartInputFile (filename.c_str ());
            fordelete.push_back (infile);
            numparts = infile->parts ();

            //copy header from all parts of input to our header array
            for (int j = 0; j < numparts; j++)
            {
                inputs.push_back (infile);

                Header h = infile->header (j);
                if (!h.hasName () || forcepartname) h.setName (partname);

                headers.push_back (h);

                if (views[i] != NULL)
                    headers[headers.size () - 1].setView (views[i]);

                partnums.push_back (j);
            }
        } // no user parts specified
        else
        {
            fornamecheck.push_back (filename);

            infile = new MultiPartInputFile (filename.c_str ());
            fordelete.push_back (infile);

            if (partnum >= infile->parts ())
            {
                std::stringstream e;
                e << "You asked for part " << partnum << " in "
                  << in[i]
                  << ", which only has " << infile->parts ()
                  << " parts";
                throw invalid_argument(e.str());
            }
            //copy header from required part of input to our header array
            inputs.push_back (infile);

            Header h = infile->header (partnum);
            if (!h.hasName () || forcepartname) h.setName (partname);

            headers.push_back (h);

            if (views[i] != NULL)
                headers[headers.size () - 1].setView (views[i]);

            partnums.push_back (partnum);
        } // user parts specified
    }

    filename_check (fornamecheck, outname);
    //
    // sort out names - make unique
    //
    if (numInputs > 1) { make_unique_names (headers); }

    //
    // do combine
    //

    // early bail if need be
    MultiPartOutputFile temp (outname, &headers[0], headers.size (), override);

    MultiPartOutputFile out (outname, &headers[0], headers.size (), override);

    for (size_t p = 0; p < partnums.size (); p++)
    {
        std::string type = headers[p].type ();
        if (type == SCANLINEIMAGE)
        {
            cout << "part " << p << ": "
                 << "scanlineimage" << endl;
            copy_scanline (*inputs[p], out, partnums[p], p);
        }
        else if (type == TILEDIMAGE)
        {
            cout << "part " << p << ": "
                 << "tiledimage" << endl;
            copy_tile (*inputs[p], out, partnums[p], p);
        }
        else if (type == DEEPSCANLINE)
        {
            cout << "part " << p << ": "
                 << "deepscanlineimage" << endl;
            copy_scanlinedeep (*inputs[p], out, partnums[p], p);
        }
        else if (type == DEEPTILE)
        {
            cout << "part " << p << ": "
                 << "deeptile" << endl;
            copy_tiledeep (*inputs[p], out, partnums[p], p);
        }
    }

    for (size_t k = 0; k < fordelete.size (); k++)
    {
        delete fordelete[k];
    }

    inputs.clear ();

    cout << "\n"
         << "Combine Success" << endl;
}

void
separate (vector<const char*> in, const char* out, bool override)
{
    if (in.size () > 1)
        throw invalid_argument("-separate only take one input file");

    //
    // parse the multipart input
    //
    string              filename (in[0]);
    MultiPartInputFile* inputimage;
    int                 numOutputs;
    vector<string>      fornamecheck;

    // add check for existence of the file
    MultiPartInputFile temp (filename.c_str ());

    inputimage = new MultiPartInputFile (filename.c_str ());
    numOutputs = inputimage->parts ();
    cout << "numOutputs: " << numOutputs << endl;

    //
    // set outputs names
    //
    for (int p = 0; p < numOutputs; p++)
    {
        string outfilename (out);

        //add number to outfilename
        std::ostringstream oss;
        oss << '.' << p + 1;
        outfilename += oss.str ();
        outfilename += ".exr";
        cout << "outputfilename: " << outfilename << endl;
        fornamecheck.push_back (outfilename);
    }

    filename_check (fornamecheck, in[0]);

    //
    // separate outputs
    //
    for (int p = 0; p < numOutputs; p++)
    {
        Header header = inputimage->header (p);

        MultiPartOutputFile out (
            fornamecheck[p].c_str (), &header, 1, override);

        std::string type = header.type ();
        if (type == "scanlineimage")
        {
            cout << "scanlineimage" << endl;
            copy_scanline (*inputimage, out, p, 0);
        }
        else if (type == "tiledimage")
        {
            cout << "tiledimage" << endl;
            copy_tile (*inputimage, out, p, 0);
        }
        else if (type == "deepscanline")
        {
            cout << "deepscanline" << endl;
            copy_scanlinedeep (*inputimage, out, p, 0);
        }
        else if (type == "deeptile")
        {
            cout << "deeptile" << endl;
            copy_tiledeep (*inputimage, out, p, 0);
        }
    }

    delete inputimage;
    cout << "\n"
         << "Separate Success" << endl;
}

void
usageMessage (ostream& stream, const char* program_name, bool verbose = false)
{
    stream << "Usage: " << program_name << "-combine -i input.exr[:partnum][::partname] [input2.exr[:partnum]][::partname] [...] -o outfile.exr [options]\n";
    stream << "  or: exrmultipart -separate -i infile.exr -o outfileBaseName [options]\n";
    stream << "  or: exrmultipart -convert -i infile.exr -o outfile.exr [options]\n";       

    if (verbose)
        stream << "\n"
            "Combine or split multipart data\n"
            "\n"
            "Options:\n"
            "  -override [0/1]   0 = do not override conflicting shared "
            "                        attributes [default]\n"
            "                    1 = override conflicting shared attributes\n"
            "  -view name        (after specifying -i) assign following inputs to view 'name'\n"
            "  -h, --help        print this message\n"
            "\n"
            "      --version     print version information\n"
            "\n"
            "Report bugs via https://github.com/AcademySoftwareFoundation/openexr/issues or email security@openexr.com\n"
            "";
}

int
main (int argc, char* argv[])
{
    try {
        
        vector<const char*> inFiles;
        vector<const char*> views;
        const char*         view     = 0;
        const char*         outFile  = 0;
        bool                override = false;

        int i = 1;
        int mode = 0; // 0-do not read input, 1-infiles, 2-outfile, 3-override, 4-view

        while (i < argc)
        {
            if (!strcmp (argv[i], "-h") || !strcmp (argv[i], "--help"))     
            {
                usageMessage (cout, "exrmultipart", true);
                return 0;
            }

            if (!strcmp (argv[i], "--version"))
            {
                const char* libraryVersion = getLibraryVersion();
            
                cout << "exrmultipart (OpenEXR) " << OPENEXR_VERSION_STRING;
                if (strcmp(libraryVersion, OPENEXR_VERSION_STRING))
                    cout << "(OpenEXR version " << libraryVersion << ")";
                cout << " https://openexr.com" << endl;
                cout << "Copyright (c) Contributors to the OpenEXR Project" << endl;
                cout << "License BSD-3-Clause" << endl;
                return 0;
            }

            if (!strcmp (argv[i], "-i"))
            {
                mode = 1;
            }
            else if (!strcmp (argv[i], "-o"))
            {
                mode = 2;
            }
            else if (!strcmp (argv[i], "-override"))
            {
                mode = 3;
            }
            else if (!strcmp (argv[i], "-view"))
            {
                if (mode != 1)
                    throw invalid_argument("-view requires -i");

                mode = 4;
            }
            else
            {
                switch (mode)
                {
                  case 1:
                      inFiles.push_back (argv[i]);
                      views.push_back (view);
                      break;
                  case 2: outFile = argv[i]; break;
                  case 3: override = atoi (argv[i]); break;
                  case 4:
                      view = argv[i];
                      mode = 1;
                      break;
                }
            }
            i++;
        }

        // check input and output files found or not
        if (inFiles.size () == 0)
        {
            usageMessage(cerr, argv[0], false);
            return 1;
        }

        cout << "input:" << endl;
        for (size_t i = 0; i < inFiles.size (); i++)
        {
            cout << "      " << inFiles[i];
            if (views[i])
                cout << " in view " << views[i];
            cout << endl;
        }

        if (!outFile)
            throw invalid_argument("No output file specified");

        cout << "output:\n      " << outFile << endl;
        cout << "override:" << override << "\n" << endl;

        if (!strcmp (argv[1], "-combine"))
        {
            cout << "-combine multipart input " << endl;
            combine (inFiles, views, outFile, override);
        }
        else if (!strcmp (argv[1], "-separate"))
        {
            cout << "-separate multipart input " << endl;
            separate (inFiles, outFile, override);
        }
        else if (!strcmp (argv[1], "-convert"))
        {
            cout << "-convert input to EXR2 multipart" << endl;
            convert (inFiles, views, outFile, override);
        }
        else
        {
            usageMessage (cerr, argv[0], false);
        }
    }
    catch (const exception& e)
    {
        cerr << argv[0] << ": " << e.what () << endl;
        return 1;
    }
    
    return 0;
}
