/*
 * 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_H
#define VISION_H

#include <kylin-ai/aiengine.h>
#include <kylin-ai/common/error.h>

#include <map>

namespace genai {

typedef enum {
    VISION_INPUT_TEXT_LENGTH_INVALID = 100,  // 文生图输入文本过长
    VISION_IMAGE_STYLE_INVALID,              // 不支持的风格
    VISION_IMAGE_SIZE_INVALID,               // 不支持的大小
    VISION_IMAGE_NUMBER_INVALID,             // 不支持的图片数量
    VISION_IMAGE_GENERATION_BLOCKED,         // 图片未过审
    VISION_IMAGE_GENERATION_FAILED,          // 生成图片失败
} VisionErrorCode;

typedef enum { IMAGE_FORMAT_UNKNOWN = -1, IMAGE_FORMAT_JPEG, IMAGE_FORMAT_PNG, IMAGE_FORMAT_WEBP } ImageFormat;

struct ImageConfig {
    int width;
    int height;
    int size;
    int number;
    std::string style;
};

inline std::map<int, int> visionEngineErrorMap{
    {(int)ai_engine::VisionEngineErrorCode::NomoreMemory, AiSdkCommonErrorCode::AISDK_SERVICE_ERROR},
    {(int)ai_engine::VisionEngineErrorCode::FileNotFound, AiSdkCommonErrorCode::AISDK_SERVICE_ERROR},
    {(int)ai_engine::VisionEngineErrorCode::VersionMismatched, AiSdkCommonErrorCode::AISDK_SERVICE_ERROR},

    {(int)ai_engine::VisionEngineErrorCode::FilePermissionDenied, AiSdkCommonErrorCode::AISDK_AUTHENTICATION_FAILURE},
    {(int)ai_engine::VisionEngineErrorCode::NetworkDisconnected, AiSdkCommonErrorCode::AISDK_NET_ERROR},
    {(int)ai_engine::VisionEngineErrorCode::NetworkLatencyTooLong, AiSdkCommonErrorCode::AISDK_NET_ERROR},

    {(int)ai_engine::VisionEngineErrorCode::FailedToConnectServer, AiSdkCommonErrorCode::AISDK_SERVICE_ERROR},
    {(int)ai_engine::VisionEngineErrorCode::ServerNoResponse, AiSdkCommonErrorCode::AISDK_SERVICE_ERROR},
    {(int)ai_engine::VisionEngineErrorCode::NoServerService, AiSdkCommonErrorCode::AISDK_SERVICE_ERROR},

    {(int)ai_engine::VisionEngineErrorCode::Unauthorized, AiSdkCommonErrorCode::AISDK_AUTHENTICATION_FAILURE},
    {(int)ai_engine::VisionEngineErrorCode::TooLittleParams, VISION_INPUT_TEXT_LENGTH_INVALID},
    {(int)ai_engine::VisionEngineErrorCode::TooManyParams, VISION_INPUT_TEXT_LENGTH_INVALID},

    /*no VISION_PARAM_ERROR*/
    {(int)ai_engine::VisionEngineErrorCode::InvalidParamType, VISION_INPUT_TEXT_LENGTH_INVALID},
    {(int)ai_engine::VisionEngineErrorCode::TooLittleData, VISION_INPUT_TEXT_LENGTH_INVALID},
    {(int)ai_engine::VisionEngineErrorCode::TooManyData, VISION_INPUT_TEXT_LENGTH_INVALID},

    {(int)ai_engine::VisionEngineErrorCode::TooFrequentRequest, AiSdkCommonErrorCode::AISDK_TOO_MANY_REQUESTS},
    {(int)ai_engine::VisionEngineErrorCode::TooManyRequest, AiSdkCommonErrorCode::AISDK_TOO_MANY_REQUESTS},

    /*no VISION_PARAM_ERROR*/
    {(int)ai_engine::VisionEngineErrorCode::Invalidwords, VISION_INPUT_TEXT_LENGTH_INVALID},
    {(int)ai_engine::VisionEngineErrorCode::ImageGenerationFailed, VISION_IMAGE_GENERATION_FAILED},
    {(int)ai_engine::VisionEngineErrorCode::ImageCensorshipFailed, VISION_IMAGE_GENERATION_BLOCKED},
    {(int)ai_engine::VisionEngineErrorCode::ImageNumberInvalid, VISION_IMAGE_NUMBER_INVALID},
    {(int)ai_engine::VisionEngineErrorCode::ImageSizeInvalid, VISION_IMAGE_SIZE_INVALID},
    {(int)ai_engine::VisionEngineErrorCode::SslCertificateError, AiSdkCommonErrorCode::AISDK_SSL_CERTIFICATE_ERROR},
};

static int getVisionErrorCode(int engineErrorCode) {
    if (genai::visionEngineErrorMap.find(engineErrorCode) != genai::visionEngineErrorMap.end()) {
        return genai::visionEngineErrorMap.at(engineErrorCode);
    }
    return AiSdkCommonErrorCode::AISDK_NO_ERROR;
}

static int getVisionImageFormat(const std::string &imageFormat) {
    std::map<std::string, ImageFormat> visionEngineImageFormatMap{
        {"jpg", ImageFormat::IMAGE_FORMAT_JPEG},  {"jpeg", ImageFormat::IMAGE_FORMAT_JPEG},
        {"JPG", ImageFormat::IMAGE_FORMAT_JPEG},  {"JPEG", ImageFormat::IMAGE_FORMAT_JPEG},
        {"png", ImageFormat::IMAGE_FORMAT_PNG},   {"PNG", ImageFormat::IMAGE_FORMAT_PNG},
        {"webp", ImageFormat::IMAGE_FORMAT_WEBP}, {"WEBP", ImageFormat::IMAGE_FORMAT_WEBP},
    };

    if (visionEngineImageFormatMap.find(imageFormat) != visionEngineImageFormatMap.end()) {
        return visionEngineImageFormatMap.at(imageFormat);
    }
    return genai::ImageFormat::IMAGE_FORMAT_UNKNOWN;
}
}  // namespace genai

#endif
