import asyncio
import json

from alibabacloud_green20220302.client import Client as AlibabaGreenClient
from alibabacloud_green20220302.models import TextModerationPlusRequest
from alibabacloud_tea_openapi.models import Config as AlibabaTeaConfig
from loguru import logger
from pydantic import BaseModel

import nonebot


class AlibabaGreenPluginConfig(BaseModel):
    module_aligreen_enable: bool = False
    module_aligreen_access_key_id: str = ""
    module_aligreen_access_key_secret: str = ""
    module_aligreen_region_id: str = "cn-shenzhen"
    module_aligreen_endpoint: str = "green-cip.cn-shenzhen.aliyuncs.com"
    module_aligreen_service: str = "llm_query_moderation"


class AlibabaGreen:
    _client: AlibabaGreenClient | None = None
    _config: AlibabaGreenPluginConfig | None = None

    @staticmethod
    def get_client() -> AlibabaGreenClient:
        assert AlibabaGreen._client is not None
        return AlibabaGreen._client

    @staticmethod
    def get_config() -> AlibabaGreenPluginConfig:
        assert AlibabaGreen._config is not None
        return AlibabaGreen._config

    @staticmethod
    def init():
        config = nonebot.get_plugin_config(AlibabaGreenPluginConfig)
        AlibabaGreen._config = config
        if not config.module_aligreen_enable:
            logger.info("该环境未启用阿里内容审查，跳过初始化")
            return
        AlibabaGreen._client = AlibabaGreenClient(AlibabaTeaConfig(
            access_key_id=config.module_aligreen_access_key_id,
            access_key_secret=config.module_aligreen_access_key_secret,
            connect_timeout=10000,
            read_timeout=3000,
            region_id=config.module_aligreen_region_id,
            endpoint=config.module_aligreen_endpoint,
        ))

    @staticmethod
    def _detect_sync(content: str) -> bool:
        if len(content) == 0:
            return True
        if not AlibabaGreen.get_config().module_aligreen_enable:
            logger.debug("该环境未启用阿里内容审查，直接跳过")
            return True

        client = AlibabaGreen.get_client()
        try:
            response = client.text_moderation_plus(TextModerationPlusRequest(
                service=AlibabaGreen.get_config().module_aligreen_service,
                service_parameters=json.dumps({
                    "content": content,
                }),
            ))
            if response.status_code == 200:
                result = response.body
                logger.info(f"检测违规内容 API 调用成功：{result}")
                risk_level: str = result.data.risk_level or "none"
                if risk_level == "high":
                    return False
                return True
            logger.error(f"检测违规内容 API 调用失败：{response}")
            return True
        except Exception as e:
            logger.error("检测违规内容 API 调用失败")
            logger.exception(e)
            return True

    @staticmethod
    async def detect(content: str) -> bool:
        return await asyncio.to_thread(AlibabaGreen._detect_sync, content)


driver = nonebot.get_driver()

@driver.on_startup
async def _():
    AlibabaGreen.init()

