From ce5d79f4ae73e8f4b26f1a9bb6f63d186c056f2c Mon Sep 17 00:00:00 2001
From: Alexander Orvik <alexao1111@gmail.com>
Date: Mon, 5 Aug 2024 22:25:34 +0200
Subject: [PATCH] feat(schedules): add first draft schedule allergy query

---
 schedules/schemas/schedules.py | 88 +++++++++++++++++++++++++++++++++-
 1 file changed, 87 insertions(+), 1 deletion(-)

diff --git a/schedules/schemas/schedules.py b/schedules/schemas/schedules.py
index 57424c1b..1ab2bc8b 100644
--- a/schedules/schemas/schedules.py
+++ b/schedules/schemas/schedules.py
@@ -22,10 +22,12 @@
 )
 from schedules.utils.schedules import normalize_shifts, send_given_shift_email
 from schedules.utils.templates import apply_schedule_template
-from users.models import User
+from users.models import User, Allergy as UserAllergy
 from django.utils import timezone
 from django.conf import settings
 
+from django.db.models import Count
+from django.db.models.functions import TruncDate
 
 class ShiftInterestNode(DjangoObjectType):
     class Meta:
@@ -82,6 +84,20 @@ def resolve_filled_slots(self: Shift, info):
     def get_node(cls, info, id):
         return Shift.objects.get(pk=id)
 
+class Allergy(graphene.ObjectType):
+    name = graphene.String()
+    count = graphene.Int()
+
+class DayAllergyNode(graphene.ObjectType):
+    """
+    
+    """
+    date = graphene.Date()
+    allergy_list = graphene.NonNull(
+        graphene.List(graphene.NonNull(Allergy))
+    )
+
+
 
 class ScheduleNode(DjangoObjectType):
     class Meta:
@@ -113,10 +129,80 @@ def get_node(cls, info, id):
 class ScheduleQuery(graphene.ObjectType):
     schedule = Node.Field(ScheduleNode)
     all_schedules = graphene.NonNull(graphene.List(ScheduleNode, required=True))
+    schedule_allergies = graphene.List(DayAllergyNode, date_from=graphene.Date(), date_to= graphene.Date())
 
     def resolve_all_schedules(self, info, *args, **kwargs):
         return Schedule.objects.all().order_by("name")
 
+    def resolve_schedule_allergies(self, info, date_from, date_to, *args, **kwargs):
+        aware_from = timezone.make_aware(
+            timezone.datetime(
+                year=date_from.year,
+                month=date_from.month,
+                day=date_from.day,
+                hour=0,
+                minute=0,
+                second=0
+            ),
+            timezone=pytz.timezone(settings.TIME_ZONE)
+        )
+        aware_to = timezone.make_aware(
+            timezone.datetime(
+                year=date_to.year,
+                month=date_to.month,
+                day=date_to.day,
+                hour=23,
+                minute=59,
+                second=59
+            ),
+            timezone=pytz.timezone(settings.TIME_ZONE)
+        )
+
+        filtered_shifts = Shift.objects.filter(datetime_start__range=(aware_from, aware_to), slots__user__isnull=False)
+        allergy_counts = (
+            filtered_shifts
+            .annotate(shift_date=TruncDate('datetime_start'))
+            .values('shift_date', 'slots__user__allergies__name')
+            .annotate(count=Count('slots__user__allergies'))
+            .order_by('shift_date', 'slots__user__allergies__name')
+        )
+
+        results = [
+            {
+                'shift_date': entry['shift_date'],
+                'allergy': entry['slots__user__allergies__name'],
+                'count': entry['count']
+            }
+            for entry in allergy_counts
+        ]
+
+        result_dict = {}
+
+        for item in results:
+            shift_date = item["shift_date"]
+            shift_date_key = shift_date.strftime("%Y-%m-%d")
+            allergy_list = result_dict.get(shift_date_key, (shift_date, []))[1]
+            allergy_name = item["allergy"]
+            allergy_count = item["count"]
+
+            allergy_list.append((allergy_name, allergy_count))
+            result_dict[shift_date_key] =  (shift_date, allergy_list)
+
+
+        final_list = []
+        for key in result_dict.keys():
+            date_obj, allergies = result_dict[key]
+
+            day_list = []
+            for allergu_tuple in allergies:
+                allergy_name, allergy_count = allergu_tuple
+                day_list.append(
+                    Allergy(name=allergy_name, count=allergy_count)
+                )
+
+            final_list.append(DayAllergyNode(date=date_obj, allergy_list=day_list))
+        return final_list
+
 
 # === Single location grouping types ===
 class ShiftDayGroup(graphene.ObjectType):