/*
 * Copyright (C) Jan 2013 Mellanox Technologies Ltd. All rights reserved.
 * Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 *
 * This software is available to you under a choice of one of two
 * licenses.  You may choose to be licensed under the terms of the GNU
 * General Public License (GPL) Version 2, available from the file
 * COPYING in the main directory of this source tree, or the
 * OpenIB.org BSD license below:
 *
 *     Redistribution and use in source and binary forms, with or
 *     without modification, are permitted provided that the following
 *     conditions are met:
 *
 *      - Redistributions of source code must retain the above
 *        copyright notice, this list of conditions and the following
 *        disclaimer.
 *
 *      - Redistributions in binary form must reproduce the above
 *        copyright notice, this list of conditions and the following
 *        disclaimer in the documentation and/or other materials
 *        provided with the distribution.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
 * BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
 * ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */

/*
 * mlxarchive_mfa2.cpp
 *
 *  Created on: March 23, 2017
 *      Author: Ahmad Soboh
 */

#include <xz_utils/xz_utils.h>
#include <mlxsign_lib/mlxsign_lib.h>
#include "mlxarchive_mfa2.h"
#include "mlxarchive_mfa2_utils.h"

using namespace mfa2;

/*void MFA2::updateSHA256()
{
    vector<u_int8_t> digest;
    MlxSignSHA256 mlxSignSHA256;

    //compute descriptors SHA256
    vector<u_int8_t> descriptorsBuff;
    packDescriptors(descriptorsBuff);
    mlxSignSHA256 << descriptorsBuff;
    mlxSignSHA256.getDigest(digest);
    _packageDescriptor.setDescriptorsSHA256(digest);

    //append the zipped components block
    packBytesArray(zippedComponentBlockBuff.data(), zippedComponentBlockBuff.size(), buff);

    //compute SHA256
    mlxSignSHA256 << zippedComponentBlockBuff;
    mlxSignSHA256.getDigest(digest);
}*/

void MFA2::packDescriptors(vector<u_int8_t>& buff) const
{
    _fingerPrint.pack(buff);
    _packageDescriptor.pack(buff);
    CONST_VECTOR_ITERATOR(DeviceDescriptor, _deviceDescriptors, it) { (*it).pack(buff); }
    CONST_VECTOR_ITERATOR(Component, _components, it) { (*it).packDescriptor(buff); }
}

void MFA2::pack(vector<u_int8_t>& buff)
{
    // find size of components block:
    u_int64_t componentsBlockSize = 0;
    CONST_VECTOR_ITERATOR(Component, _components, it) { componentsBlockSize += (*it).getComponentBinarySize(); }
    _packageDescriptor.setComponentsBlockSize(componentsBlockSize);

    _fingerPrint.pack(buff);

    _packageDescriptor.pack(buff);

    CONST_VECTOR_ITERATOR(DeviceDescriptor, _deviceDescriptors, it) { (*it).pack(buff); }

    CONST_VECTOR_ITERATOR(Component, _components, it) { (*it).packDescriptor(buff); }

    _packageDescriptor.setComponentsBlockOffset(buff.size());

    // compress components block
    vector<u_int8_t> componentsBlockBuff;
    VECTOR_ITERATOR(Component, _components, it)
    {
        (*it).setComponentBinaryOffset(componentsBlockBuff.size());
        (*it).packData(componentsBlockBuff);
    }
    u_int32_t zippedSize = componentsBlockBuff.size();
    zippedSize = xz_compress_crc32(8, componentsBlockBuff.data(), componentsBlockBuff.size(), NULL, 0);
    if (zippedSize <= 0)
    {
        // TODO throw exception
        printf("-E- Error while compressing\n");
        exit(1);
    }
    _packageDescriptor.setComponentsBlockArchiveSize(zippedSize);
    vector<u_int8_t> zippedComponentBlockBuff(zippedSize);
    xz_compress_crc32(9, componentsBlockBuff.data(), componentsBlockBuff.size(), zippedComponentBlockBuff.data(),
                      zippedSize);

    // compute descriptors SHA256
    vector<u_int8_t> descriptorsBuff;
    packDescriptors(descriptorsBuff);
    vector<u_int8_t> digest;
    MlxSignSHA256 mlxSignSHA256;
    mlxSignSHA256 << descriptorsBuff;
    mlxSignSHA256.getDigest(digest);
    _packageDescriptor.setDescriptorsSHA256(digest);

    // append the zipped components block
    packBytesArray(zippedComponentBlockBuff.data(), zippedComponentBlockBuff.size(), buff);

    // compute SHA256
    mlxSignSHA256 << zippedComponentBlockBuff;
    mlxSignSHA256.getDigest(digest);
    _packageDescriptor.setSHA256(digest);
}

/*void MFA2::update(vector<u_int8_t>& buff)
{
    //update components descriptors with information about components blocks
    u_int64_t offset = 0x0;

    VECTOR_ITERATOR(_components, );
}*/

void MFA2::generateBinary(vector<u_int8_t>& buff)
{
    vector<u_int8_t> tmpBuff;

    // do first pass
    pack(tmpBuff);

    // do second pass
    pack(buff);

    // compute SHA256
}

MFA2* MFA2::LoadMFA2Package(const string& file_name)
{
    Mfa2Buffer mfa2buff;
    if (!mfa2buff.loadFile(file_name))
    {
        fprintf(stderr, "Failed to load mfa2 package: %s\n", file_name.c_str());
        return NULL;
    }
    FingerPrint finger_print("");
    VersionExtension version("");
    PackageDescriptor packageDescriptor(0, 0, version);
    vector<DeviceDescriptor> deviceDescriptors;
    vector<Component> components;
    MFA2* mfa2pkg = new MFA2(packageDescriptor, deviceDescriptors, components);
    if (!mfa2pkg->unpack(mfa2buff))
    {
        delete mfa2pkg;
        mfa2pkg = NULL;
        return NULL;
    }
    mfa2pkg->setBufferAndZipOffset(mfa2buff.getBuffer(), mfa2buff.getSize(), mfa2buff.tell());
    return mfa2pkg;
}

bool MFA2::unpack(Mfa2Buffer& buff)
{
    if (!_fingerPrint.unpack(buff))
    {
        return false;
    }
    _packageDescriptor.unpack(buff);
    u_int16_t devCount = _packageDescriptor.getDeviceDescriptorsCount();
    u_int16_t compCount = _packageDescriptor.getComponentsCount();

    for (u_int16_t devIdx = 0; devIdx < devCount; devIdx++)
    {
        vector<ComponentPointerExtension> componentPointers;
        PSIDExtension psidExt("");
        DeviceDescriptor devDescriptor(componentPointers, psidExt);
        devDescriptor.unpack(buff);
        _deviceDescriptors.push_back(devDescriptor);
    }

    for (u_int16_t compIdx = 0; compIdx < compCount; compIdx++)
    {
        VersionExtension version("");
        ComponentDescriptor compDescriptor(version, "");
        compDescriptor.unpack(buff);
        Component comp(compDescriptor);
        _components.push_back(comp);
    }

    return true;
}

void MFA2::minidump()
{
    printf("Package Version : %s\n", _packageDescriptor.getVersionExtension().getVersion(false).c_str());
}

void MFA2::dump()
{
    // printf("Finger print    : %s\n", _fingerPrint.toString().c_str());
    u_int16_t devCount = _packageDescriptor.getDeviceDescriptorsCount();
    // u_int16_t compCount = _packageDescriptor.getComponentsCount();
    printf("Package Version : %s\n", _packageDescriptor.getVersionExtension().getVersion(false).c_str());
    printf("Creation Time   : %s\n", _packageDescriptor.getVersionExtension().getDateAndTime().c_str());
    // printf("Extensions      : %u\n", _packageDescriptor.getExtensionsCount());
    printf("Devices         : %u\n", devCount);
    for (u_int16_t index = 0; index < devCount; index++)
    {
        DeviceDescriptor& devDescriptor = _deviceDescriptors[index];
        string psid = devDescriptor.getPSIDExtension().getString();
        u_int8_t compPtrCount = devDescriptor.getComponentPointerExtensionsCount();
        printf("    PSID       : %s\n", psid.c_str());
        printf("    Num of Images : %u\n", compPtrCount);
        for (u_int8_t comp = 0; comp < compPtrCount; comp++)
        {
            const ComponentPointerExtension& compPtr = devDescriptor.getComponentPointerExtension(comp);
            u_int16_t compIndex = compPtr.getComponentIndex();
            Component compObj = _components[compIndex];
            const ComponentDescriptor& compDescr = compObj.getComponentDescriptor();
            printf("        Index   : %u\n", compIndex);
            printf("        Version : %s\n", compDescr.getVersionExtension().getVersion(true).c_str());
            char buffer[32] = {0};
            compDescr.getVersionExtension().getDateAndTime(buffer);
            printf("        Date    : %s\n", buffer);
        }
    }
}

bool MFA2::extractComponent(Component* requiredComponent, vector<u_int8_t>& fwBinaryData)
{
    PackageDescriptor packageDescriptor = getPackageDescriptor();
    u_int64_t totalSize = packageDescriptor.getComponentsBlockSize();
    u_int32_t zipOffset = packageDescriptor.getComponentsBlockOffset();

    vector<u_int8_t> unzippedBlockBuff;
    unzippedBlockBuff.resize(totalSize);
    vector<u_int8_t> zippedBuff = getBuffer();
    u_int8_t* zippedData = (u_int8_t*)zippedBuff.data() + zipOffset;
    int32_t retVal =
      xz_decompress_crc32(zippedData, zippedBuff.size() - zipOffset, unzippedBlockBuff.data(), totalSize);
    if (retVal != (int32_t)totalSize)
    {
        printf("Decompress error occurred %s\n", xz_get_error(retVal));
        return false;
    }
    u_int32_t requiredOffset = requiredComponent->getBinaryComponentOffset();
    u_int32_t componentBinarySize = requiredComponent->getComponentBinarySize() - strlen(FINGERPRINT_MFA2);
    fwBinaryData.resize(componentBinarySize);
    for (unsigned int index = 0; index < componentBinarySize; index++)
    {
        fwBinaryData[index] =
          unzippedBlockBuff[requiredOffset + index + strlen(FINGERPRINT_MFA2)]; // skip 16 bytes from decompressed file
                                                                                // !
    }
    return true;
}

bool MFA2::unzipLatestVersionComponent(map_string_to_component& matchingComponentsMap, vector<u_int8_t>& fwBinaryData)
{
    map_string_to_component::iterator itAtKey = matchingComponentsMap.find(_latestComponentKey);
    if (itAtKey == matchingComponentsMap.end())
    {
        return false;
    }
    Component* requiredComponent = &itAtKey->second;
    return extractComponent(requiredComponent, fwBinaryData);
}

bool MFA2::unzipComponent(map_string_to_component& matchingComponentsMap,
                          u_int32_t choice,
                          vector<u_int8_t>& fwBinaryData)
{
    if (choice >= matchingComponentsMap.size())
    {
        return false;
    }
    map_string_to_component::iterator itAtOffset = matchingComponentsMap.begin();
    std::advance(itAtOffset, choice);
    Component* requiredComponent = &itAtOffset->second;
    return extractComponent(requiredComponent, fwBinaryData);
}

map_string_to_component MFA2::getMatchingComponents(char* device_psid, u_int16_t fw_ver[3])
{
    map_string_to_component matchingComponentsMap;
    PackageDescriptor packageDescriptor = getPackageDescriptor();
    u_int16_t devCount = packageDescriptor.getDeviceDescriptorsCount();
    Version currentLatestVersion;
    u_int8_t deviceMajorVer = (u_int8_t)fw_ver[0];
    u_int8_t deviceMinorVer = (u_int8_t)fw_ver[1];
    u_int16_t deviceSubMinorVer = fw_ver[2];
    for (u_int16_t index = 0; index < devCount; index++)
    {
        DeviceDescriptor devDescriptor = getDeviceDescriptor(index);
        string psid = devDescriptor.getPSIDExtension().getString();
        if (device_psid != NULL)
        {
            if (strcmp((char*)psid.c_str(), device_psid))
            {
                continue;
            }
        }
        u_int8_t compPtrCount = devDescriptor.getComponentPointerExtensionsCount();
        for (u_int8_t comp = 0; comp < compPtrCount; comp++)
        {
            const ComponentPointerExtension& compPtr = devDescriptor.getComponentPointerExtension(comp);
            u_int16_t compIndex = compPtr.getComponentIndex();
            Component compObj = getComponentObject(compIndex);
            const ComponentDescriptor& compDescr = compObj.getComponentDescriptor();
            u_int8_t majorVer = compDescr.getVersionExtension().getMajor();
            u_int8_t minorVer = compDescr.getVersionExtension().getMinor();
            u_int16_t subMinorVer = compDescr.getVersionExtension().getSubMinor();
            if (deviceMajorVer != 0xff)
            {
                if (deviceMajorVer != majorVer)
                {
                    continue;
                }
            }
            if (deviceMinorVer != 0xff)
            {
                if (minorVer != deviceMinorVer)
                {
                    continue;
                }
            }
            if (deviceSubMinorVer != 0xffff)
            {
                if (subMinorVer != deviceSubMinorVer)
                {
                    continue;
                }
            }
            char dateTimeBuffer[32] = {0};
            compDescr.getVersionExtension().getDateAndTime(dateTimeBuffer);
            VersionExtension currentVersion = compDescr.getVersionExtension();
            string currentVersionString = currentVersion.getVersion(true);
            string mapKey = currentVersionString + "  " + dateTimeBuffer;
            matchingComponentsMap.insert(std::pair<string, Component>(mapKey, compObj));
            Version version(currentVersionString);
            // find the latest matching FW version
            if (version > currentLatestVersion)
            {
                currentLatestVersion = version;
                _latestComponentKey = mapKey;
            }
        }
    }
    return matchingComponentsMap;
}