modules/models.py

from dataclasses import dataclass
from datetime import date, datetime
from pydantic import BaseModel, Field
from typing import List, Literal, Optional, TypedDict, Union


class PatientInfo(BaseModel):
    name: Optional[str] = Field(None, description="Patient's full name")
    age: Optional[int] = Field(None, description="Patient's age in years")
    sex: Optional[str] = Field(None, description="Male/Female/Other")
    medical_record_number: Optional[str] = None

    class Config:
        extra = "forbid"  # 🚨 ensures schema matches OpenAI’s strict rules


class TestResultReferenceRange(BaseModel):
    min: Optional[float] = None
    max: Optional[float] = None


class LabResult(BaseModel):
    test_name: str
    result_value: str
    test_unit: str
    test_reference_range: Optional[TestResultReferenceRange] = None
    test_date: Optional[str] = Field(
        None,
        description="Date of test in ISO format (YYYY-MM-DD). Convert all dates into this format.",
    )
    inferred_range: Literal["low", "normal", "high"]

    class Config:
        extra = "forbid"


class CompositeLabResult(BaseModel):
    section_name: str  # e.g., "CUE - COMPLETE URINE ANALYSIS"
    sub_results: List[LabResult]  # each individual parameter under the section

    class Config:
        extra = "forbid"


class StandardizedReport(BaseModel):
    original_report_file_name: str
    patient_info: PatientInfo
    lab_results: List[
        Union[LabResult, CompositeLabResult]
    ]  # ✅ supports both flat + grouped results
    diagnosis: List[str]
    recommendations: List[str]

    class Config:
        extra = "forbid"


@dataclass
class HealthReport:
    report_file_name_with_path: str
    report_file_name: str
    report_contents: str


@dataclass
class SheamiMilestone:
    step_name: str
    status: str  # e.g., "started", "in_progress", "completed", "failed"
    start_time: Optional[datetime] = None
    end_time: Optional[datetime] = None

    @property
    def time_taken(self) -> Optional[float]:
        """Return time taken in seconds if available."""
        if self.start_time and self.end_time:
            return (self.end_time - self.start_time).total_seconds()
        return None

    @property
    def status_icon(self) -> str:
        if self.status == "started" or self.status == "inprogress":
            return "⏳"
        elif self.status == "completed":
            return "✅"
        elif self.status == "failed":
            return "❌"
        else:
            return self.status


class SheamiState(TypedDict):
    user_email: str
    patient_id: str
    report_id_list: str
    run_id: str
    messages: list[str]
    thread_id: str
    uploaded_reports: List[HealthReport]
    standardized_reports: List[StandardizedReport]
    trends_json: dict
    pdf_path: str
    current_index: int
    process_desc: str
    units_processed: int
    units_total: int
    overall_units_processed: int
    overall_units_total: int
    milestones: list[SheamiMilestone]
    interpretation_html: str

@dataclass
class SheamiUser:
    email : str
    name : str
    picture_url : str