// 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.

#include "service/http/action/tablets_distribution_action.h"

#include <glog/logging.h>

#include <exception>

#include <map>

#include <ostream>

#include <string>

#include <utility>

#include <vector>

#include "absl/strings/substitute.h"

#include "common/status.h"

#include "service/backend_options.h"

#include "service/http/http_channel.h"

#include "service/http/http_headers.h"

#include "service/http/http_request.h"

#include "service/http/http_status.h"

#include "storage/data_dir.h"

#include "storage/olap_common.h"

#include "storage/storage_engine.h"

#include "storage/tablet/tablet_manager.h"

namespace doris {

const static std::string HEADER_JSON = "application/json";

TabletsDistributionAction::TabletsDistributionAction(ExecEnv* exec_env, StorageEngine& engine,

                                                     TPrivilegeHier::type hier,

                                                     TPrivilegeType::type type)

        : HttpHandlerWithAuth(exec_env, hier, type), _engine(engine) {

    _host = BackendOptions::get_localhost();

}

void TabletsDistributionAction::handle(HttpRequest* req) {

    req->add_output_header(HttpHeaders::CONTENT_TYPE, HEADER_JSON.c_str());

    std::string req_group_method = req->param("group_by");

    if (req_group_method == "partition") {

        std::string req_partition_id = req->param("partition_id");

        uint64_t partition_id = 0;

        if (!req_partition_id.empty()) {

            try {

                partition_id = std::stoull(req_partition_id);

            } catch (const std::exception& e) {

                Status status = Status::InternalError("invalid argument: {}, reason:{}",

                                                      req_partition_id, e.what());

                std::string status_result = status.to_json();

                HttpChannel::send_reply(req, HttpStatus::INTERNAL_SERVER_ERROR, status_result);

                return;

            }

        }

        HttpChannel::send_reply(

                req, HttpStatus::OK,

                get_tablets_distribution_group_by_partition(partition_id).ToString());

        return;

    }

    LOG(WARNING) << "invalid argument. group_by:" << req_group_method;

    Status status = Status::InternalError(absl::Substitute("invalid argument: group_by"));

    std::string status_result = status.to_json();

    HttpChannel::send_reply(req, HttpStatus::INTERNAL_SERVER_ERROR, status_result);

}

EasyJson TabletsDistributionAction::get_tablets_distribution_group_by_partition(

        uint64_t partition_id) {

    std::map<int64_t, std::map<DataDir*, int64_t>> tablets_num_on_disk;

    std::map<int64_t, std::map<DataDir*, std::vector<TabletSize>>> tablets_info_on_disk;

    TabletManager* tablet_manager = _engine.tablet_manager();

    tablet_manager->get_tablets_distribution_on_different_disks(tablets_num_on_disk,

                                                                tablets_info_on_disk);

    EasyJson tablets_distribution_ej;

    tablets_distribution_ej["msg"] = "OK";

    tablets_distribution_ej["code"] = 0;

    EasyJson data = tablets_distribution_ej.Set("data", EasyJson::kObject);

    data["host"] = _host;

    EasyJson tablets_distribution = data.Set("tablets_distribution", EasyJson::kArray);

    int64_t tablet_total_number = 0;

    for (auto& [part_id, disk_tablets_num] : tablets_num_on_disk) {

        if (partition_id != 0 && partition_id != part_id) {

            continue;

        }

        EasyJson partition = tablets_distribution.PushBack(EasyJson::kObject);

        partition["partition_id"] = part_id;

        EasyJson disks = partition.Set("disks", EasyJson::kArray);

        for (auto& [data_dir, tablets_num] : disk_tablets_num) {

            EasyJson disk = disks.PushBack(EasyJson::kObject);

            disk["disk_path"] = data_dir->path();

            disk["tablets_num"] = tablets_num;

            tablet_total_number += tablets_num;

            if (partition_id != 0) {

                EasyJson tablets = disk.Set("tablets", EasyJson::kArray);

                for (auto& i : tablets_info_on_disk[part_id][data_dir]) {

                    EasyJson tablet = tablets.PushBack(EasyJson::kObject);

                    tablet["tablet_id"] = i.tablet_id;

                    tablet["tablet_size"] = i.tablet_size;

                }

            }

        }

    }

    tablets_distribution_ej["count"] = tablet_total_number;

    return tablets_distribution_ej;

}

} // namespace doris