/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */
package org.apache.bookkeeper.tools.cli.commands.bookieid;

import com.beust.jcommander.Parameter;
import com.google.common.util.concurrent.RateLimiter;
import java.util.ArrayList;
import java.util.List;
import lombok.Setter;
import lombok.experimental.Accessors;
import org.apache.bookkeeper.client.BookKeeperAdmin;
import org.apache.bookkeeper.client.LedgerMetadataBuilder;
import org.apache.bookkeeper.client.api.BookKeeper;
import org.apache.bookkeeper.client.api.LedgerMetadata;
import org.apache.bookkeeper.meta.LedgerManager;
import org.apache.bookkeeper.net.BookieId;
import org.apache.bookkeeper.tools.cli.helpers.ClientCommand;
import org.apache.bookkeeper.tools.framework.CliFlags;
import org.apache.bookkeeper.tools.framework.CliSpec;
import org.apache.bookkeeper.versioning.Versioned;

/**
 * Search and replace a bookie id in ledger metadata.
 */
public class SearchReplaceBookieIdCommand extends ClientCommand<SearchReplaceBookieIdCommand.Flags> {

    private static final String NAME = "searchreplace";
    private static final String DESC = "Search all ledgers for a bookie ID and replace";

    /**
     * Flags for replace bookie id.
     */
    @Accessors(fluent = true)
    @Setter
    public static class Flags extends CliFlags {

        @Parameter(names = { "-f", "--from" }, description = "Bookie ID to search for", required = true)
        private String from;
        @Parameter(names = { "-t", "--to" }, description = "Bookie ID to replace with", required = true)
        private String to;
        @Parameter(names = { "-m", "--max" }, description = "Maximum number of replacements to make")
        private long max = Long.MAX_VALUE;
        @Parameter(names = { "-r", "--rate" }, description = "Rate limit (updates per second)")
        private int rate = Integer.MAX_VALUE;
        @Parameter(names = { "--dry-run" }, description = "Don't actually write anything")
        private boolean dryRun = false;
        @Parameter(names = { "-v", "--verbose" }, description = "Verbose output")
        private boolean verbose = false;
    }

    public SearchReplaceBookieIdCommand() {
        this(new Flags());
    }

    public SearchReplaceBookieIdCommand(Flags flags) {
        super(CliSpec.<Flags>newBuilder()
            .withName(NAME)
            .withDescription(DESC)
            .withFlags(flags)
            .build());
    }

    @Override
    protected void run(BookKeeper bk, Flags flags) throws Exception {
        try (BookKeeperAdmin admin = new BookKeeperAdmin((org.apache.bookkeeper.client.BookKeeper) bk)) {
            LedgerManager ledgerManager = ((org.apache.bookkeeper.client.BookKeeper) bk).getLedgerManager();
            long i = 0;

            BookieId fromAddr = BookieId.parse(flags.from);
            BookieId toAddr = BookieId.parse(flags.to);
            System.out.println(String.format("Replacing bookie id %s with %s in metadata", fromAddr, toAddr));
            RateLimiter limiter = RateLimiter.create(flags.rate);
            for (Long lid : admin.listLedgers()) {
                Versioned<LedgerMetadata> md = ledgerManager.readLedgerMetadata(lid).get();
                if (md.getValue().getAllEnsembles().entrySet()
                        .stream().anyMatch(e -> e.getValue().contains(fromAddr))) {
                    limiter.acquire();

                    LedgerMetadataBuilder builder = LedgerMetadataBuilder.from(md.getValue());
                    md.getValue().getAllEnsembles().entrySet().stream()
                        .filter(e -> e.getValue().contains(fromAddr))
                        .forEach(e -> {
                                List<BookieId> ensemble = new ArrayList<>(e.getValue());
                                ensemble.replaceAll((a) -> {
                                        if (a.equals(fromAddr)) {
                                            return toAddr;
                                        } else {
                                            return a;
                                        }
                                    });
                                builder.replaceEnsembleEntry(e.getKey(), ensemble);
                            });
                    LedgerMetadata newMeta = builder.build();
                    if (flags.verbose) {
                        System.out.println("Replacing ledger " + lid + " metadata ...");
                        System.out.println(md.getValue().toSafeString());
                        System.out.println("with ...");
                        System.out.println(newMeta.toSafeString());
                    }
                    i++;
                    if (!flags.dryRun) {
                        ledgerManager.writeLedgerMetadata(lid, newMeta, md.getVersion()).get();
                    }
                }
                if (i >= flags.max) {
                    System.out.println("Max number of ledgers processed, exiting");
                    break;
                }
            }
            System.out.println("Replaced bookie ID in " + i + " ledgers");
        }
    }
}
