/***********************************************************************************************************************************
S3 Storage File Write
***********************************************************************************************************************************/
#include <build.h>

#include "common/debug.h"
#include "common/log.h"
#include "common/type/object.h"
#include "common/type/xml.h"
#include "storage/s3/write.h"
#include "storage/write.h"

/***********************************************************************************************************************************
Part defaults based on limits at https://docs.aws.amazon.com/AmazonS3/latest/userguide/qfacts.html
***********************************************************************************************************************************/
#define STORAGE_S3_SPLIT_DEFAULT                                    205
#define STORAGE_S3_SPLIT_MAX                                        9557

/***********************************************************************************************************************************
S3 query tokens
***********************************************************************************************************************************/
STRING_STATIC(S3_QUERY_PART_NUMBER_STR,                             "partNumber");
STRING_STATIC(S3_QUERY_UPLOADS_STR,                                 "uploads");
STRING_STATIC(S3_QUERY_UPLOAD_ID_STR,                               "uploadId");

/***********************************************************************************************************************************
XML tags
***********************************************************************************************************************************/
STRING_STATIC(S3_XML_TAG_ETAG_STR,                                  "ETag");
STRING_STATIC(S3_XML_TAG_UPLOAD_ID_STR,                             "UploadId");
STRING_STATIC(S3_XML_TAG_COMPLETE_MULTIPART_UPLOAD_STR,             "CompleteMultipartUpload");
STRING_STATIC(S3_XML_TAG_PART_STR,                                  "Part");
STRING_STATIC(S3_XML_TAG_PART_NUMBER_STR,                           "PartNumber");

/***********************************************************************************************************************************
Object type
***********************************************************************************************************************************/
struct StorageWriteS3
{
    const StorageWriteInterface *interface;                         // Interface
    StorageS3 *storage;                                             // Storage that created this object
    const String *name;                                             // File name

    HttpRequest *request;                                           // Async request
    size_t partSize;
    Buffer *partBuffer;
    const String *uploadId;
    StringList *uploadPartList;
};

/***********************************************************************************************************************************
Flush bytes to upload part
***********************************************************************************************************************************/
static void
storageWriteS3Part(StorageWriteS3 *const this)
{
    FUNCTION_LOG_BEGIN(logLevelTrace);
        FUNCTION_LOG_PARAM(STORAGE_WRITE_S3, this);
    FUNCTION_LOG_END();

    ASSERT(this != NULL);

    // If there is an outstanding async request then wait for the response and store the part id
    if (this->request != NULL)
    {
        HttpResponse *const response = storageS3ResponseP(this->request);

        strLstAdd(this->uploadPartList, httpHeaderGet(httpResponseHeader(response), HTTP_HEADER_ETAG_STR));
        ASSERT(strLstGet(this->uploadPartList, strLstSize(this->uploadPartList) - 1) != NULL);

        httpResponseFree(response);
        httpRequestFree(this->request);
        this->request = NULL;
    }

    FUNCTION_LOG_RETURN_VOID();
}

static void
storageWriteS3PartAsync(StorageWriteS3 *const this)
{
    FUNCTION_LOG_BEGIN(logLevelTrace);
        FUNCTION_LOG_PARAM(STORAGE_WRITE_S3, this);
    FUNCTION_LOG_END();

    ASSERT(this != NULL);
    ASSERT(this->partBuffer != NULL);
    ASSERT(bufSize(this->partBuffer) > 0);

    MEM_CONTEXT_TEMP_BEGIN()
    {
        // Complete prior async request, if any
        storageWriteS3Part(this);

        // Get the upload id if we have not already
        if (this->uploadId == NULL)
        {
            // Initiate mult-part upload
            const XmlNode *const xmlRoot = xmlDocumentRoot(
                xmlDocumentNewBuf(
                    httpResponseContent(
                        storageS3RequestP(
                            this->storage, HTTP_VERB_POST_STR, this->name,
                            .query = httpQueryAdd(httpQueryNewP(), S3_QUERY_UPLOADS_STR, EMPTY_STR), .sseKms = true,
                            .sseC = true, .tag = true))));

            // Store the upload id
            MEM_CONTEXT_OBJ_BEGIN(this)
            {
                this->uploadId = xmlNodeContent(xmlNodeChild(xmlRoot, S3_XML_TAG_UPLOAD_ID_STR, true));
                this->uploadPartList = strLstNew();
            }
            MEM_CONTEXT_OBJ_END();
        }

        // Upload the part async
        HttpQuery *const query = httpQueryNewP();
        httpQueryAdd(query, S3_QUERY_UPLOAD_ID_STR, this->uploadId);
        httpQueryAdd(query, S3_QUERY_PART_NUMBER_STR, strNewFmt("%u", strLstSize(this->uploadPartList) + 1));

        MEM_CONTEXT_OBJ_BEGIN(this)
        {
            this->request = storageS3RequestAsyncP(
                this->storage, HTTP_VERB_PUT_STR, this->name, .query = query, .content = this->partBuffer, .sseC = true);
        }
        MEM_CONTEXT_OBJ_END();
    }
    MEM_CONTEXT_TEMP_END();

    FUNCTION_LOG_RETURN_VOID();
}

/***********************************************************************************************************************************
Write to internal buffer
***********************************************************************************************************************************/
static void
storageWriteS3(THIS_VOID, const Buffer *const buffer)
{
    THIS(StorageWriteS3);

    FUNCTION_LOG_BEGIN(logLevelTrace);
        FUNCTION_LOG_PARAM(STORAGE_WRITE_S3, this);
        FUNCTION_LOG_PARAM(BUFFER, buffer);
    FUNCTION_LOG_END();

    ASSERT(this != NULL);
    ASSERT(buffer != NULL);

    // Resize chunk buffer
    storageWriteChunkBufferResize(buffer, this->partBuffer, this->partSize);

    // Continue until the write buffer has been exhausted
    size_t bytesTotal = 0;

    do
    {
        // Copy as many bytes as possible into the part buffer
        const size_t bytesNext =
            bufRemains(this->partBuffer) > bufUsed(buffer) - bytesTotal ?
                bufUsed(buffer) - bytesTotal : bufRemains(this->partBuffer);

        bufCatSub(this->partBuffer, buffer, bytesTotal, bytesNext);
        bytesTotal += bytesNext;

        // If the part buffer is full then write it
        if (bufUsed(this->partBuffer) == this->partSize)
        {
            storageWriteS3PartAsync(this);
            bufUsedZero(this->partBuffer);

            this->partSize = storageWriteChunkSize(
                this->partSize, STORAGE_S3_SPLIT_DEFAULT, STORAGE_S3_SPLIT_MAX, strLstSize(this->uploadPartList));
        }
    }
    while (bytesTotal != bufUsed(buffer));

    FUNCTION_LOG_RETURN_VOID();
}

/***********************************************************************************************************************************
Close the file
***********************************************************************************************************************************/
static void
storageWriteS3Close(THIS_VOID)
{
    THIS(StorageWriteS3);

    FUNCTION_LOG_BEGIN(logLevelTrace);
        FUNCTION_LOG_PARAM(STORAGE_WRITE_S3, this);
    FUNCTION_LOG_END();

    ASSERT(this != NULL);

    // Close if the file has not already been closed
    if (this->partBuffer != NULL)
    {
        MEM_CONTEXT_TEMP_BEGIN()
        {
            // If a multi-part upload was started we need to finish that way
            if (this->uploadId != NULL)
            {
                // If there is anything left in the part buffer then write it
                if (!bufEmpty(this->partBuffer))
                    storageWriteS3PartAsync(this);

                // Complete prior async request, if any
                storageWriteS3Part(this);

                // Generate the xml part list
                XmlDocument *const partList = xmlDocumentNew(S3_XML_TAG_COMPLETE_MULTIPART_UPLOAD_STR);

                for (unsigned int partIdx = 0; partIdx < strLstSize(this->uploadPartList); partIdx++)
                {
                    XmlNode *const partNode = xmlNodeAdd(xmlDocumentRoot(partList), S3_XML_TAG_PART_STR);
                    xmlNodeContentSet(xmlNodeAdd(partNode, S3_XML_TAG_PART_NUMBER_STR), strNewFmt("%u", partIdx + 1));
                    xmlNodeContentSet(xmlNodeAdd(partNode, S3_XML_TAG_ETAG_STR), strLstGet(this->uploadPartList, partIdx));
                }

                // Finalize the multi-part upload
                HttpRequest *const request = storageS3RequestAsyncP(
                    this->storage, HTTP_VERB_POST_STR, this->name,
                    .query = httpQueryAdd(httpQueryNewP(), S3_QUERY_UPLOAD_ID_STR, this->uploadId),
                    .content = xmlDocumentBuf(partList));
                HttpResponse *const response = storageS3ResponseP(request);

                // Error when no etag in the result. This indicates that the request did not succeed despite the success code.
                if (xmlNodeChild(
                        xmlDocumentRoot(xmlDocumentNewBuf(httpResponseContent(response))), S3_XML_TAG_ETAG_STR, false) == NULL)
                {
                    httpRequestError(request, response);
                }
            }
            // Else upload all the data in a single put
            else
            {
                storageS3RequestP(
                    this->storage, HTTP_VERB_PUT_STR, this->name, .content = this->partBuffer, .sseKms = true,
                    .sseC = true, .tag = true);
            }

            bufFree(this->partBuffer);
            this->partBuffer = NULL;
        }
        MEM_CONTEXT_TEMP_END();
    }

    FUNCTION_LOG_RETURN_VOID();
}

/**********************************************************************************************************************************/
static const StorageWriteInterface storageWriteS3Interface =
{
    .close = storageWriteS3Close,
    .write = storageWriteS3,
};

FN_EXTERN StorageWriteS3 *
storageWriteS3New(StorageS3 *const storage, const String *const name, const size_t partSize)
{
    FUNCTION_LOG_BEGIN(logLevelTrace);
        FUNCTION_LOG_PARAM(STORAGE_S3, storage);
        FUNCTION_LOG_PARAM(STRING, name);
        FUNCTION_LOG_PARAM(SIZE, partSize);
    FUNCTION_LOG_END();

    ASSERT(storage != NULL);
    ASSERT(name != NULL);

    OBJ_NEW_BEGIN(StorageWriteS3, .childQty = MEM_CONTEXT_QTY_MAX)
    {
        *this = (StorageWriteS3)
        {
            .interface = &storageWriteS3Interface,
            .storage = storage,
            .name = strDup(name),
            .partSize = partSize,
            .partBuffer = bufNew(0),
        };
    }
    OBJ_NEW_END();

    FUNCTION_LOG_RETURN(STORAGE_WRITE_S3, this);
}
