from django.conf import settings
from django.core.management.base import BaseCommand
from gm_types.gaia import DIARY_CONTENT_LEVEL

from utils.execel import ExcelWriter
from talos.models.diary import Diary, DiaryTag

CONTENT_LEVEL = [
    DIARY_CONTENT_LEVEL.GENERAL, DIARY_CONTENT_LEVEL.FINE,
    DIARY_CONTENT_LEVEL.EXCELLENT, DIARY_CONTENT_LEVEL.BETTER]

# 指定的标签id
tag_ids = [10, 161, 162, 163, 164, 165, 166, 167, 767, 768, 866, 867, 955, 956, 959, 960,
           4282, 6684, 7091, 7175, 7176, 7177, 7178, 7179]


class Command(BaseCommand):
    def handle(self, *args, **options):
        """
        根据指定tag_id 导出日记本
        :param args:
        :param options:
        :return:export_diary_by_tag_ids
        """
        print('START')
        excel = ExcelWriter("diary.xlsx")
        diary_ids = set()
        for tag_id in tag_ids:
            d_ids = set(DiaryTag.objects.using(settings.SLAVE_DB_NAME).filter(
                tag_id=tag_id).values_list('diary_id', flat=True))
            diary_ids.update(d_ids)

        diary_ids = Diary.objects.using(settings.SLAVE_DB_NAME).filter(
            id__in=diary_ids, content_level__in=CONTENT_LEVEL, is_online=True
        ).values_list('id', flat=True).distinct()

        excel.write_header(['日记本ID'])
        row = 2
        for d_id in diary_ids.iterator():
            excel.write_row(row_num=row, values=[d_id])
            row += 1

        excel.save()
        print('DONE，日记本总数：%s' % (row-1))
