// Copyright (C) 2025 The Android Open Source Project
//
// 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 <string.h>
#include <sysexits.h>
#include <unistd.h>
#include <charconv>
#include <optional>
#include <string_view>

#include <android-base/file.h>
#include <android-base/parseint.h>
#include <android-base/unique_fd.h>
#include <liblp/builder.h>
#include <liblp/liblp.h>
#include <iostream>

using android::base::unique_fd;
using android::fs_mgr::LpMetadata;
using android::fs_mgr::MetadataBuilder;
using android::fs_mgr::PartitionOpener;
using android::fs_mgr::ReadMetadata;
using android::fs_mgr::UpdatePartitionTable;

enum class SubCommand {
    kUnknown,
    kAdd,
    kRemove,
    kAddGroup,
    kRemoveGroup,
    kChangeGroup,
};

static int usage(const char* program, SubCommand cmd = SubCommand::kUnknown) {
    std::cerr << program
              << " - Modifies partitions in a super partition block device (in-place metadata "
                 "modification).\n";
    std::cerr << "\n";
    std::cerr << "Usage:\n";
    std::cerr << "  " << program << " add [options] SUPER_DEVICE PARTNAME PARTGROUP\n";
    std::cerr << "  " << program << " remove [options] SUPER_DEVICE PARTNAME\n";
    std::cerr << "\n";

    if (cmd == SubCommand::kUnknown || cmd == SubCommand::kAdd) {
        std::cerr << "Add sub-command options:\n";
        std::cerr << "  SUPER_DEVICE                  Path to the super partition block device to "
                     "modify (e.g., /dev/block/by-name/super).\n";
        std::cerr << "  PARTNAME                      Name of the partition to add or replace in "
                     "SUPER_DEVICE.\n";
        std::cerr
                << "  PARTGROUP                     Name of the partition group to use. If the\n"
                << "                                partition can be updated over OTA, the group\n"
                << "                                should match its updatable group.\n";
        std::cerr << "  --replace                     If PARTNAME already exists, it will be "
                     "replaced.\n";
        std::cerr << "  --size <bytes>                Size in bytes for the new empty partition. "
                     "Defaults to 0 if not specified.\n";
    }
    if (cmd == SubCommand::kUnknown || cmd == SubCommand::kRemove) {
        std::cerr << "\nRemove sub-command options:\n";
        std::cerr << "  SUPER_DEVICE                  Path to the super block device (e.g., "
                     "/dev/block/by-name/super).\n";
        std::cerr << "  PARTNAME                      Name of the partition to remove.\n";
    }
    if (cmd == SubCommand::kUnknown || cmd == SubCommand::kAddGroup) {
        std::cerr << "\naddgrp sub-command options:\n";
        std::cerr << "  SUPER_DEVICE                  Path to the super block device.\n";
        std::cerr << "  GROUPNAME                     Name of the new partition group.\n";
        std::cerr << "  SIZE                          Maximum size in bytes (0 for unlimited).\n";
    }
    if (cmd == SubCommand::kUnknown || cmd == SubCommand::kRemoveGroup) {
        std::cerr << "\nrmgrp sub-command options:\n";
        std::cerr << "  SUPER_DEVICE                  Path to the super block device.\n";
        std::cerr << "  GROUPNAME                     Name of the partition group to remove.\n";
        std::cerr << "  --force                       Remove all partitions in the group first.\n";
    }
    if (cmd == SubCommand::kUnknown || cmd == SubCommand::kChangeGroup) {
        std::cerr << "\nchgrp sub-command options:\n";
        std::cerr << "  SUPER_DEVICE                  Path to the super block device.\n";
        std::cerr << "  PARTNAME                      Name of the partition to move.\n";
        std::cerr << "  GROUPNAME                     Name of the new group for the partition.\n";
    }
    std::cerr << "\nCommon options:\n";
    std::cerr << "  --slot <number>               Metadata slot number to use (default is 0).\n";
    std::cerr << "  -h, --help                    Show this help message.\n";
    std::cerr << "\n";
    return EX_USAGE;
}

class SuperModifyHelper final {
  public:
    explicit SuperModifyHelper(const std::string& super_device_path, uint32_t slot)
        : super_device_path_(super_device_path), slot_(slot) {}

    bool Open();
    bool AddOrReplaceEmptyPartition(const std::string& partition_name,
                                    const std::string& group_name, uint32_t attributes,
                                    uint64_t partition_size, bool replace);
    bool RemovePartition(const std::string& partition_name);
    bool AddGroup(const std::string& group_name, uint64_t maximum_size);
    bool RemoveGroup(const std::string& group_name, bool force);
    bool ChangePartitionGroup(const std::string& partition_name, const std::string& group_name);
    bool Finalize();

  private:
    bool OpenSuperDevice();
    bool UpdateSuperMetadata();

    std::string super_device_path_;
    unique_fd super_device_fd_;
    std::unique_ptr<LpMetadata> metadata_;
    std::unique_ptr<MetadataBuilder> builder_;
    uint32_t slot_;
};

bool SuperModifyHelper::Open() {
    if (!OpenSuperDevice()) {
        return false;
    }
    metadata_ = android::fs_mgr::ReadMetadata(super_device_path_, slot_);
    if (!metadata_) {
        std::cerr << "Could not read metadata from super device: " << super_device_path_ << " slot "
                  << slot_ << "\n";
        return false;
    }
    builder_ = MetadataBuilder::New(*metadata_.get());
    if (!builder_) {
        std::cerr << "Could not create MetadataBuilder for super device: " << super_device_path_
                  << "\n";
        return false;
    }
    return true;
}

bool SuperModifyHelper::AddOrReplaceEmptyPartition(const std::string& partition_name,
                                                   const std::string& group_name,
                                                   uint32_t attributes, uint64_t partition_size,
                                                   bool replace) {
    auto existing_partition = builder_->FindPartition(partition_name);
    if (existing_partition) {
        if (!replace) {
            std::cerr << "Partition " << partition_name
                      << " already exists and --replace not specified. Aborting.\n";
            return false;
        }
        std::cout << "Replacing existing partition: " << partition_name << std::endl;
        builder_->RemovePartition(partition_name);
    }

    auto new_partition = builder_->AddPartition(partition_name, group_name, attributes);
    if (!new_partition) {
        std::cerr << "Could not add partition metadata for: " << partition_name << "\n";
        return false;
    }

    if (!builder_->ResizePartition(new_partition, partition_size)) {
        std::cerr << "Failed to resize partition " << partition_name << " to " << partition_size
                  << " bytes.\n";
        return false;
    }
    std::cout << "Successfully configured partition " << partition_name << " with size "
              << partition_size << " bytes." << std::endl;

    return UpdateSuperMetadata();
}

bool SuperModifyHelper::RemovePartition(const std::string& partition_name) {
    auto partition = builder_->FindPartition(partition_name);
    if (!partition) {
        std::cerr << "Could not find partition to remove: " << partition_name << "\n";
        return false;
    }
    builder_->RemovePartition(partition_name);
    std::cout << "Successfully marked partition " << partition_name << " for removal." << std::endl;
    return UpdateSuperMetadata();
}

bool SuperModifyHelper::AddGroup(const std::string& group_name, uint64_t maximum_size) {
    if (group_name.size() >= sizeof(LpMetadataPartitionGroup::name)) {
        std::cerr << "Group name too long: " << group_name << " (max 35 characters)\n";
        return false;
    }
    if (builder_->FindGroup(group_name)) {
        std::cerr << "Group already exists: " << group_name << "\n";
        return false;
    }
    if (maximum_size > builder_->AllocatableSpace()) {
        std::cerr << "Warning: group size " << maximum_size
                  << " exceeds total allocatable space " << builder_->AllocatableSpace() << "\n";
    }
    if (!builder_->AddGroup(group_name, maximum_size)) {
        std::cerr << "Failed to add group: " << group_name << "\n";
        return false;
    }
    std::cout << "Successfully added group " << group_name << " with max size " << maximum_size
              << " bytes." << std::endl;
    return UpdateSuperMetadata();
}

bool SuperModifyHelper::RemoveGroup(const std::string& group_name, bool force) {
    if (group_name == android::fs_mgr::kDefaultGroup) {
        std::cerr << "Cannot remove the default group.\n";
        return false;
    }
    auto group = builder_->FindGroup(group_name);
    if (!group) {
        std::cerr << "Could not find group to remove: " << group_name << "\n";
        return false;
    }
    auto partitions = builder_->ListPartitionsInGroup(group_name);
    if (!partitions.empty() && !force) {
        std::cerr << "Group " << group_name << " is not empty. Use --force to remove it and its "
                  << partitions.size() << " partitions.\n";
        return false;
    }
    builder_->RemoveGroupAndPartitions(group_name);
    std::cout << "Successfully removed group " << group_name << "." << std::endl;
    return UpdateSuperMetadata();
}

bool SuperModifyHelper::ChangePartitionGroup(const std::string& partition_name,
                                             const std::string& group_name) {
    auto partition = builder_->FindPartition(partition_name);
    if (!partition) {
        std::cerr << "Could not find partition: " << partition_name << "\n";
        return false;
    }
    if (!builder_->FindGroup(group_name)) {
        std::cerr << "Could not find group: " << group_name << "\n";
        return false;
    }
    if (!builder_->ChangePartitionGroup(partition, group_name)) {
        std::cerr << "Failed to change partition " << partition_name << " to group " << group_name
                  << ".\n";
        return false;
    }
    std::cout << "Successfully moved partition " << partition_name << " to group " << group_name
              << "." << std::endl;
    return UpdateSuperMetadata();
}

bool SuperModifyHelper::OpenSuperDevice() {
    super_device_fd_.reset(open(super_device_path_.c_str(), O_RDWR | O_CLOEXEC));
    if (super_device_fd_ < 0) {
        std::cerr << "Open failed for super device: " << super_device_path_ << ": "
                  << strerror(errno) << "\n";
        return false;
    }
    return true;
}

bool SuperModifyHelper::UpdateSuperMetadata() {
    metadata_ = builder_->Export();
    if (!metadata_) {
        std::cerr << "Failed to export new metadata for super device.\n";
        return false;
    }
    PartitionOpener opener;  // Default opener for block devices
    for (uint32_t i = 0; i < metadata_->geometry.metadata_slot_count; i++) {
        if (!UpdatePartitionTable(opener, super_device_path_, *metadata_.get(), i)) {
            std::cerr << "Could not write metadata to super device slot " << i << ".\n";
            return false;
        }
    }
    std::cout << "Successfully updated super partition metadata in all slots." << std::endl;
    return true;
}

bool SuperModifyHelper::Finalize() {
    if (super_device_fd_ >= 0 && fsync(super_device_fd_.get()) < 0) {
        std::cerr << "fsync on super device failed: " << strerror(errno) << "\n";
        return false;
    }
    return true;
}

enum class AddOptionCode : int {
    kReplace = 1,
    kSize = 2,
    kSlot = 3,  // Common option, but needs to be in specific subcommand options for getopt_long
    kHelp = 'h',
};

enum class RemoveOptionCode : int {
    kSlot = 1,  // Common option
    kHelp = 'h',
};

enum class AddGroupOptionCode : int {
    kSlot = 1,
    kHelp = 'h',
};

enum class RemoveGroupOptionCode : int {
    kForce = 1,
    kSlot = 2,
    kHelp = 'h',
};

enum class ChangeGroupOptionCode : int {
    kSlot = 1,
    kHelp = 'h',
};

int HandleAddCommand(int argc, char* argv[], const char* program_name) {
    struct option add_options[] = {
            {"replace", no_argument, nullptr, static_cast<int>(AddOptionCode::kReplace)},
            {"size", required_argument, nullptr, static_cast<int>(AddOptionCode::kSize)},
            {"slot", required_argument, nullptr, static_cast<int>(AddOptionCode::kSlot)},
            {"help", no_argument, nullptr, static_cast<int>(AddOptionCode::kHelp)},
            {nullptr, 0, nullptr, 0},
    };

    bool replace_flag = false;
    uint64_t specified_size = 0;
    uint32_t slot = 0;

    int rv;
    // Reset optind for subcommand parsing
    optind = 1;
    while ((rv = getopt_long(argc, argv, "h", add_options, nullptr)) != -1) {
        switch (rv) {
            case static_cast<int>(AddOptionCode::kHelp):
                usage(program_name, SubCommand::kAdd);
                return EX_OK;
            case static_cast<int>(AddOptionCode::kReplace):
                replace_flag = true;
                break;
            case static_cast<int>(AddOptionCode::kSize): {
                if (!android::base::ParseUint(optarg, &specified_size)) {
                    std::cerr << "Error: Invalid argument for --size: " << optarg
                              << ". Must be a non-negative integer.\n";
                    return usage(program_name, SubCommand::kAdd);
                }
            } break;
            case static_cast<int>(AddOptionCode::kSlot): {
                if (!android::base::ParseUint(optarg, &slot)) {
                    std::cerr << "Error: Invalid argument for --slot: " << optarg
                              << ". Must be a non-negative integer.\n";
                    return usage(program_name, SubCommand::kAdd);
                }
            } break;
            default:
                return usage(program_name, SubCommand::kAdd);
        }
    }

    if (argc - optind != 3) {
        std::cerr << "Error: Incorrect number of arguments for 'add'. Expected SUPER_DEVICE "
                     "PARTNAME PARTGROUP.\n\n";
        return usage(program_name, SubCommand::kAdd);
    }

    std::string super_device_arg = argv[optind++];
    std::string partition_name_arg = argv[optind++];
    std::string group_name_arg = argv[optind++];

    SuperModifyHelper super_modifier(super_device_arg, slot);
    if (!super_modifier.Open()) {
        return EX_SOFTWARE;
    }

    uint32_t attributes = LP_PARTITION_ATTR_NONE;

    if (!super_modifier.AddOrReplaceEmptyPartition(partition_name_arg, group_name_arg, attributes,
                                                   specified_size, replace_flag)) {
        return EX_SOFTWARE;
    }

    if (!super_modifier.Finalize()) {
        return EX_SOFTWARE;
    }

    std::cout << "lpmodify add: Successfully processed partition " << partition_name_arg << " on "
              << super_device_arg << ".\n";
    return EX_OK;
}

int HandleRemoveCommand(int argc, char* argv[], const char* program_name) {
    struct option remove_options[] = {
            {"slot", required_argument, nullptr, static_cast<int>(RemoveOptionCode::kSlot)},
            {"help", no_argument, nullptr, static_cast<int>(RemoveOptionCode::kHelp)},
            {nullptr, 0, nullptr, 0},
    };

    uint32_t slot = 0;

    int rv;
    // Reset optind for subcommand parsing
    optind = 1;
    while ((rv = getopt_long(argc, argv, "h", remove_options, nullptr)) != -1) {
        switch (rv) {
            case static_cast<int>(RemoveOptionCode::kHelp):
                return usage(program_name, SubCommand::kRemove);
            case static_cast<int>(RemoveOptionCode::kSlot): {
                if (!android::base::ParseUint(optarg, &slot)) {
                    std::cerr << "Error: Invalid argument for --slot: " << optarg
                              << ". Must be a non-negative integer.\n";
                    return usage(program_name, SubCommand::kRemove);
                }
            } break;
            default:
                return usage(program_name, SubCommand::kRemove);
        }
    }

    if (argc - optind != 2) {
        std::cerr << "Error: Incorrect number of arguments for 'remove'. Expected SUPER_DEVICE "
                     "PARTNAME.\n\n";
        return usage(program_name, SubCommand::kRemove);
    }

    std::string super_device_path = argv[optind++];
    std::string partition_name = argv[optind++];

    SuperModifyHelper remover(super_device_path, slot);
    if (!remover.Open()) {
        return EX_SOFTWARE;
    }

    if (!remover.RemovePartition(partition_name)) {
        remover.Finalize();
        return EX_SOFTWARE;
    }

    if (!remover.Finalize()) {
        return EX_SOFTWARE;
    }

    std::cout << "lpmodify remove: Partition " << partition_name << " removed successfully from "
              << super_device_path << ".\n";
    return EX_OK;
}

int HandleAddGroupCommand(int argc, char* argv[], const char* program_name) {
    struct option options[] = {
            {"slot", required_argument, nullptr, static_cast<int>(AddGroupOptionCode::kSlot)},
            {"help", no_argument, nullptr, static_cast<int>(AddGroupOptionCode::kHelp)},
            {nullptr, 0, nullptr, 0},
    };

    uint32_t slot = 0;
    int rv;
    optind = 1;
    while ((rv = getopt_long(argc, argv, "h", options, nullptr)) != -1) {
        switch (rv) {
            case static_cast<int>(AddGroupOptionCode::kHelp):
                return usage(program_name, SubCommand::kAddGroup);
            case static_cast<int>(AddGroupOptionCode::kSlot):
                if (!android::base::ParseUint(optarg, &slot)) {
                    std::cerr << "Error: Invalid argument for --slot: " << optarg << "\n";
                    return usage(program_name, SubCommand::kAddGroup);
                }
                break;
            default:
                return usage(program_name, SubCommand::kAddGroup);
        }
    }

    if (argc - optind != 3) {
        std::cerr << "Error: Incorrect number of arguments for 'addgrp'.\n\n";
        return usage(program_name, SubCommand::kAddGroup);
    }

    std::string super_device_path = argv[optind++];
    std::string group_name = argv[optind++];
    uint64_t group_size = 0;
    if (!android::base::ParseUint(argv[optind++], &group_size)) {
        std::cerr << "Error: Invalid group size.\n";
        return usage(program_name, SubCommand::kAddGroup);
    }

    SuperModifyHelper helper(super_device_path, slot);
    if (!helper.Open()) return EX_SOFTWARE;
    if (!helper.AddGroup(group_name, group_size)) return EX_SOFTWARE;
    if (!helper.Finalize()) return EX_SOFTWARE;

    std::cout << "lpmodify addgrp: Group " << group_name << " added successfully.\n";
    return EX_OK;
}

int HandleRemoveGroupCommand(int argc, char* argv[], const char* program_name) {
    struct option options[] = {
            {"force", no_argument, nullptr, static_cast<int>(RemoveGroupOptionCode::kForce)},
            {"slot", required_argument, nullptr, static_cast<int>(RemoveGroupOptionCode::kSlot)},
            {"help", no_argument, nullptr, static_cast<int>(RemoveGroupOptionCode::kHelp)},
            {nullptr, 0, nullptr, 0},
    };

    bool force = false;
    uint32_t slot = 0;
    int rv;
    optind = 1;
    while ((rv = getopt_long(argc, argv, "h", options, nullptr)) != -1) {
        switch (rv) {
            case static_cast<int>(RemoveGroupOptionCode::kHelp):
                return usage(program_name, SubCommand::kRemoveGroup);
            case static_cast<int>(RemoveGroupOptionCode::kForce):
                force = true;
                break;
            case static_cast<int>(RemoveGroupOptionCode::kSlot):
                if (!android::base::ParseUint(optarg, &slot)) {
                    std::cerr << "Error: Invalid argument for --slot: " << optarg << "\n";
                    return usage(program_name, SubCommand::kRemoveGroup);
                }
                break;
            default:
                return usage(program_name, SubCommand::kRemoveGroup);
        }
    }

    if (argc - optind != 2) {
        std::cerr << "Error: Incorrect number of arguments for 'rmgrp'.\n\n";
        return usage(program_name, SubCommand::kRemoveGroup);
    }

    std::string super_device_path = argv[optind++];
    std::string group_name = argv[optind++];

    SuperModifyHelper helper(super_device_path, slot);
    if (!helper.Open()) return EX_SOFTWARE;
    if (!helper.RemoveGroup(group_name, force)) return EX_SOFTWARE;
    if (!helper.Finalize()) return EX_SOFTWARE;

    std::cout << "lpmodify rmgrp: Group " << group_name << " removed successfully.\n";
    return EX_OK;
}

int HandleChangeGroupCommand(int argc, char* argv[], const char* program_name) {
    struct option options[] = {
            {"slot", required_argument, nullptr, static_cast<int>(ChangeGroupOptionCode::kSlot)},
            {"help", no_argument, nullptr, static_cast<int>(ChangeGroupOptionCode::kHelp)},
            {nullptr, 0, nullptr, 0},
    };

    uint32_t slot = 0;
    int rv;
    optind = 1;
    while ((rv = getopt_long(argc, argv, "h", options, nullptr)) != -1) {
        switch (rv) {
            case static_cast<int>(ChangeGroupOptionCode::kHelp):
                return usage(program_name, SubCommand::kChangeGroup);
            case static_cast<int>(ChangeGroupOptionCode::kSlot):
                if (!android::base::ParseUint(optarg, &slot)) {
                    std::cerr << "Error: Invalid argument for --slot: " << optarg << "\n";
                    return usage(program_name, SubCommand::kChangeGroup);
                }
                break;
            default:
                return usage(program_name, SubCommand::kChangeGroup);
        }
    }

    if (argc - optind != 3) {
        std::cerr << "Error: Incorrect number of arguments for 'chgrp'.\n\n";
        return usage(program_name, SubCommand::kChangeGroup);
    }

    std::string super_device_path = argv[optind++];
    std::string partition_name = argv[optind++];
    std::string group_name = argv[optind++];

    SuperModifyHelper helper(super_device_path, slot);
    if (!helper.Open()) return EX_SOFTWARE;
    if (!helper.ChangePartitionGroup(partition_name, group_name)) return EX_SOFTWARE;
    if (!helper.Finalize()) return EX_SOFTWARE;

    std::cout << "lpmodify chgrp: Partition " << partition_name << " moved to group " << group_name
              << " successfully.\n";
    return EX_OK;
}

int main(int argc, char* argv[]) {
    const char* program_name = argv[0];

    if (argc < 2) {
        std::cerr << "Error: No subcommand specified.\n";
        return usage(program_name);
    }

    std::string subcommand_str = argv[1];

    if (subcommand_str == "add") {
        // Shift argv and adjust argc for subcommand handlers
        return HandleAddCommand(argc - 1, argv + 1, program_name);
    } else if (subcommand_str == "remove") {
        return HandleRemoveCommand(argc - 1, argv + 1, program_name);
    } else if (subcommand_str == "addgrp") {
        return HandleAddGroupCommand(argc - 1, argv + 1, program_name);
    } else if (subcommand_str == "rmgrp") {
        return HandleRemoveGroupCommand(argc - 1, argv + 1, program_name);
    } else if (subcommand_str == "chgrp") {
        return HandleChangeGroupCommand(argc - 1, argv + 1, program_name);
    } else if (subcommand_str == "--help" || subcommand_str == "-h") {
        return usage(program_name);
    } else {
        std::cerr << "Error: Unknown subcommand: " << subcommand_str << "\n";
        return usage(program_name);
    }
}
