zhenxun_bot/plugins/statistics/statistics_handle.py

from nonebot import on_command
from nonebot.adapters.onebot.v11 import Bot, GroupMessageEvent, MessageEvent, Message
from models.group_info import GroupInfo
from configs.path_config import DATA_PATH, IMAGE_PATH
from nonebot.params import CommandArg, Command
from utils.image_utils import BuildMat
from utils.message_builder import image
from utils.manager import plugins2settings_manager
from typing import Tuple
import asyncio
import os

try:
    import ujson as json
except ModuleNotFoundError:
    import json


__zx_plugin_name__ = "<EFBFBD><EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD>ƿ<EFBFBD><EFBFBD>ӻ<EFBFBD>"
__plugin_usage__ = """
usage<EFBFBD><EFBFBD>
    <EFBFBD><EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD>ƿ<EFBFBD><EFBFBD>ӻ<EFBFBD>
    ָ<EFBFBD>
        <EFBFBD><EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD>
        <EFBFBD>չ<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD>
        <EFBFBD>ܹ<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD> ?[<EFBFBD><EFBFBD><EFBFBD><EFBFBD>]
        <EFBFBD>¹<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD> ?[<EFBFBD><EFBFBD><EFBFBD><EFBFBD>]
        <EFBFBD>ҵĹ<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD>
        <EFBFBD>ҵ<EFBFBD><EFBFBD>չ<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD> ?[<EFBFBD><EFBFBD><EFBFBD><EFBFBD>]
        <EFBFBD>ҵ<EFBFBD><EFBFBD>ܹ<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD> ?[<EFBFBD><EFBFBD><EFBFBD><EFBFBD>]
        <EFBFBD>ҵ<EFBFBD><EFBFBD>¹<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD> ?[<EFBFBD><EFBFBD><EFBFBD><EFBFBD>]
""".strip()
__plugin_superuser_usage__ = """
usage<EFBFBD><EFBFBD>
    <EFBFBD><EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD>ƿ<EFBFBD><EFBFBD>ӻ<EFBFBD>
    ָ<EFBFBD>
        ȫ<EFBFBD>ֹ<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD>
        ȫ<EFBFBD><EFBFBD><EFBFBD>չ<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD>
        ȫ<EFBFBD><EFBFBD><EFBFBD>ܹ<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD> ?[<EFBFBD><EFBFBD><EFBFBD><EFBFBD>]
        ȫ<EFBFBD><EFBFBD><EFBFBD>¹<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD> ?[<EFBFBD><EFBFBD><EFBFBD><EFBFBD>]
""".strip()
__plugin_des__ = "<EFBFBD><EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD>ƿ<EFBFBD><EFBFBD>ӻ<EFBFBD>"
__plugin_cmd__ = [
    "<EFBFBD><EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD>",
    "ȫ<EFBFBD>ֹ<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD> [_superuser]",
    "ȫ<EFBFBD><EFBFBD><EFBFBD>չ<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD> [_superuser]",
    "ȫ<EFBFBD><EFBFBD><EFBFBD>ܹ<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD> ?[<5B><><EFBFBD><EFBFBD>] [_superuser]",
    "ȫ<EFBFBD><EFBFBD><EFBFBD>¹<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD> ?[<5B><><EFBFBD><EFBFBD>] [_superuser]",
    "<EFBFBD>ܹ<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD> ?[<5B><><EFBFBD><EFBFBD>]",
    "<EFBFBD>¹<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD> ?[<5B><><EFBFBD><EFBFBD>]",
    "<EFBFBD>ҵĹ<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD>",
    "<EFBFBD>ҵ<EFBFBD><EFBFBD>չ<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD> ?[<5B><><EFBFBD><EFBFBD>]",
    "<EFBFBD>ҵ<EFBFBD><EFBFBD>ܹ<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD> ?[<5B><><EFBFBD><EFBFBD>]",
    "<EFBFBD>ҵ<EFBFBD><EFBFBD>¹<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD> ?[<5B><><EFBFBD><EFBFBD>]",
]
__plugin_type__ = ("<EFBFBD><EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD>ƿ<EFBFBD><EFBFBD>ӻ<EFBFBD>", 1)
__plugin_version__ = 0.1
__plugin_author__ = "HibiKier"
__plugin_settings__ = {
    "level": 5,
    "default_status": True,
    "limit_superuser": False,
    "cmd": ["<EFBFBD><EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD>"],
}


statistics = on_command(
    "<EFBFBD><EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD>",
    aliases={
        "ȫ<EFBFBD>ֹ<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD>",
        "ȫ<EFBFBD><EFBFBD><EFBFBD>չ<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD>",
        "ȫ<EFBFBD><EFBFBD><EFBFBD>ܹ<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD>",
        "ȫ<EFBFBD><EFBFBD><EFBFBD>¹<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD>",
        "<EFBFBD>չ<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD>",
        "<EFBFBD>ܹ<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD>",
        "<EFBFBD>¹<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD>",
        "<EFBFBD>ҵĹ<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD>",
        "<EFBFBD>ҵ<EFBFBD><EFBFBD>չ<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD>",
        "<EFBFBD>ҵ<EFBFBD><EFBFBD>ܹ<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD>",
        "<EFBFBD>ҵ<EFBFBD><EFBFBD>¹<EFBFBD><EFBFBD>ܵ<EFBFBD><EFBFBD><EFBFBD>ͳ<EFBFBD><EFBFBD>",
    },
    priority=5,
    block=True,
)


statistics_group_file = DATA_PATH / "statistics" / "_prefix_count.json"
statistics_user_file = DATA_PATH / "statistics" / "_prefix_user_count.json"


@statistics.handle()
async def _(bot: Bot, event: MessageEvent, cmd: Tuple[str, ...] = Command(), arg: Message = CommandArg()):
    msg = arg.extract_plain_text().strip()
    if cmd[0][:2] == "ȫ<EFBFBD><EFBFBD>":
        if str(event.user_id) in bot.config.superusers:
            data: dict = json.load(open(statistics_group_file, "r", encoding="utf8"))
            if cmd[0][2] == '<EFBFBD><EFBFBD>':
                _type = 'day_statistics'
            elif cmd[0][2] == '<EFBFBD><EFBFBD>':
                _type = 'week_statistics'
            elif cmd[0][2] == '<EFBFBD><EFBFBD>':
                _type = 'month_statistics'
            else:
                _type = 'total_statistics'
            tmp_dict = {}
            data = data[_type]
            if _type in ["day_statistics", "total_statistics"]:
                for key in data['total']:
                    tmp_dict[key] = data['total'][key]
            else:
                for group in data.keys():
                    if group != 'total':
                        for day in data[group].keys():
                            for plugin_name in data[group][day].keys():
                                if data[group][day][plugin_name] is not None:
                                    if tmp_dict.get(plugin_name) is None:
                                        tmp_dict[plugin_name] = 1
                                    else:
                                        tmp_dict[plugin_name] += data[group][day][plugin_name]
            bar_graph = await init_bar_graph(tmp_dict, cmd[0])
            await asyncio.get_event_loop().run_in_executor(None, bar_graph.gen_graph)
            await statistics.finish(image(b64=bar_graph.pic2bs4()))
        return
    if cmd[0][:2] == "<EFBFBD>ҵ<EFBFBD>":
        _type = "user"
        key = str(event.user_id)
        cmd = list(cmd)
        cmd[0] = cmd[0][2:]
        if not statistics_user_file.exists():
            await statistics.finish("ͳ<EFBFBD><EFBFBD><EFBFBD>ļ<EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>...", at_sender=True)
    else:
        if not isinstance(event, GroupMessageEvent):
            await statistics.finish("<EFBFBD><EFBFBD><EFBFBD><EFBFBD>Ⱥ<EFBFBD>ڵ<EFBFBD><EFBFBD>ô˹<EFBFBD><EFBFBD><EFBFBD>...")
        _type = "group"
        key = str(event.group_id)
        if not statistics_group_file.exists():
            await statistics.finish("ͳ<EFBFBD><EFBFBD><EFBFBD>ļ<EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>...", at_sender=True)
    plugin = ""
    if cmd[0][0] == "<EFBFBD><EFBFBD>":
        arg = "day_statistics"
    elif cmd[0][0] == "<EFBFBD><EFBFBD>":
        arg = "week_statistics"
    elif cmd[0][0] == "<EFBFBD><EFBFBD>":
        arg = "month_statistics"
    else:
        arg = "total_statistics"
    if msg:
        plugin = plugins2settings_manager.get_plugin_module(msg)
        if not plugin:
            if arg not in ["day_statistics", "total_statistics"]:
                await statistics.finish("δ<EFBFBD>ҵ<EFBFBD><EFBFBD>˹<EFBFBD><EFBFBD>ܵĵ<EFBFBD><EFBFBD><EFBFBD>...", at_sender=True)
    if _type == "group":
        data: dict = json.load(open(statistics_group_file, "r", encoding="utf8"))
        if not data[arg].get(str(event.group_id)):
            await statistics.finish("<EFBFBD><EFBFBD>Ⱥͳ<EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ݲ<EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>...", at_sender=True)
    else:
        data: dict = json.load(open(statistics_user_file, "r", encoding="utf8"))
        if not data[arg].get(str(event.user_id)):
            await statistics.finish("<EFBFBD><EFBFBD><EFBFBD>û<EFBFBD>ͳ<EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>ݲ<EFBFBD><EFBFBD><EFBFBD><EFBFBD><EFBFBD>...", at_sender=True)
    day_index = data["day_index"]
    data = data[arg][key]
    if _type == "group":
        name = await GroupInfo.get_group_info(event.group_id)
        name = name.group_name if name else str(event.group_id)
    else:
        name = event.sender.card or event.sender.nickname
    img = await generate_statistics_img(data, arg, name, plugin, day_index)
    await statistics.send(image(b64=img))


async def generate_statistics_img(
    data: dict, arg: str, name: str, plugin: str, day_index: int
):
    try:
        plugin = plugins2settings_manager.get_plugin_data(plugin)['cmd'][0]
    except (KeyError, IndexError):
        pass
    bar_graph = None
    if arg == "day_statistics":
        bar_graph = await init_bar_graph(data, f"{name} <20>չ<EFBFBD><D5B9>ܵ<EFBFBD><DCB5><EFBFBD>ͳ<EFBFBD><CDB3>")
    elif arg == "week_statistics":
        if plugin:
            current_week = day_index % 7
            week_lst = []
            if current_week == 0:
                week_lst = [1, 2, 3, 4, 5, 6, 7]
            else:
                for i in range(current_week + 1, 7):
                    week_lst.append(str(i))
                for i in range(current_week + 1):
                    week_lst.append(str(i))
            count = []
            for i in range(7):
                if int(week_lst[i]) == 7:
                    try:
                        count.append(data[str(0)][plugin])
                    except KeyError:
                        count.append(0)
                else:
                    try:
                        count.append(data[str(week_lst[i])][plugin])
                    except KeyError:
                        count.append(0)
            week_lst = ["7" if i == "0" else i for i in week_lst]
            bar_graph = BuildMat(
                y=count,
                mat_type="line",
                title=f"{name} <20><> {plugin} <20><><EFBFBD>ܵ<EFBFBD><DCB5><EFBFBD>ͳ<EFBFBD>ơ<EFBFBD>Ϊ7<CEAA><37>ͳ<EFBFBD>ơ<EFBFBD>",
                x_index=week_lst,
                display_num=True,
                background=[
                    f"{IMAGE_PATH}/background/create_mat/{x}"
                    for x in os.listdir(f"{IMAGE_PATH}/background/create_mat")
                ],
                bar_color=["*"],
            )
        else:
            bar_graph = await init_bar_graph(update_data(data), f"{name} <20>ܹ<EFBFBD><DCB9>ܵ<EFBFBD><DCB5><EFBFBD>ͳ<EFBFBD>ơ<EFBFBD>Ϊ7<CEAA><37>ͳ<EFBFBD>ơ<EFBFBD>")
    elif arg == "month_statistics":
        if plugin:
            day_index = day_index % 30
            day_lst = []
            for i in range(day_index + 1, 30):
                day_lst.append(i)
            for i in range(day_index + 1):
                day_lst.append(i)
            count = [data[str(day_lst[i])][plugin] for i in range(30)]
            day_lst = [str(x + 1) for x in day_lst]
            bar_graph = BuildMat(
                y=count,
                mat_type="line",
                title=f"{name} <20><> {plugin} <20><><EFBFBD>ܵ<EFBFBD><DCB5><EFBFBD>ͳ<EFBFBD>ơ<EFBFBD>Ϊ30<33><30>ͳ<EFBFBD>ơ<EFBFBD>",
                x_index=day_lst,
                display_num=True,
                background=[
                    f"{IMAGE_PATH}/background/create_mat/{x}"
                    for x in os.listdir(f"{IMAGE_PATH}/background/create_mat")
                ],
                bar_color=["*"],
            )
        else:
            bar_graph = await init_bar_graph(update_data(data), f"{name} <20>¹<EFBFBD><C2B9>ܵ<EFBFBD><DCB5><EFBFBD>ͳ<EFBFBD>ơ<EFBFBD>Ϊ30<33><30>ͳ<EFBFBD>ơ<EFBFBD>")
    elif arg == "total_statistics":
        bar_graph = await init_bar_graph(data, f"{name} <20><><EFBFBD>ܵ<EFBFBD><DCB5><EFBFBD>ͳ<EFBFBD><CDB3>")
    await asyncio.get_event_loop().run_in_executor(None, bar_graph.gen_graph)
    return bar_graph.pic2bs4()


async def init_bar_graph(data: dict, title: str) -> BuildMat:
    return await asyncio.get_event_loop().run_in_executor(None, _init_bar_graph, data, title)


def _init_bar_graph(data: dict, title: str) -> BuildMat:
    bar_graph = BuildMat(
        y=[data[x] for x in data.keys() if data[x] != 0],
        mat_type="barh",
        title=title,
        x_index=[x for x in data.keys() if data[x] != 0],
        display_num=True,
        background=[
            f"{IMAGE_PATH}/background/create_mat/{x}"
            for x in os.listdir(f"{IMAGE_PATH}/background/create_mat")
        ],
        bar_color=["*"],
    )
    return bar_graph


def update_data(data: dict):
    tmp_dict = {}
    for day in data.keys():
        for plugin_name in data[day].keys():
            # print(f'{day}<7D><>{plugin_name} = {data[day][plugin_name]}')
            if data[day][plugin_name] is not None:
                if tmp_dict.get(plugin_name) is None:
                    tmp_dict[plugin_name] = 1
                else:
                    tmp_dict[plugin_name] += data[day][plugin_name]
    return tmp_dict