/*
 * Copyright 2024 KylinSoft Co., Ltd.
 *
 * This program is free software: you can redistribute it and/or modify it under
 * the terms of the GNU General Public License as published by the Free Software
 * Foundation, either version 3 of the License, or (at your option) any later
 * version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
 * FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
 * details.
 *
 * You should have received a copy of the GNU General Public License along with
 * this program. If not, see <https://www.gnu.org/licenses/>.
 */

#ifndef VISION_IMAGE_GENERATOR_H
#define VISION_IMAGE_GENERATOR_H

#include <kylin-ai/aiengine.h>

#include <atomic>
#include <memory>
#include <string>

#include "concurrentqueue/blockingconcurrentqueue.h"
#include "taskmanager/taskrunner.h"
#include "vision.h"

struct ImageGenerationResult {
    int format{genai::ImageFormat::IMAGE_FORMAT_UNKNOWN};
    int width{0};
    int height{0};
    int total{0};
    int index{-1};
    std::string imageBase64Data;
    int errorCode{0};
    std::string errorMessage;
};

class ImageGenerator : public std::enable_shared_from_this<ImageGenerator> {
private:
    using ImageGenerationEngine = ai_engine::gen_ai::vision::AbstractImageGenerationEngine;

public:
    using ImageGeneratorResultCallback = std::function<void(const ImageGenerationResult &)>;
    using ImageGenerationFinishedCallback = std::function<void()>;

    ImageGenerator(std::unique_ptr<ImageGenerationEngine> visionEngine);

    ~ImageGenerator();

    void setImageGeneratorCallback(const ImageGeneratorResultCallback &callback);
    void setVisionEngine(std::unique_ptr<ImageGenerationEngine> visionEngine);

    void generateImage(const std::string &prompt, ImageGeneratorResultCallback resultCallback,
                       ImageGenerationFinishedCallback finisedCallback);
    void stopImageGeneration();

    int getSupportedImageNumber();
    std::string getSupportedImageSize();
    std::string getSupportedImageStyle();

    void setImageConfig(const genai::ImageConfig &config);

    bool engineInUse() { return engineInUse_; }

private:
    std::shared_ptr<ImageGenerator> getSharedObject() { return shared_from_this(); }

    void initVisionEngineResultCallback();
    void onVisionEngineResultCallback(const ai_engine::gen_ai::vision::ImageGenerationResult &imageData);
    void setVisionEngineConfig(const genai::ImageConfig &config);
    void onErrorOccurred(const ai_engine::EngineError &error);

    int getImageNumberFromJsonString(const std::string &params);

    void generateImageAsync();

    void generateImageTask();

private:
    struct ImageGenerationRequest {
        std::string prompt;
        ImageGeneratorResultCallback resultCallback;
        ImageGenerationFinishedCallback finishedCallback;
    };

private:
    genai::ImageConfig imageConfig_{};
    std::unique_ptr<ImageGenerationEngine> visionEngine_{nullptr};
    ImageGeneratorResultCallback imageGeneratorResultCallback_{nullptr};
    std::atomic_bool engineInUse_{false};
    std::atomic<task::TaskId> taskId_{task::INVALID_TASK_ID};
    moodycamel::BlockingConcurrentQueue<std::string> promptQueue_{};
    moodycamel::BlockingConcurrentQueue<ImageGenerationRequest> requestQueue_{};
    ImageGenerationRequest currentRequest_;
};

#endif
