/**
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
 * CANN Open Software License Agreement Version 2.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */

/*!
 * \file buffers_policy.h
 * \brief 综合管理buffer的内存和同步
 */
#ifndef BUFFERS_POLICY_H
#define BUFFERS_POLICY_H

#include "buffer_manager.h"
#define NUM_2 2
#define NUM_3 3
#define NUM_4 4
// Q复用 KV复用
// 申请单块buffer
namespace fa_base_matmul {
template <BufferType bufferType, SyncType syncType = SyncType::INNER_CORE_SYNC,
          SyncMode syncMode = SyncMode::SET_WAIT_FLAG, IdSource idSource = IdSource::INTERNAL>
class BuffersPolicySingleBuffer {
public:
    __aicore__ inline void Init(BufferManager<bufferType> &bufferManager, uint32_t size, uint32_t id = 0U)
    {
        buffer_ = bufferManager.template AllocBuffer<syncType, syncMode>(size);
        if constexpr (idSource == IdSource::INTERNAL) {
            buffer_.template Init<idSource>();
        } else if constexpr (idSource == IdSource::EXTERNAL) {
            buffer_.template Init<idSource>(id);
        }
    }

    __aicore__ inline void Uninit(BufferManager<bufferType> &bufferManager)
    {
        buffer_.template UnInit<idSource>();
        bufferManager.template FreeBuffer<syncType, syncMode>(buffer_);
    }

    __aicore__ inline Buffer<bufferType, syncType, syncMode> &Get()
    {
        return buffer_;
    }

    __aicore__ inline Buffer<bufferType, syncType, syncMode> &GetPre()
    {
        return Get();
    }

    __aicore__ inline Buffer<bufferType, syncType, syncMode> &GetReused()
    {
        return Get();
    }

private:
    Buffer<bufferType, syncType, syncMode> buffer_;
};

// 申请2个buffer，乒乓轮转
template <BufferType bufferType, SyncType syncType = SyncType::INNER_CORE_SYNC,
          SyncMode syncMode = SyncMode::SET_WAIT_FLAG, IdSource idSource = IdSource::INTERNAL>
class BuffersPolicyDB {
public:
    __aicore__ inline void Init(BufferManager<bufferType> &bufferManager, uint32_t size, uint32_t pingId = 0U,
                                uint32_t pongId = 0U)
    {
        ping_ = bufferManager.template AllocBuffer<syncType, syncMode>(size);
        pong_ = bufferManager.template AllocBuffer<syncType, syncMode>(size);

        if constexpr (idSource == IdSource::INTERNAL) {
            ping_.template Init<idSource>();
            pong_.template Init<idSource>();
        } else if constexpr (idSource == IdSource::EXTERNAL) {
            ping_.template Init<idSource>(pingId);
            pong_.template Init<idSource>(pongId);
        }
    }

    __aicore__ inline void Uninit(BufferManager<bufferType> &bufferManager)
    {
        ping_.template UnInit<idSource>();
        pong_.template UnInit<idSource>();

        bufferManager.template FreeBuffer<syncType, syncMode>(ping_);
        bufferManager.template FreeBuffer<syncType, syncMode>(pong_);
    }

    __aicore__ inline Buffer<bufferType, syncType, syncMode> &Get()
    {
        if (flag1_) { // 1
            flag1_ = 0;
            return ping_;
        } else { // 0
            flag1_ = 1;
            return pong_;
        }
    }

    // 需要与Get联用， 首次调用Get，第二次调用GetPre(Q复用)
    __aicore__ inline Buffer<bufferType, syncType, syncMode> &GetPre()
    {
        if (flag1_) { // 0->1
            return pong_;
        } else { // 1->0
            return ping_;
        }
    }

    // 需要与Get,GetPre联用， 首次调用Get，第二次调用GetPre,第三次复用时GetReused(KV复用)
    __aicore__ inline Buffer<bufferType, syncType, syncMode> &GetReused()
    {
        if (flag2_ == 0) {
            flag2_ = 1;
            return pong_;
        } else {
            flag2_ = 0;
            return ping_;
        }
    }

    __aicore__ inline Buffer<bufferType, syncType, syncMode> &GetReused(bool isNextS2IdxNoChange)
    {
        if (isNextS2IdxNoChange) {
            if (flag2_ == 0) {
                return pong_;
            } else {
                return ping_;
            }
        } else {
            return GetReused();
        }
    }

private:
    Buffer<bufferType, syncType, syncMode> ping_;
    Buffer<bufferType, syncType, syncMode> pong_;
    uint32_t flag1_ = 0;
    uint32_t flag2_ = 0;
};

// 申请3个buffer, 轮转
template <BufferType bufferType, SyncType syncType = SyncType::INNER_CORE_SYNC,
          SyncMode syncMode = SyncMode::SET_WAIT_FLAG, IdSource idSource = IdSource::INTERNAL>
class BuffersPolicy3buff {
public:
    __aicore__ inline void Init(BufferManager<bufferType> &bufferManager, uint32_t size, uint32_t aId = 0U,
                                uint32_t bId = 0U, uint32_t cId = 0U)
    {
        a_ = bufferManager.template AllocBuffer<syncType, syncMode>(size);
        b_ = bufferManager.template AllocBuffer<syncType, syncMode>(size);
        c_ = bufferManager.template AllocBuffer<syncType, syncMode>(size);

        if constexpr (idSource == IdSource::INTERNAL) {
            a_.template Init<idSource>();
            b_.template Init<idSource>();
            c_.template Init<idSource>();
        } else if constexpr (idSource == IdSource::EXTERNAL) {
            a_.template Init<idSource>(aId);
            b_.template Init<idSource>(bId);
            c_.template Init<idSource>(cId);
        }
    }

    __aicore__ inline void Uninit(BufferManager<bufferType> &bufferManager)
    {
        a_.template UnInit<idSource>();
        b_.template UnInit<idSource>();
        c_.template UnInit<idSource>();

        bufferManager.template FreeBuffer<syncType, syncMode>(a_);
        bufferManager.template FreeBuffer<syncType, syncMode>(b_);
        bufferManager.template FreeBuffer<syncType, syncMode>(c_);
    }

    __aicore__ inline Buffer<bufferType, syncType, syncMode> &Get()
    {
        if (flag1_ == 0) {
            flag1_ = 1;
            return a_;
        } else if (flag1_ == 1) {
            flag1_ = NUM_2;
            return b_;
        } else {
            flag1_ = 0;
            return c_;
        }
    }

    __aicore__ inline Buffer<bufferType, syncType, syncMode> &GetVec()
    { // mixcore architecture
        if (flag1_vec1_ == 0) {
            flag1_vec1_ = 1;
            return a_;
        } else if (flag1_vec1_ == 1) {
            flag1_vec1_ = NUM_2;
            return b_;
        } else {
            flag1_vec1_ = 0;
            return c_;
        }
    }

    __aicore__ inline Buffer<bufferType, syncType, syncMode> &GetCube()
    { // mixcore architecture
        if (flag1_bmm2_ == 0) {
            flag1_bmm2_ = 1;
            return a_;
        } else if (flag1_bmm2_ == 1) {
            flag1_bmm2_ = NUM_2;
            return b_;
        } else {
            flag1_bmm2_ = 0;
            return c_;
        }
    }

    // Q复用
    __aicore__ inline Buffer<bufferType, syncType, syncMode> &GetPre()
    {
        if (flag1_ == 0) {
            return c_;
        } else if (flag1_ == 1) {
            return a_;
        } else {
            return b_;
        }
    }

    // KV复用
    __aicore__ inline Buffer<bufferType, syncType, syncMode> &GetReused()
    {
        if (flag2_ == 0) {
            flag2_ = 1;
            return a_;
        } else if (flag2_ == 1) {
            flag2_ = NUM_2;
            return b_;
        } else {
            flag2_ = 0;
            return c_;
        }
    }

private:
    Buffer<bufferType, syncType, syncMode> a_;
    Buffer<bufferType, syncType, syncMode> b_;
    Buffer<bufferType, syncType, syncMode> c_;
    uint32_t flag1_ = 0;
    uint32_t flag1_vec1_ = 0;
    uint32_t flag1_bmm2_ = 0;
    uint32_t flag2_ = 0;
};

// 申请4个buffer + kv复用
template <BufferType bufferType, SyncType syncType = SyncType::INNER_CORE_SYNC,
          SyncMode syncMode = SyncMode::SET_WAIT_FLAG, IdSource idSource = IdSource::INTERNAL>
class BuffersPolicy4buff {
public:
    __aicore__ inline void Init(BufferManager<bufferType> &bufferManager, uint32_t size, uint32_t aId = 0U,
                                uint32_t bId = 0U, uint32_t cId = 0U, uint32_t dId = 0U)
    {
        a_ = bufferManager.template AllocBuffer<syncType, syncMode>(size);
        b_ = bufferManager.template AllocBuffer<syncType, syncMode>(size);
        c_ = bufferManager.template AllocBuffer<syncType, syncMode>(size);
        d_ = bufferManager.template AllocBuffer<syncType, syncMode>(size);

        if constexpr (idSource == IdSource::INTERNAL) {
            a_.template Init<idSource>();
            b_.template Init<idSource>();
            c_.template Init<idSource>();
            d_.template Init<idSource>();
        } else if constexpr (idSource == IdSource::EXTERNAL) {
            a_.template Init<idSource>(aId);
            b_.template Init<idSource>(bId);
            c_.template Init<idSource>(cId);
            d_.template Init<idSource>(dId);
        }
    }

    __aicore__ inline void Uninit(BufferManager<bufferType> &bufferManager)
    {
        a_.template UnInit<idSource>();
        b_.template UnInit<idSource>();
        c_.template UnInit<idSource>();
        d_.template UnInit<idSource>();

        bufferManager.template FreeBuffer<syncType, syncMode>(a_);
        bufferManager.template FreeBuffer<syncType, syncMode>(b_);
        bufferManager.template FreeBuffer<syncType, syncMode>(c_);
        bufferManager.template FreeBuffer<syncType, syncMode>(d_);
    }

    __aicore__ inline Buffer<bufferType, syncType, syncMode> &Get(uint32_t id)
    {
        uint32_t flag = id % 4;
        if (flag == 0) {
            return a_;
        } else if (flag == 1) {
            return b_;
        } else if (flag == 2) { // 2:c_
            return c_;
        } else {
            return d_;
        }
    }

    __aicore__ inline Buffer<bufferType, syncType, syncMode> &Get()
    {
        auto &buffer = Get(head_);
        head_++;
        return buffer;
    }

    __aicore__ inline Buffer<bufferType, syncType, syncMode> &GetReused()
    {
        auto &buffer = Get(used_);
        used_ = (used_ - tail_ + 1) % (head_ - tail_) + tail_;
        return buffer;
    }

    __aicore__ inline Buffer<bufferType, syncType, syncMode> &GetFree()
    {
        if (tail_ == used_) {
            used_++;
        }
        auto &buffer = Get(tail_);
        tail_++;
        return buffer;
    }

private:
    Buffer<bufferType, syncType, syncMode> a_;
    Buffer<bufferType, syncType, syncMode> b_;
    Buffer<bufferType, syncType, syncMode> c_;
    Buffer<bufferType, syncType, syncMode> d_;
    uint32_t tail_ = 0; // 表示当前正在使用的buffer队列队尾
    uint32_t head_ = 0; // 表示当前正在使用的buffer队列队首+1
    uint32_t used_ = 0; // 表示当前正在使用的buffer，于首尾间，左闭右开
};

template <BufferType bufferType, SyncType syncType = SyncType::INNER_CORE_SYNC>
class Matrix2x2BufferPolicy { // 4buffer
    // 二维buffer管理，地址行优先，使用列优先
    // MracBuffer:memory address with row first, alloc/use/free with column first
public:
    __aicore__ inline void Init(BufferManager<bufferType> &bufferManager, uint32_t size)
    {
        bufferM0k0_ = bufferManager.template AllocBuffer<syncType>(size);
        bufferM0k1_ = bufferManager.template AllocBuffer<syncType>(size);
        bufferM1k0_ = bufferManager.template AllocBuffer<syncType>(size);
        bufferM1k1_ = bufferManager.template AllocBuffer<syncType>(size);

        bufferM0k0_.template Init<IdSource::INTERNAL>();
        bufferM0k1_.template Init<IdSource::INTERNAL>();
        bufferM1k0_.template Init<IdSource::INTERNAL>();
        bufferM1k1_.template Init<IdSource::INTERNAL>();
    }

    __aicore__ inline void Uninit(BufferManager<bufferType> &bufferManager)
    {
        bufferM0k0_.template UnInit<IdSource::INTERNAL>();
        bufferM0k1_.template UnInit<IdSource::INTERNAL>();
        bufferM1k0_.template UnInit<IdSource::INTERNAL>();
        bufferM1k1_.template UnInit<IdSource::INTERNAL>();

        bufferManager.FreeBuffer(bufferM0k0_);
        bufferManager.FreeBuffer(bufferM0k1_);
        bufferManager.FreeBuffer(bufferM1k0_);
        bufferManager.FreeBuffer(bufferM1k1_);
    }

    __aicore__ inline void SetMExtent(int32_t mExtent)
    {
        aIdx_ = -1;
        amIdx_ = (amIdx_ + mSize_ - 1) % mSize_; // 翻转 0->1, 1->0
        akIdx_ = 0;

        uIdx_ = -1;
        umIdx_ = (umIdx_ + mSize_ - 1) % mSize_;
        ukIdx_ = 0;

        fIdx_ = -1;
        fmIdx_ = (fmIdx_ + mSize_ - 1) % mSize_;
        fkIdx_ = 0;

        mExtent_ = mExtent;
    }

    __aicore__ inline Buffer<bufferType, syncType> &AllocNext()
    {
        aIdx_++;
        return GetBuffer(aIdx_, amIdx_, akIdx_);
    }

    __aicore__ inline Buffer<bufferType, syncType> &ReuseNext()
    {
        uIdx_++;
        return GetBuffer(uIdx_, umIdx_, ukIdx_);
    }

    __aicore__ inline Buffer<bufferType, syncType> &FreeNext()
    {
        fIdx_++;
        return GetBuffer(fIdx_, fmIdx_, fkIdx_);
    }

    __aicore__ inline Buffer<bufferType, syncType> &PeekNextK()
    {                                            // 在Alloc阶段使用，k方向取下一个
        return PeekBuffer(amIdx_, (1 - akIdx_)); // k翻转
    }

private:
    __aicore__ inline Buffer<bufferType, syncType> &GetBuffer(int32_t xIdx, int32_t &mIdx, int32_t &kIdx)
    {
        // xIdx为入参，表示当前alloc/use/free的idx，mIdx和kIdx为下标出参，移动到下一个buffer并获取
        mIdx = (mIdx + mExtent_ - 1) % mExtent_;
        kIdx = (xIdx / mExtent_) % kSize_;
        if (mIdx == 0 && kIdx == 0) {
            return bufferM0k0_;
        } else if (mIdx == 0 && kIdx == 1) {
            return bufferM0k1_;
        } else if (mIdx == 1 && kIdx == 0) {
            return bufferM1k0_;
        } else { // 该分支条件为：mIdx == 1 && kIdx == 1
            return bufferM1k1_;
        }
    }

    __aicore__ inline Buffer<bufferType, syncType> &PeekBuffer(int32_t mIdx, int32_t kIdx)
    {
        // 只访问buffer，不进行下标移动
        if (mIdx == 0 && kIdx == 0) {
            return bufferM0k0_;
        } else if (mIdx == 0 && kIdx == 1) {
            return bufferM0k1_;
        } else if ((mIdx == 1) && (kIdx == 0)) {
            return bufferM1k0_;
        } else { // mIdx == 1 && kIdx == 1
            return bufferM1k1_;
        }
    }

    Buffer<bufferType, syncType> bufferM0k0_;
    Buffer<bufferType, syncType> bufferM0k1_;
    Buffer<bufferType, syncType> bufferM1k0_;
    Buffer<bufferType, syncType> bufferM1k1_;
    int32_t mSize_ = 2; // m的总buffer数
    int32_t kSize_ = 2; // k的总buffer数

    // Alloc
    int32_t aIdx_ = -1; // 当前第几次Alloc Buffer
    int32_t amIdx_ = 0; // 当前Alloc Buffer的m下标
    int32_t akIdx_ = 0; // 当前Alloc Buffer的k下标

    // Reuse
    int32_t uIdx_ = -1;
    int32_t umIdx_ = 0;
    int32_t ukIdx_ = 0;

    // Free
    int32_t fIdx_ = -1;
    int32_t fmIdx_ = 0;
    int32_t fkIdx_ = 0;

    int32_t mExtent_ = 0; // m实际使用的大小，可以为1或者2
};
} // namespace fa_base_matmul
#endif