#ifndef SERVICES_ASSISTANT_MODEL_INFERENCE_H
#define SERVICES_ASSISTANT_MODEL_INFERENCE_H

#include "assistant/sessionbase.h"
#include "engineselector/engineselector.h"
#include "modelinferenceinterface.h"

class ModelInference : public ModelInferenceInterface {
public:
    ModelInference() = default;

    ~ModelInference() override = default;

    std::future<bool> callChatModelInference(const std::string &modelInput) override;

    DialogueData::CallSuccessIdentifier callIntentionModelInference(const std::string &modelInput) override;

    std::string intentModelName() override;

    void initChatEngine(const std::shared_ptr<assistant::TextGenerationEngine> &nlpEngine,
                        EngineSelector &engineSelector_);

    void initIntentEngine(const std::shared_ptr<assistant::TextGenerationEngine> &nlpEngine);

    int getCurrentModelMaxTokens() override;
    bool isCloudModel() override;

private:
    void handleChatResult(const ai_engine::gen_ai::nlp::TextGenerationResult &result);

    bool handleNetworkDisconnection(
        const std::string &modelInput,
        const std::function<void(const ai_engine::gen_ai::nlp::TextGenerationResult &result)> &callback);

private:
    std::shared_ptr<assistant::TextGenerationEngine> chatEngine_;
    std::shared_ptr<assistant::TextGenerationEngine> intentEngine_;
    std::string intentReplyCache_;
    std::string intentResult_;
    std::string intentModelName_;
    EngineSelector *engineSelector_{nullptr};
    std::atomic_bool isInited_{false};
};

#endif  // SERVICES_ASSISTANT_MODEL_INFERENCE_H
