// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <getopt.h>

#include <flasher/device.hpp>
#include <flashupdate/args.hpp>
#include <flashupdate/config.hpp>
#include <flashupdate/flash.hpp>
#include <flashupdate/info.hpp>
#include <flashupdate/validator.hpp>
#include <flashupdate/validator/cr51.hpp>
#include <stdplus/print.hpp>

#include <charconv>
#include <optional>
#include <stdexcept>
#include <string>
#include <string_view>
#include <unordered_map>

namespace
{

std::optional<uint8_t> getSecondaryIndex(std::string_view arg)
{
    if (arg == "primary")
    {
        return std::nullopt;
    }
    if (!arg.starts_with("secondary/"))
    {
        throw std::runtime_error(
            "FLASH_TYPE must be primary or secondary/<partition number>");
    }

    std::string_view indexStr =
        arg.substr(std::string_view("secondary/").size());

    uint32_t index;
    auto [ptr, ec]{std::from_chars(indexStr.begin(), indexStr.end(), index)};
    if (ec != std::errc())
    {
        throw std::runtime_error(
            std::format("failed to convert string `{}` to uint32_t: {}",
                        indexStr, std::make_error_code(ec).message()));
    }
    if (ptr != indexStr.end())
    {
        throw std::runtime_error(
            std::format("converted invalid characters: {}", indexStr));
    }

    return index;
}

} // namespace

namespace flashupdate
{

const std::unordered_map<std::string_view, Args::Op> stringToOp = {
    {"hash_descriptor", Args::Op::HashDescriptor},
    {"info", Args::Op::Info},
    {"invalidate", Args::Op::Invalidate},
    {"read", Args::Op::Read},
    {"update_state", Args::Op::UpdateState},
    {"update_version", Args::Op::UpdateVersion},
    {"validate_config", Args::Op::ValidateConfig},
    {"write", Args::Op::Write},
    {"verify_staging", Args::Op::VerifyStaging},
    {"fetch_version", Args::Op::FetchVersion},
    {"copy_partition", Args::Op::CopyPartition},
    {"erase", Args::Op::Erase},
};

void Args::setValidatorHelper(validator::Validator* validatorHelper)
{
    this->validatorHelper = validatorHelper;
}

void Args::setFlashHelper(flash::Flash* flashHelper)
{
    this->flashHelper = flashHelper;
}

Args::Args()
{
    validatorHelperPtr = std::make_unique<validator::cr51::Cr51>();
    validatorHelper = validatorHelperPtr.get();

    flashHelperPtr = std::make_unique<flash::Flash>();
    flashHelper = flashHelperPtr.get();
}

Args::Args(int argc, char* argv[])
{
    static const char opts[] = ":j:i:acIfkSsvxz";
    static const struct option longopts[] = {
        {"active_version", no_argument, nullptr, 'a'},
        {"clean_output", no_argument, nullptr, 'c'},
        {"config", required_argument, nullptr, 'j'},
        {"skip_staging_index_update", no_argument, nullptr, 'x'},
        {"skip_validation", no_argument, nullptr, 'z'},
        {"stage_state", no_argument, nullptr, 'S'},
        {"staged_version", no_argument, nullptr, 's'},
        {"staged_index", no_argument, nullptr, 'I'},
        {"keep_mux", no_argument, nullptr, 'k'},
        {"verbose", no_argument, nullptr, 'v'},
        {nullptr, 0, nullptr, 0},
    };

    int c;
    optind = 0;
    while ((c = getopt_long(argc, argv, opts, longopts, nullptr)) > 0)
    {
        switch (c)
        {
            case 'a':
                checkActiveVersion = true;
                break;
            case 's':
                checkStagedVersion = true;
                break;
            case 'S':
                checkStageState = true;
                break;
            case 'v':
                verbose++;
                break;
            case 'c':
                cleanOutput = true;
                break;
            case 'k':
                keepMux = true;
                break;
            case 'I':
                checkStagedIndex = true;
                break;
            case 'j':
                configFile.emplace(optarg);
                break;
            case 'x':
                updateStagingIndex = false;
                break;
            case 'z':
                skipValidation = true;
                break;
            case ':':
                throw std::runtime_error(
                    std::format("Missing argument for `{}`", argv[optind - 1]));
                break;
            default:
                throw std::runtime_error(std::format(
                    "Invalid command line argument `{}`", argv[optind - 1]));
        }
    }

    // Enabled all UpdateInfo output when using `Info` command
    // If no specific targeted information, print out all avaliable information
    // for general debugging purpose.
    if (!checkActiveVersion && !checkStagedVersion && !checkStageState &&
        !checkStagedIndex)
    {
        checkActiveVersion = true;
        checkStagedVersion = true;
        checkStageState = true;
        checkStagedIndex = true;
        otherInfo = true;
    }

    if (optind == argc)
    {
        throw std::runtime_error("Missing flashupdate operation");
    }
    auto it = stringToOp.find(argv[optind]);
    if (it == stringToOp.end())
    {
        throw std::runtime_error(
            std::format("Invalid operation: {}", argv[optind]));
    }

    optind++;
    op = it->second;

    std::optional<uint8_t> secondaryIndex = std::nullopt;
    switch (op)
    {
        case Args::Op::ValidateConfig:
            break;
        case Args::Op::HashDescriptor:
            printHelp = printHashDescriptorHelp;

            if (optind + 1 != argc)
            {
                throw std::runtime_error("Must specify FILE");
            }
            file.emplace(argv[optind]);
            break;
        case Args::Op::Read:
            printHelp = printReadHelp;

            if (optind + 2 != argc)
            {
                throw std::runtime_error("Must specify FLASH_TYPE and FILE");
            }

            secondaryIndex = getSecondaryIndex(argv[optind]);
            primary = !secondaryIndex.has_value();
            stagingIndex = secondaryIndex.value_or(0);
            file.emplace(argv[optind + 1]);
            break;
        case Args::Op::Write:
            printHelp = printWriteHelp;

            if (optind + 2 != argc)
            {
                throw std::runtime_error("Must specify FILE and FLASH_TYPE");
            }
            file.emplace(argv[optind]);
            secondaryIndex = getSecondaryIndex(argv[optind + 1]);
            primary = !secondaryIndex.has_value();
            stagingIndex = secondaryIndex.value_or(0);
            break;
        case Args::Op::UpdateState:
            printHelp = printUpdateStateHelp;

            if (optind + 1 != argc)
            {
                throw std::runtime_error("Must specify STATE");
            }

            state = argv[optind];
            break;
        case Args::Op::UpdateVersion:
            printHelp = printUpdateVersionHelp;

            if (optind + 2 != argc)
            {
                throw std::runtime_error("Must specify VERSION and FLASH_TYPE");
            }
            version = argv[optind];
            secondaryIndex = getSecondaryIndex(argv[optind + 1]);
            primary = !secondaryIndex.has_value();
            stagingIndex = secondaryIndex.value_or(0);
            break;
        case Args::Op::Info:
            printHelp = printInfoHelp;
            break;
        case Args::Op::Invalidate:
            printHelp = [this](const char* arg0) { printInvalidateHelp(arg0); };
            if (optind + 1 != argc)
            {
                throw std::runtime_error("Must specify FLASH_TYPE");
            }
            secondaryIndex = getSecondaryIndex(argv[optind]);
            primary = !secondaryIndex.has_value();
            stagingIndex = secondaryIndex.value_or(0);

            break;
        case Args::Op::VerifyStaging:
            printHelp = printVerifyStagingHelp;
            break;
        case Args::Op::FetchVersion:
            printHelp = printFetchVersionHelp;

            if (optind + 1 != argc)
            {
                throw std::runtime_error("Must specify FLASH_TYPE");
            }
            secondaryIndex = getSecondaryIndex(argv[optind]);
            primary = !secondaryIndex.has_value();
            stagingIndex = secondaryIndex.value_or(0);
            break;
        case Args::Op::CopyPartition:
            printHelp = printCopyPartitionHelp;

            if (optind + 2 != argc)
            {
                throw std::runtime_error("Must specify two FLASH_TYPEs");
            }
            fromPartition = getSecondaryIndex(argv[optind]);
            toPartition = getSecondaryIndex(argv[optind + 1]);
            if (fromPartition == toPartition)
            {
                throw std::runtime_error("From and to partition are the same.");
            }

            if (toPartition.has_value())
            {
                // Anything to Secondary
                primary = false;
                stagingIndex = toPartition.value();
            }
            else
            {
                // Secondary to Primary
                primary = true;
                stagingIndex = fromPartition.value();
            }
            break;
        case Args::Op::Erase:
            printHelp = printEraseHelp;

            if (optind + 1 != argc)
            {
                throw std::runtime_error("Must specify the FLASH_TYPEs");
            }

            auto partition = getSecondaryIndex(argv[optind]);
            if (!partition.has_value())
            {
                throw std::runtime_error(
                    "Erase command is not supported on primary partition.");
            }
            primary = false;
            stagingIndex = *partition;
            break;
    };

    config = createConfig(configFile, stagingIndex);

    // The validatorHelper will be set to Cr51 for now.
    if (config.cr51.has_value())
    {
        validatorHelperPtr =
            std::make_unique<validator::cr51::Cr51>(*config.cr51);
        validatorHelper = validatorHelperPtr.get();
    }

    // The flashHelper needs the config to be valid before setting it up.
    flashHelperPtr = std::make_unique<flash::Flash>(config, keepMux);
    flashHelper = flashHelperPtr.get();
}

void Args::printHashDescriptorHelp(const char* arg0)
{
    stdplus::println(stderr, "Usage: {} [OPTION]... hash_descriptor FILE",
                     arg0);
    stdplus::println(stderr, "");

    stdplus::println(stderr, "Ex: {} hash_descriptor image.bin", arg0);
    stdplus::println(stderr, "");
}

void Args::printUpdateStateHelp(const char* arg0)
{
    stdplus::println(stderr, "Usage: {} [OPTION]... update_state STATE", arg0);
    stdplus::println(stderr, "");

    stdplus::println(stderr, "STATE options");
    stdplus::println(stderr, "{}", info::listStates());

    stdplus::println(stderr, "");

    stdplus::println(stderr, "Ex: {} update_state STAGED", arg0);
    stdplus::println(stderr, "");
}

void Args::printUpdateVersionHelp(const char* arg0)
{
    stdplus::println(stderr,
                     "Usage: {} [OPTION]... update_version VERSION FLASH_TYPE",
                     arg0);
    stdplus::println(stderr, "");

    stdplus::println(stderr, "VERSION format");
    stdplus::println(stderr, "  `x.y.z.w`");
    stdplus::println(stderr, "");
    stdplus::println(stderr, "FLASH_TYPE options");
    stdplus::println(stderr, "  `primary`");
    stdplus::println(stderr, "  `secondary/<partition number>`");
    stdplus::println(stderr, "");

    stdplus::println(stderr, "Ex: {} update_version 10.10.10.0 secondary/0",
                     arg0);
    stdplus::println(stderr, "");
}

void Args::printWriteHelp(const char* arg0)
{
    stdplus::println(stderr, "Usage: {} [OPTION]... write FILE FLASH_TYPE",
                     arg0);
    stdplus::println(stderr, "");

    stdplus::println(stderr, "FLASH_TYPE options");
    stdplus::println(stderr, "  `primary`");
    stdplus::println(stderr, "  `secondary/<partition number>`");
    stdplus::println(stderr, "");

    stdplus::println(stderr, "Ex: {} write image.bin primary", arg0);
    stdplus::println(stderr, "");
}

void Args::printReadHelp(const char* arg0)
{
    stdplus::println(stderr, "Usage: {} [OPTION]... read FLASH_TYPE FILE",
                     arg0);
    stdplus::println(stderr, "");

    stdplus::println(stderr, "FLASH_TYPE options");
    stdplus::println(stderr, "  `primary`");
    stdplus::println(stderr, "  `secondary/<partition number>`");
    stdplus::println(stderr, "");

    stdplus::println(stderr, "Ex: {} read primary image.bin", arg0);
    stdplus::println(stderr, "");
}

void Args::printVerifyStagingHelp(const char* arg0)
{
    stdplus::println(stderr, "Usage: {} verify_staging", arg0);
    stdplus::println(stderr, "");
    stdplus::println(stderr,
                     "Verify the staging partition to make sure it matches the "
                     "version in metadata. ");
    stdplus::println(stderr, "Ex: {} verify_staging", arg0);
    stdplus::println(stderr, "");
}

void Args::printFetchVersionHelp(const char* arg0)
{
    stdplus::println(stderr, "Usage: {} fetch_version", arg0);
    stdplus::println(stderr, "");
    stdplus::println(stderr,
                     "Read the CR51 descriptor of the partition and return "
                     "the validated version ");

    stdplus::println(stderr, "FLASH_TYPE options");
    stdplus::println(stderr, "  `primary`");
    stdplus::println(stderr, "  `secondary/<partition number>`");
    stdplus::println(stderr, "");
    stdplus::println(stderr, "Ex: {} fetch_version", arg0);
    stdplus::println(stderr, "");
}

void Args::printCopyPartitionHelp(const char* arg0)
{
    stdplus::println(stderr, "Usage: {} copy_partition FLASH_TYPE", arg0);
    stdplus::println(stderr, "");
    stdplus::println(stderr, "Copy the firmware from one partition to another");

    stdplus::println(stderr, "FLASH_TYPE options");
    stdplus::println(stderr, "  `primary`");
    stdplus::println(stderr, "  `secondary/<partition number>`");
    stdplus::println(stderr, "");
    stdplus::println(stderr, "Ex: {} copy_partition secondary/1 primary", arg0);
    stdplus::println(stderr, "");
}

void Args::printEraseHelp(const char* arg0)
{
    stdplus::println(stderr, "Usage: {} erase FLASH_TYPE", arg0);
    stdplus::println(stderr, "");
    stdplus::println(stderr, "Erase the target partition (staging only)");

    stdplus::println(stderr, "FLASH_TYPE options");
    stdplus::println(stderr, "  `secondary/<partition number>`");
    stdplus::println(stderr, "");
    stdplus::println(stderr, "Ex: {} erase secondary/1", arg0);
    stdplus::println(stderr, "");
}

void Args::printInfoHelp(const char* arg0)
{
    stdplus::println(stderr, "Usage: {} [OPTION]... info", arg0);
    stdplus::println(stderr, "");

    stdplus::println(stderr, "Optional Arguments for `info` command:");
    stdplus::println(stderr,
                     "  -a, --active_version   Print the active version with "
                     "the `info command");
    stdplus::println(stderr,
                     "  -s, --staged_version   Print the stage version with "
                     "the `info` command");
    stdplus::println(stderr, "  -S, --stage_state      Print the Stage "
                             "stage of the BIOS image.");
    stdplus::println(stderr,
                     "  -I, --staged_index     Print the index of the lasted "
                     "staged secondary partition");
    stdplus::println(stderr, "  -c, --clean_output     Print the `info` "
                             "message with no prefixes");
    stdplus::println(stderr, "");

    stdplus::println(stderr, "Ex: {} info -avS", arg0);
    stdplus::println(stderr, "");
}

void Args::printInvalidateHelp(const char* arg0)
{
    stdplus::println(stderr, "Usage: {} [OPTION]... invalidate FLASH_TYPE",
                     arg0);
    stdplus::println(stderr, "");

    stdplus::println(stderr, "FLASH_TYPE options");
    stdplus::println(stderr, "  `primary`");
    stdplus::println(stderr, "  `secondary`");
    stdplus::println(stderr, "");
    stdplus::println(stderr, "Ex: {} invalidate primary", arg0);
    stdplus::println(stderr, "");
}

std::function<void(const char* arg0)> Args::printHelp = [](const char* arg0) {
    stdplus::println(stderr, "Usage: {} [OPTION]... validate_config", arg0);
    stdplus::println(stderr, "   or: {} [OPTION]... read FLASH_TYPE FILE",
                     arg0);
    stdplus::println(stderr, "   or: {} [OPTION]... write FILE FLASH_TYPE",
                     arg0);
    stdplus::println(stderr, "   or: {} [OPTION]... update_state STATE", arg0);
    stdplus::println(stderr,
                     "   or: {} [OPTION]... update_version VERSION FLASH_TYPE",
                     arg0);
    stdplus::println(stderr, "   or: {} [OPTION]... info", arg0);
    stdplus::println(stderr, "   or: {} [OPTION]... hash_descriptor", arg0);
    stdplus::println(stderr, "   or: {} [OPTION]... verify_staging", arg0);
    stdplus::println(stderr, "   or: {} [OPTION]... invalidate FLASH_TYPE",
                     arg0);
    stdplus::println(stderr, "   or: {} [OPTION]... fetch_version FLASH_TYPE",
                     arg0);
    stdplus::println(
        stderr, "   or: {} [OPTION]... copy_partition FLASH_TYPE FLASH_TYPE",
        arg0);
    stdplus::println(stderr, "   or: {} [OPTION]... erase FLASH_TYPE", arg0);
    stdplus::println(stderr, "");

    stdplus::println(stderr, "Command Specifc Arguments:");
    stdplus::println(stderr, "  Run `{} $CMD` to see the command arguments",
                     arg0);
    stdplus::println(stderr, "");
    stdplus::println(stderr, "General Optional Arguments:");
    stdplus::println(stderr, "  -v, --verbose          Increases the verbosity "
                             "level of error message output");
    stdplus::println(stderr, "  -j, --config[=JSON]    Path for json config. "
                             "For example, `--config bios` will look for "
                             "/usr/share/flashupdate/bios.json. You can also "
                             "target specific config file with full path");
    stdplus::println(stderr,
                     "  -k, --keep_mux         Keep the mux state to select "
                     "the current flash.");
    stdplus::println(stderr, "");

    stdplus::println(stderr, "Ex: {} validate_config", arg0);
    stdplus::println(stderr, "Ex: {} read primary image.bin", arg0);
    stdplus::println(stderr, "Ex: {} write image.bin primary", arg0);
    stdplus::println(stderr, "Ex: {} update_state STAGED", arg0);
    stdplus::println(stderr, "Ex: {} update_version 10.0.0.0 secondary/0",
                     arg0);
    stdplus::println(stderr, "Ex: {} info", arg0);
    stdplus::println(stderr, "Ex: {} hash_descriptor image.bin", arg0);
    stdplus::println(stderr, "Ex: {} verify_staging", arg0);
    stdplus::println(stderr, "Ex: {} invalidate secondary/0", arg0);
    stdplus::println(stderr, "Ex: {} fetch_version secondary/1", arg0);
    stdplus::println(stderr, "Ex: {} copy_partition secondary/1 secondary/0",
                     arg0);
    stdplus::println(stderr, "Ex: {} erase secondary/1", arg0);
    stdplus::println(stderr, "");
};

Args Args::argsOrHelp(int argc, char* argv[])
{
    try
    {
        return Args(argc, argv);
    }
    catch (...)
    {
        printHelp(argv[0]);
        throw;
    }
}

} // namespace flashupdate