"""给日记本打公立tag"""

import math
from datetime import time
from multiprocessing import Pool, Manager

from django.core.management.base import BaseCommand
from django.conf import settings
from django.db.models import Q, Max
from django import db

from gm_types.gaia import HOSPITAL_TYPE

from talos.models.diary import Diary
from talos.services import DoctorService
from talos.services.hospital import HospitalService
from talos.models.diary import Diary



def has_public_tag(diaries):

    doctor_ids = [diary.doctor_id for diary in diaries if diary.doctor_id]
    doctors = DoctorService.get_doctor_from_doctor_ids(doctor_ids)
    doctors_info = {doctor.id: doctor for doctor in doctors}
    hospital_ids = [diary.hospital_id for diary in diaries if diary.hospital_id]
    hospitals = HospitalService.get_hospital_from_hospital_ids(hospital_ids)
    hospitals_info = {hospital.id: hospital for hospital in hospitals}

    res = {}
    for diary in diaries:

        if diary.doctor_id:
            doctor = doctors_info.get(diary.doctor_id)
            if doctor:
                try:
                    res[diary.id] = True if doctor.hospital.hospital_type == HOSPITAL_TYPE.PUBLIC else False
                    continue
                except:
                    pass

        if diary.hospital_id:
            hospital = hospitals_info.get(diary.hospital_id)
            if hospital:
                try:
                    res[diary.id] = True if hospital.hospital_type == HOSPITAL_TYPE.PUBLIC else False
                    continue
                except:
                    pass

        res[diary.id] = False

    return res


def update_diary_tag(queue, limit):
    """更新日记本标签"""

    start_id = queue.get()
    print(start_id)

    diaries = Diary.objects.filter(Q(pk__gt=start_id))[: limit]
    max_id = diaries.aggregate(max_id=Max('id'))
    queue.put(max_id["max_id"])

    if not diaries:
        return

    diary_has_tag = has_public_tag(diaries)
    for diary in diaries:
        if diary_has_tag.get(diary.id):
            diary.add_tags([settings.NEW_PUBLIC_TAG])


class Command(BaseCommand):

    def handle(self, *args, **options):
        print('------ starting -----')
        start_time = time.time()
        print("start at: ", start_time)

        queue = Manager().Queue(maxsize=4)
        queue.put(0)  # 触发程序开始

        args_list = []
        per_num = 200
        count = Diary.objects.filter(is_online=True).count()
        cnt = int(math.ceil(count / per_num))
        for _ in range(cnt):
            args_list.append((queue, per_num))

        db.connections.close_all()
        pool = Pool(processes=4)
        pool.starmap(update_diary_tag, args_list)
        pool.close()
        pool.join()

        end_time = time.time()
        print("end at: ", end_time)
        print('total use {} s.'.format(end_time - start_time))
        print('Done!')
