from decimal import Decimal
from typing import List

from sqlalchemy.orm import selectinload
from sqlmodel import select

from app.model.ApproveModel import ApproveService, ApproveModel
from app.model.ProjectModel import ProjectModel
from app.utils.create_module_service import create_model_service
from app.utils.db_utils import AsyncSessionDep

ProjectService = create_model_service(
  Cls=ProjectModel,
  custom_query=lambda: select(ProjectModel).options(
    selectinload(ProjectModel.leader),
  ),
  after_query_list=lambda query_cls_list, has_next, query_param, session: set_list_spent(query_cls_list, session),
  after_query_item=lambda item_cls, row_dict, session: set_item_spent(item_cls, session),
  after_insert=lambda insert_cls, row_dict, session: set_item_spent(insert_cls, session),
  after_update=lambda update_cls, row_dict, session: set_item_spent(update_cls, session),
  after_batch_insert=lambda refresh_cls_list, row_dict_list, session: set_list_spent(refresh_cls_list, session),
  after_batch_update=lambda refresh_cls_list, row_dict_list, session: set_list_spent(refresh_cls_list, session),
)


async def set_list_spent(cls_list: List[ProjectModel], session: AsyncSessionDep):
  project_dict_list, project_dict_map = await query_proj_spent([item.id for item in cls_list], session)
  for item in cls_list:
    target = project_dict_map.get(item.id, None)
    if target:
      item.spent = target['spent']
      item.balance = target['balance']

  for item in cls_list:
    print("==============>>>>>>>>>>>>")
    print(item.id, item.spent, item.balance)
    dict_val = item.model_dump()
    print(dict_val.get('spent'), dict_val.get('balance'))


async def set_item_spent(cls: ProjectModel, session: AsyncSessionDep):
  project_dict_list, project_dict_map = await query_proj_spent([cls.id], session)
  target = project_dict_map.get(cls.id, None)
  if target:
    cls.spent = target['spent']
    cls.balance = target['balance']


async def query_proj_spent(proj_id_list: List[str], session: AsyncSessionDep):
  result = {item: Decimal(0) for item in proj_id_list}
  approve_cls_list = (await session.execute(ApproveService.select_cls().where(ApproveModel.proj_id.in_(proj_id_list)))).scalars().all()
  for approve_item in approve_cls_list:
    if approve_item.status == 'approved':
      result[approve_item.proj_id] += Decimal(approve_item.amount)

  project_cls_list = (await session.execute(ProjectService.select_cls().where(ProjectModel.id.in_(proj_id_list)))).scalars().all()
  project_dict_list = [item.model_dump() for item in project_cls_list]
  for dict_item in project_dict_list:
    # 已花费金额
    dict_item['spent'] = result[dict_item['id']]
    # 剩余金额
    dict_item['balance'] = Decimal(dict_item['budget']) - Decimal(dict_item['spent'])
  project_dict_map = {item['id']: item for item in project_dict_list}
  return project_dict_list, project_dict_map