From 13d9a76bedcd03965b54b6f39d5500c74015ef45 Mon Sep 17 00:00:00 2001
From: ipu <uwu@ipu.moe>
Date: Fri, 3 Oct 2025 22:53:17 +0300
Subject: [PATCH] new estimation data model

---
 src/api/v1/router.py                  |  2 +-
 src/models.py                         |  6 ++++--
 src/services/estimation_service_v2.py | 14 ++++++++++----
 3 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/src/api/v1/router.py b/src/api/v1/router.py
index 646442c..81537c0 100644
--- a/src/api/v1/router.py
+++ b/src/api/v1/router.py
@@ -85,7 +85,7 @@ async def estimate(request: models.EstimationRequest):
             )
         
         estimation_service = EstimationService()
-        estimation_response = await estimation_service.estimate_insurance(request.applicants, request.phq, request.plans)
+        estimation_response = await estimation_service.estimate_insurance(request.applicants, request.phq, request.plans, request.coverage)
 
         return estimation_response
         
diff --git a/src/models.py b/src/models.py
index c0ee547..7cff934 100644
--- a/src/models.py
+++ b/src/models.py
@@ -17,7 +17,7 @@ class Applicant(BaseModel):
 
 class Plan(BaseModel):
     id: int
-    priceId: int
+    priceId: int | None = None
 
 class Medication(BaseModel):
     applicant: int
@@ -60,6 +60,7 @@ class Address(BaseModel):
 
 class EstimationRequest(BaseModel):
     userId: str | int | None = Field(None, description="Unique identifier")
+    coverage: int = 1
     applicants: List[Applicant]
     plans: List[Plan]
     phq: PHQ
@@ -69,7 +70,8 @@ class EstimationRequest(BaseModel):
 class EstimationDetails(BaseModel):
     dtq: bool
     reason: str
-    price_id: int
+    price_id: int = -1
+    tier: str
 
 class EstimationResult(BaseModel):
     name: str
diff --git a/src/services/estimation_service_v2.py b/src/services/estimation_service_v2.py
index 553c85f..2f7da6c 100644
--- a/src/services/estimation_service_v2.py
+++ b/src/services/estimation_service_v2.py
@@ -458,7 +458,7 @@ class EstimationService:
             return None
 
         
-    async def estimate_insurance(self, applicants: list[Applicant],  phq: PHQ, plans: list[Plan]):
+    async def estimate_insurance(self, applicants: list[Applicant],  phq: PHQ, plans: list[Plan], plan_coverage: int):
         estimation_results = []
         is_review = False
         review_reasons = []
@@ -478,7 +478,9 @@ class EstimationService:
                 base_tier = tier
                 break
 
-        plan_coverage = self.get_plan_coverage(plans[0])
+        if not plan_coverage:
+            plan_coverage = self.get_plan_coverage(plans[0])
+
         rx_spend = 0
         for applicant_id, applicant in enumerate(applicants):
             applicant_review_reasons = []
@@ -559,7 +561,7 @@ class EstimationService:
                     message=final_reason,
                 )
             )
-        
+
         plan_price_id = self.get_plan_price(plans[0], base_tier, plan_coverage)
         
         if is_dtq:
@@ -570,6 +572,7 @@ class EstimationService:
                     dtq=is_dtq,
                     reason=reason,
                     price_id=plan_price_id,
+                    tier=f"Tier {base_tier.value}",
                 ),
                 results=estimation_results
             )
@@ -582,6 +585,7 @@ class EstimationService:
                     dtq=is_dtq,
                     reason=reason,
                     price_id=plan_price_id,
+                    tier=f"Tier {base_tier.value}",
                 ),
                 results=estimation_results
             )
@@ -597,6 +601,7 @@ class EstimationService:
                     dtq=True,
                     reason=reason,
                     price_id=plan_price_id,
+                    tier=f"Tier {base_tier.value}",
                 ),
                 results=estimation_results
             )
@@ -614,6 +619,7 @@ class EstimationService:
                     dtq=is_dtq,
                     reason=reason,
                     price_id=plan_price_id,
+                    tier=f"Tier {base_tier.value}",
                 ),
                 results=estimation_results
             )
@@ -625,7 +631,7 @@ class EstimationService:
                     dtq=is_dtq,
                     reason=reason,
                     price_id=plan_price_id,
+                    tier=f"Tier {base_tier.value}",
                 ),
                 results=estimation_results
             )
-        
\ No newline at end of file