from tqdm.asyncio import tqdm
from applications.utils import run_tasks_with_asyncio_task_group


class DataAnalysis:
    def __init__(self, pool):
        self.pool = pool

    async def get_articles(self):
        query = """
            select id, title
            from publish_account_category_detail
            where category is null;
        """
        return await self.pool.async_fetch(query=query)

    async def get_title_category(self, title: str):
        query = f"""
            select category from article_category
            where title = %s and status = 2 and version = 2
            order by id desc limit 1;
        """
        return await self.pool.async_fetch(query=query, params=(title,))

    async def set_each_record(self, article):
        id_ = article['id']
        title = article['title']
        result = await self.get_title_category(title)
        if not result:
            category = "empty"
        else:
            category = result[0]['category']
        update_query = """
            update publish_account_category_detail
            set category = %s
            where id = %s;
        """
        await self.pool.async_save(query=update_query, params=(category, id_))

    async def deal(self):
        articles = await self.get_articles()
        return await run_tasks_with_asyncio_task_group(
            task_list=articles,
            handler=self.set_each_record,
            max_concurrency=1,
            fail_fast=False,
            description="更新文章分类",
            unit="id",
        )
        for article in tqdm(articles):
            await self.set_each_record(article)