// Copyright (c) Huawei Technologies Co., Ltd. 2025-2025. All rights reserved.

#pragma once

#include <memory>

#include "yacl/link/context.h"

#include "kcal/core/context.h"

namespace kcal {

class YaclLinker {
public:
    using LinkContext = yacl::link::Context;

    YaclLinker(Config config, const yacl::link::ContextDesc &desc, size_t rank, bool hasFile);
    ~YaclLinker();

    YaclLinker(const YaclLinker &) = delete;
    YaclLinker &operator=(const YaclLinker &) = delete;
    YaclLinker(YaclLinker &&) = delete;
    YaclLinker &operator=(YaclLinker &&) = delete;

    std::shared_ptr<Context> GetKcalContext() const { return kcalCtx_; }

    std::shared_ptr<LinkContext> GetYaclContext() const { return yaclCtx_; }

    static std::shared_ptr<YaclLinker> Create(Config config, const yacl::link::ContextDesc &desc, size_t rank);

private:
    // TeeNodeInfo(nodeId) ---map---> yacl rank
    int GetRankByNodeId(int nodeId) const;
    int readInputData(const char *fileName, size_t startIndex, size_t itemCount, DG_TeeInput *input);
    int readPairListData(const char *fileName, size_t startIndex, size_t itemCount, DG_PairList *input);
    int writeOutputData(const char *fileName, size_t startIndex, size_t itemCount, DG_TeeOutput *output);

    static int SendDataThunk(TeeNodeInfo *nodeInfo, unsigned char *buf, u64 len);
    static int RecvDataThunk(TeeNodeInfo *nodeInfo, unsigned char *buf, u64 *len);
    static int ReadInputThunk(const char *fileName, size_t startIndex, size_t itemCount, DG_TeeInput *input);
    static int ReadPairListThunk(const char *fileName, size_t startIndex, size_t itemCount, DG_PairList *input);
    static int WriteOutputThunk(const char *fileName, size_t startIndex, size_t itemCount, DG_TeeOutput *output);

    Config config_;
    std::shared_ptr<LinkContext> yaclCtx_;
    std::shared_ptr<Context> kcalCtx_;

    static YaclLinker *currentLinker_;
};

} // namespace kcal