from apps.mtm.models import Material
from typing import List
from django.db.models import Sum
from apps.inm.models import MaterialBatch
from apps.wpm.models import WMaterial

def cal_material_count(materialId_list: List[str]=None):
    """
    计算物料总数量
    """
    if materialId_list is None:
        materialId_list = []
    if materialId_list:
        objs = Material.objects.filter(id__in=set(materialId_list))
    else:
        objs = Material.objects.all()
    for material in objs:
        mb_count = MaterialBatch.objects.filter(material=material, state=10).aggregate(total=Sum("count"))["total"] or 0
        wm_count = WMaterial.objects.filter(material=material, state=10).aggregate(total=Sum("count"))["total"] or 0
        if mb_count is None:
            mb_count = 0
        if wm_count is None:
            wm_count = 0
        Material.objects.filter(id=material.id).update(
            count_wm=wm_count,
            count_mb=mb_count,
            count=mb_count + wm_count)