// 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 "recycler/util.h"

#include <glog/logging.h>

#include <cstdint>

#include "common/util.h"

#include "meta-service/keys.h"

#include "meta-service/meta_service_schema.h"

#include "meta-service/txn_kv.h"

#include "meta-service/txn_kv_error.h"

namespace doris::cloud {

namespace config {

extern int32_t recycle_job_lease_expired_ms;

} // namespace config

int get_all_instances(TxnKv* txn_kv, std::vector<InstanceInfoPB>& res) {

    InstanceKeyInfo key0_info {""};

    InstanceKeyInfo key1_info {"\xff"}; // instance id are human readable strings

    std::string key0;

    std::string key1;

    instance_key(key0_info, &key0);

    instance_key(key1_info, &key1);

    std::unique_ptr<Transaction> txn;

    TxnErrorCode err = txn_kv->create_txn(&txn);

    if (err != TxnErrorCode::TXN_OK) {

        LOG(INFO) << "failed to init txn, err=" << err;

        return -1;

    }

    std::unique_ptr<RangeGetIterator> it;

    do {

        TxnErrorCode err = txn->get(key0, key1, &it);

        if (err != TxnErrorCode::TXN_OK) {

            LOG(WARNING) << "failed to get instance, err=" << err;

            return -1;

        }

        while (it->has_next()) {

            auto [k, v] = it->next();

            if (!it->has_next()) key0 = k;

            InstanceInfoPB instance_info;

            if (!instance_info.ParseFromArray(v.data(), v.size())) {

                LOG(WARNING) << "malformed instance info, key=" << hex(k);

                return -1;

            }

            res.push_back(std::move(instance_info));

        }

        key0.push_back('\x00'); // Update to next smallest key for iteration

    } while (it->more());

    return 0;

}

int prepare_instance_recycle_job(TxnKv* txn_kv, std::string_view key,

                                 const std::string& instance_id, const std::string& ip_port,

                                 int64_t interval_ms) {

    std::string val;

    std::unique_ptr<Transaction> txn;

    TxnErrorCode err = txn_kv->create_txn(&txn);

    if (err != TxnErrorCode::TXN_OK) {

        LOG(WARNING) << "failed to create txn";

        return -1;

    }

    err = txn->get(key, &val);

    if (err != TxnErrorCode::TXN_OK && err != TxnErrorCode::TXN_KEY_NOT_FOUND) {

        LOG(WARNING) << "failed to get kv, err=" << err << " key=" << hex(key);

        return -1;

    }

    using namespace std::chrono;

    auto now = duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();

    JobRecyclePB job_info;

    auto is_expired = [&]() {

        if (!job_info.ParseFromString(val)) {

            LOG(WARNING) << "failed to parse JobRecyclePB, key=" << hex(key);

            // if failed to parse, just recycle it.

            return true;

        }

        DCHECK(job_info.instance_id() == instance_id);

        if (job_info.status() == JobRecyclePB::BUSY) {

            if (now < job_info.expiration_time_ms()) {

                LOG(INFO) << "job is busy. host=" << job_info.ip_port()

                          << " expiration=" << job_info.expiration_time_ms()

                          << " instance=" << instance_id;

                return false;

            }

        }

        bool finish_expired = now - job_info.last_ctime_ms() > interval_ms;

        if (!finish_expired) {

            LOG(INFO) << "the time since last finished job is too short. host="

                      << job_info.ip_port() << " ctime=" << job_info.last_ctime_ms()

                      << " instance=" << instance_id;

        }

        return finish_expired;

    };

    if (err == TxnErrorCode::TXN_KEY_NOT_FOUND || is_expired()) {

        job_info.set_status(JobRecyclePB::BUSY);

        job_info.set_instance_id(instance_id);

        job_info.set_ip_port(ip_port);

        job_info.set_expiration_time_ms(now + config::recycle_job_lease_expired_ms);

        val = job_info.SerializeAsString();

        txn->put(key, val);

        err = txn->commit();

        if (err != TxnErrorCode::TXN_OK) {

            LOG(WARNING) << "failed to commit, err=" << err << " key=" << hex(key);

            return -1;

        }

        return 0;

    }

    return 1;

}

void finish_instance_recycle_job(TxnKv* txn_kv, std::string_view key,

                                 const std::string& instance_id, const std::string& ip_port,

                                 bool success, int64_t ctime_ms) {

    std::string val;

    int retry_times = 0;

    do {

        std::unique_ptr<Transaction> txn;

        TxnErrorCode err = txn_kv->create_txn(&txn);

        if (err != TxnErrorCode::TXN_OK) {

            LOG(WARNING) << "failed to create txn";

            return;

        }

        err = txn->get(key, &val);

        if (err != TxnErrorCode::TXN_OK) {

            LOG(WARNING) << "failed to get kv, err=" << err << " key=" << hex(key);

            return;

        }

        using namespace std::chrono;

        auto now = duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();

        JobRecyclePB job_info;

        if (!job_info.ParseFromString(val)) {

            LOG(WARNING) << "failed to parse JobRecyclePB, key=" << hex(key);

            return;

        }

        DCHECK(job_info.instance_id() == instance_id);

        if (job_info.ip_port() != ip_port) {

            LOG(WARNING) << "job is doing at other machine: " << job_info.ip_port()

                         << " key=" << hex(key);

            return;

        }

        if (job_info.status() != JobRecyclePB::BUSY) {

            LOG(WARNING) << "job is not busy, key=" << hex(key);

            return;

        }

        job_info.set_status(JobRecyclePB::IDLE);

        job_info.set_instance_id(instance_id);

        job_info.set_last_finish_time_ms(now);

        job_info.set_last_ctime_ms(ctime_ms);

        if (success) {

            job_info.set_last_success_time_ms(now);

        }

        val = job_info.SerializeAsString();

        txn->put(key, val);

        err = txn->commit();

        if (err == TxnErrorCode::TXN_OK) {

            LOG(INFO) << "succ to commit to finish recycle job, key=" << hex(key);

            return;

        }

        // maybe conflict with the commit of the leased thread

        LOG(WARNING) << "failed to commit to finish recycle job, err=" << err << " key=" << hex(key)

                     << " retry_times=" << retry_times;

    } while (retry_times++ < 3);

    LOG(WARNING) << "finally failed to commit to finish recycle job, key=" << hex(key);

}

int lease_instance_recycle_job(TxnKv* txn_kv, std::string_view key, const std::string& instance_id,

                               const std::string& ip_port) {

    std::string val;

    std::unique_ptr<Transaction> txn;

    TxnErrorCode err = txn_kv->create_txn(&txn);

    if (err != TxnErrorCode::TXN_OK) {

        LOG(WARNING) << "failed to create txn";

        return -1;

    }

    err = txn->get(key, &val);

    if (err != TxnErrorCode::TXN_OK) {

        LOG(WARNING) << "failed to get kv, err=" << err << " key=" << hex(key);

        return -1;

    }

    using namespace std::chrono;

    auto now = duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();

    JobRecyclePB job_info;

    if (!job_info.ParseFromString(val)) {

        LOG(WARNING) << "failed to parse JobRecyclePB, key=" << hex(key);

        return 1;

    }

    DCHECK(job_info.instance_id() == instance_id);

    if (job_info.ip_port() != ip_port) {

        LOG(WARNING) << "job is doing at other machine: " << job_info.ip_port()

                     << " key=" << hex(key);

        return 1;

    }

    if (job_info.status() != JobRecyclePB::BUSY) {

        LOG(WARNING) << "job is not busy, key=" << hex(key);

        return 1;

    }

    job_info.set_expiration_time_ms(now + config::recycle_job_lease_expired_ms);

    val = job_info.SerializeAsString();

    txn->put(key, val);

    err = txn->commit();

    if (err != TxnErrorCode::TXN_OK) {

        LOG(WARNING) << "failed to commit, failed to lease recycle job, err=" << err

                     << " key=" << hex(key);

        return -1;

    }

    return 0;

}

int get_tablet_idx(TxnKv* txn_kv, const std::string& instance_id, int64_t tablet_id,

                   TabletIndexPB& tablet_idx) {

    std::unique_ptr<Transaction> txn;

    TxnErrorCode err = txn_kv->create_txn(&txn);

    if (err != TxnErrorCode::TXN_OK) {

        LOG(WARNING) << "failed to create txn";

        return -1;

    }

    std::string key, val;

    meta_tablet_idx_key({instance_id, tablet_id}, &key);

    err = txn->get(key, &val);

    if (err != TxnErrorCode::TXN_OK) {

        LOG(WARNING) << fmt::format("failed to get tablet_idx, err={} tablet_id={} key={}", err,

                                    tablet_id, hex(key));

        return -1;

    }

    if (!tablet_idx.ParseFromString(val)) [[unlikely]] {

        LOG(WARNING) << fmt::format("malformed tablet index value, tablet_id={} key={}", tablet_id,

                                    hex(key));

        return -1;

    }

    if (tablet_id != tablet_idx.tablet_id()) [[unlikely]] {

        LOG(WARNING) << "unexpected error given_tablet_id=" << tablet_id

                     << " idx_pb_tablet_id=" << tablet_idx.tablet_id() << " key=" << hex(key);

        return -1;

    }

    return 0;

}

int get_tablet_meta(TxnKv* txn_kv, const std::string& instance_id, int64_t tablet_id,

                    TabletMetaCloudPB& tablet_meta) {

    TabletIndexPB tablet_idx;

    int ret = get_tablet_idx(txn_kv, instance_id, tablet_id, tablet_idx);

    if (ret < 0) {

        return ret;

    }

    std::unique_ptr<Transaction> txn;

    TxnErrorCode err = txn_kv->create_txn(&txn);

    if (err != TxnErrorCode::TXN_OK) {

        LOG(WARNING) << "failed to create txn";

        return -1;

    }

    std::string key, val;

    meta_tablet_key({instance_id, tablet_idx.table_id(), tablet_idx.index_id(),

                     tablet_idx.partition_id(), tablet_id},

                    &key);

    err = txn->get(key, &val);

    if (err != TxnErrorCode::TXN_OK) {

        LOG(WARNING) << fmt::format(

                "failed to get tablet, err={}, table_id={}, index_id={}, partition_id={}, "

                "tablet_id={} key={}",

                err, tablet_idx.table_id(), tablet_idx.index_id(), tablet_idx.partition_id(),

                tablet_id, hex(key));

        return -1;

    }

    if (!tablet_meta.ParseFromString(val)) [[unlikely]] {

        LOG(WARNING) << fmt::format("malformed tablet meta, tablet_id={} key={}", tablet_id,

                                    hex(key));

        return -1;

    }

    return 0;

}

} // namespace doris::cloud