#include "tool.h"
#include "print.h"

#include <maps/libs/common/include/retry.h>
#include <maps/libs/http/include/request.h>
#include <maps/libs/common/include/file_utils.h>
#include <yandex/maps/proto/firmware_updater/storage/firmware.pb.h>
#include <yandex/maps/proto/firmware_updater/storage/upload.pb.h>

#include <fstream>

namespace proto = yandex::maps::proto::firmware_updater::storage;

namespace maps::fw_updater::storage::cli {

namespace {

std::string startUpload(
    Context& ctx,
    const std::string& hardware,
    const std::string& slot,
    const std::string& version)
{
    http::URL url(ctx.baseUrl());
    url.setPath("/v2/firmware/upload/start")
        .addParam("hardware", hardware)
        .addParam("slot", slot)
        .addParam("version", version);

    auto response = common::retry([&] {
            http::Request request(ctx.httpClient(), http::PUT, url);
            addAuth(request, ctx);
            auto response = request.perform();
            return response;
        },
        ctx.retryPolicy(),
        ctx.responseValidator(),
        ctx.errorReporter());

    REQUIRE(response.status() == 201, errorMessage(response));
    proto::Upload upload;
    Y_PROTOBUF_SUPPRESS_NODISCARD upload.ParseFromString(TString(response.readBody()));
    return upload.id();
}

void uploadPart(
    Context& ctx,
    const std::string& uploadId,
    size_t partNumber,
    std::string&& content)
{
    http::URL url(ctx.baseUrl());
    url.setPath("/v2/firmware/upload/add_part")
        .addParam("id", uploadId)
        .addParam("part", partNumber);

    INFO() << "Uploading part " << partNumber << ", size = " << content.size();

    auto response = common::retry([&] {
            http::Request request(ctx.httpClient(), http::PUT, url);
            addAuth(request, ctx);
            request.setContent(content);
            auto response = request.perform();
            return response;
        },
        ctx.retryPolicy(),
        ctx.responseValidator(),
        ctx.errorReporter());

    REQUIRE(response.status() == 200, errorMessage(response));
    proto::UploadPart part;
    REQUIRE(part.ParseFromString(TString(response.readBody())),
            "Failed to parse server response on /add_part");
    INFO() << "Part " << partNumber << " uploaded";
}

proto::Firmware finishUpload(Context& ctx, const std::string& uploadId)
{
    http::URL url(ctx.baseUrl());
    url.setPath("/v2/firmware/upload/finish")
        .addParam("id", uploadId);

    auto response = common::retry([&] {
            http::Request request(ctx.httpClient(), http::POST, url);
            addAuth(request, ctx);
            auto response = request.perform();
            return response;
        },
        ctx.retryPolicy(),
        ctx.responseValidator(),
        ctx.errorReporter());

    REQUIRE(response.status() == 200, errorMessage(response));
    proto::Firmware firmware;
    Y_PROTOBUF_SUPPRESS_NODISCARD firmware.ParseFromString(TString(response.readBody()));
    return firmware;
}

void abortUpload(Context& ctx, const std::string& uploadId)
{
    http::URL url(ctx.baseUrl());
    url.setPath("/v2/firmware/upload/delete")
        .addParam("id", uploadId);

    auto response = common::retry([&] {
            http::Request request(ctx.httpClient(), http::DELETE, url);
            addAuth(request, ctx);
            auto response = request.perform();
            return response;
        },
        ctx.retryPolicy(),
        ctx.responseValidator(),
        ctx.errorReporter());

    REQUIRE(response.status() == 200, errorMessage(response));
}


void uploadFirmware(
    Context& ctx,
    const std::string& hardware,
    const std::string& slot,
    const std::string& version,
    const std::string& filePath)
{
    constexpr size_t CHUNK_SIZE = 50 * 1024 * 1024; // 50 MB

    const auto size = maps::common::getFileSize(filePath);
    REQUIRE(size > 0, "Empty or invalid file: " << filePath);

    std::filebuf file;
    REQUIRE(file.open(filePath.c_str(), std::ios_base::in | std::ios_base::binary),
        "Error while opening file " << filePath);

    auto uploadId = startUpload(ctx, hardware, slot, version);

    size_t totalBytesRead = 0;
    try {
        for (size_t partNumber = 1; totalBytesRead < size; ++partNumber) {
            std::string chunk(CHUNK_SIZE, '\0');
            const auto bytesRead = file.sgetn(reinterpret_cast<char*>(&chunk[0]), CHUNK_SIZE);
            totalBytesRead += bytesRead;
            REQUIRE(bytesRead == CHUNK_SIZE || totalBytesRead == size,
                    "Error while reading file: " << filePath);
            chunk.resize(bytesRead);
            uploadPart(ctx, uploadId, partNumber, std::move(chunk));
        }
        INFO() << "Finishing upload...";
        auto firmware = finishUpload(ctx, uploadId);
        print(firmware);
    } catch(const std::exception& e) {
        ERROR() << "Error while uploading file: " << e.what();
        abortUpload(ctx, uploadId);
        return;
    }
}


void deleteFirmware(
    Context& ctx,
    const std::string& hardware,
    const std::string& slot,
    const std::string& version)
{
    http::URL url(ctx.baseUrl());
    url.setPath("/v2/firmware/delete")
        .addParam("hardware", hardware)
        .addParam("slot", slot)
        .addParam("version", version);

    auto response = common::retry([&] {
            http::Request request(ctx.httpClient(), http::DELETE, url);
            addAuth(request, ctx);
            auto response = request.perform();
            return response;
        },
        ctx.retryPolicy(),
        ctx.responseValidator(),
        ctx.errorReporter());

    REQUIRE(response.status() == 200, errorMessage(response));
}


void listFirmwares(
    Context& ctx,
    const std::string& hardware,
    const std::optional<std::string>& slot,
    const std::optional<uint32_t>& results,
    const std::optional<uint32_t>& skip)
{
    http::URL url(ctx.baseUrl());
    url.setPath("/v2/firmware/list")
        .addParam("hardware", hardware);

    if (slot) {
        url.addParam("slot", *slot);
    }
    if (results) {
        url.addParam("results", *results);
    }
    if (skip) {
        url.addParam("skip", *skip);
    }

    auto response = common::retry([&] {
            http::Request request(ctx.httpClient(), http::GET, url);
            addAuth(request, ctx);
            auto response = request.perform();
            return response;
        },
        ctx.retryPolicy(),
        ctx.responseValidator(),
        ctx.errorReporter());

    REQUIRE(response.responseClass() == http::ResponseClass::Success,
            errorMessage(response));
    proto::FirmwaresList firmwaresList;
    Y_PROTOBUF_SUPPRESS_NODISCARD firmwaresList.ParseFromString(TString(response.readBody()));
    print(firmwaresList);
}


} // namespace


void handleFirmwareCommand(
    Context& ctx,
    const Options& options)
{
    ASSERT(options.command() == Command::Firmware);

    auto subCommand = options.subCommand<FirmwareSubCommand>();
    switch(subCommand) {
        case FirmwareSubCommand::Upload:
            uploadFirmware(ctx, options.hardware(), options.slot(), options.version(), options.filePath());
            break;
        case FirmwareSubCommand::Delete:
            deleteFirmware(ctx, options.hardware(), options.slot(), options.version());
            break;
        case FirmwareSubCommand::List:
            listFirmwares(ctx, options.hardware(), options.slot.optional(),
                          options.results.optional(), options.skip.optional());
            break;
    }
}

} // namespace maps::fw_updater::storage::cli
