From c31f8e82e4d15be1e6001e61f126fe4849215eaa Mon Sep 17 00:00:00 2001
From: ipu <uwu@ipu.moe>
Date: Tue, 19 Aug 2025 13:22:19 +0300
Subject: [PATCH] get coverage from priceId

---
 src/models.py                         |  3 +--
 src/services/estimation_service_v2.py | 26 +++++++++++++++-----------
 2 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/src/models.py b/src/models.py
index 3c2f6ff..4c2b753 100644
--- a/src/models.py
+++ b/src/models.py
@@ -17,8 +17,7 @@ class Applicant(BaseModel):
 
 class Plan(BaseModel):
     id: int
-    coverage: int
-    tier: Optional[str] = Field(None, description="Tier assignment")
+    priceId: int
 
 class Medication(BaseModel):
     applicant: int
diff --git a/src/services/estimation_service_v2.py b/src/services/estimation_service_v2.py
index 72f51db..c7b1007 100644
--- a/src/services/estimation_service_v2.py
+++ b/src/services/estimation_service_v2.py
@@ -213,12 +213,19 @@ class EstimationService:
         
         return None
 
+    def get_plan_coverage(self, plan: Plan):
+        plan_data = get_plan_by_id(plan.id)
+        prices = plan_data["prices"]
+        for price in prices:
+            if price["id"] == plan.priceId:
+                return price["coverage"]
+        return 0
+
     def get_plan_price(self, plan: Plan, tier: Tier, coverage: int):
         plan_data = get_plan_by_id(plan.id)
         prices = plan_data["prices"]
-        tier_str = f"Tier {tier.value}"
         for price in prices:
-            if price["tier"] == tier_str and price["coverage"] == coverage:
+            if price["coverage"] == coverage and price["tier"] == f"Tier {tier.value}":
                 return price["price"]
         return 0
 
@@ -242,6 +249,7 @@ class EstimationService:
                 base_tier = tier
                 break
         
+        plan_coverage = self.get_plan_coverage(plans[0])
         rx_spend = 0
         for applicant_id, applicant in enumerate(applicants):
             applicant_age = self.calculate_age(applicant.dob)
@@ -266,7 +274,7 @@ class EstimationService:
             rx_spend_applicant = await self.calculate_rx_spend(phq, applicant_id)
             rx_spend += rx_spend_applicant
 
-            applicant_new_tier = self.get_tier(plans[0].coverage, rx_spend_applicant)
+            applicant_new_tier = self.get_tier(plan_coverage, rx_spend_applicant)
             if applicant_new_tier is None:
                 is_dtq = True
                 reason = "Declined due to high Rx spend"
@@ -287,7 +295,7 @@ class EstimationService:
                 )
             )
         
-        plan_price = self.get_plan_price(plans[0], base_tier, plans[0].coverage)
+        plan_price = self.get_plan_price(plans[0], base_tier, plan_coverage)
         
         if is_dtq:
             return EstimationResponse(
@@ -301,7 +309,7 @@ class EstimationService:
                 results=estimation_results
             )
         
-        new_tier = self.get_tier(plans[0].coverage, rx_spend)
+        new_tier = self.get_tier(plan_coverage, rx_spend)
         
         if new_tier is None:
             return EstimationResponse(
@@ -318,7 +326,7 @@ class EstimationService:
         if new_tier > base_tier:
             base_tier = new_tier
         
-        plan_price = self.get_plan_price(plans[0], base_tier, plans[0].coverage)
+        plan_price = self.get_plan_price(plans[0], base_tier, plan_coverage)
         
         status = "accepted" if base_tier is not None else "rejected"
         
@@ -332,8 +340,4 @@ class EstimationService:
             ),
             results=estimation_results
         )
-        
-
-
-
-    
\ No newline at end of file
+        
\ No newline at end of file