health_data/pydantic_models.py

"""
Pydantic Models for Health Data
Provides automatic validation and parsing
"""

from datetime import datetime
from typing import Optional, List, Dict, Any, Union
from pydantic import BaseModel, Field, field_validator, model_validator
from enum import Enum

from .validators import HealthDataParser, HealthDataValidator


class Gender(str, Enum):
    """Gender enum"""
    MALE = "male"
    FEMALE = "female"
    OTHER = "other"


class ActivityLevel(str, Enum):
    """Activity level enum"""
    SEDENTARY = "sedentary"  # Ít vận động
    LIGHT = "light"  # Vận động nhẹ
    MODERATE = "moderate"  # Vận động vừa
    ACTIVE = "active"  # Vận động nhiều
    VERY_ACTIVE = "very_active"  # Vận động rất nhiều


class FitnessLevel(str, Enum):
    """Fitness level enum"""
    BEGINNER = "beginner"
    INTERMEDIATE = "intermediate"
    ADVANCED = "advanced"


class RecordType(str, Enum):
    """Health record type"""
    NUTRITION = "nutrition"
    EXERCISE = "exercise"
    SYMPTOM = "symptom"
    MENTAL_HEALTH = "mental_health"
    GENERAL_HEALTH = "general_health"


class HealthRecord(BaseModel):
    """
    Health Record with Pydantic validation
    Automatically validates and normalizes health data
    """
    
    record_id: str = Field(default_factory=lambda: str(__import__('uuid').uuid4()))
    user_id: str
    record_type: RecordType
    data: Dict[str, Any] = Field(default_factory=dict)
    timestamp: datetime = Field(default_factory=datetime.now)
    agent_name: Optional[str] = None
    confidence: float = Field(default=0.5, ge=0.0, le=1.0)
    
    # Health metrics (optional, extracted from data)
    height: Optional[float] = Field(None, description="Height in cm")
    weight: Optional[float] = Field(None, description="Weight in kg")
    age: Optional[int] = Field(None, description="Age in years")
    gender: Optional[Gender] = None
    bmi: Optional[float] = Field(None, description="BMI")
    
    class Config:
        use_enum_values = True
        json_encoders = {
            datetime: lambda v: v.isoformat()
        }
    
    @field_validator('height', mode='before')
    @classmethod
    def parse_height(cls, v):
        """Parse height from various formats"""
        if v is None:
            return None
        parsed = HealthDataParser.parse_height(v)
        if parsed is not None:
            is_valid, error = HealthDataValidator.validate_height(parsed)
            if not is_valid:
                raise ValueError(error)
        return parsed
    
    @field_validator('weight', mode='before')
    @classmethod
    def parse_weight(cls, v):
        """Parse weight from various formats"""
        if v is None:
            return None
        parsed = HealthDataParser.parse_weight(v)
        if parsed is not None:
            is_valid, error = HealthDataValidator.validate_weight(parsed)
            if not is_valid:
                raise ValueError(error)
        return parsed
    
    @field_validator('age', mode='before')
    @classmethod
    def parse_age(cls, v):
        """Parse age from various formats"""
        if v is None:
            return None
        parsed = HealthDataParser.parse_age(v)
        if parsed is not None:
            is_valid, error = HealthDataValidator.validate_age(parsed)
            if not is_valid:
                raise ValueError(error)
        return parsed
    
    @field_validator('bmi', mode='before')
    @classmethod
    def parse_bmi(cls, v):
        """Parse BMI"""
        if v is None:
            return None
        parsed = HealthDataParser.parse_bmi(v)
        if parsed is not None:
            is_valid, error = HealthDataValidator.validate_bmi(parsed)
            if not is_valid:
                raise ValueError(error)
        return parsed
    
    @model_validator(mode='after')
    def calculate_bmi_if_missing(self):
        """Auto-calculate BMI if weight and height are provided"""
        if self.bmi is None and self.weight and self.height:
            self.bmi = HealthDataValidator.calculate_bmi(self.weight, self.height)
        return self


class UserHealthProfile(BaseModel):
    """
    User Health Profile with Pydantic validation
    """
    
    user_id: str
    age: Optional[int] = Field(None, ge=13, le=150)
    gender: Optional[Gender] = None
    weight: Optional[float] = Field(None, ge=20, le=300, description="Weight in kg")
    height: Optional[float] = Field(None, ge=50, le=300, description="Height in cm")
    bmi: Optional[float] = Field(None, ge=10, le=60)
    activity_level: Optional[ActivityLevel] = None
    fitness_level: Optional[FitnessLevel] = None
    health_conditions: List[str] = Field(default_factory=list)
    medications: List[str] = Field(default_factory=list)
    allergies: List[str] = Field(default_factory=list)
    dietary_restrictions: List[str] = Field(default_factory=list)
    created_at: datetime = Field(default_factory=datetime.now)
    updated_at: datetime = Field(default_factory=datetime.now)
    
    class Config:
        use_enum_values = True
        json_encoders = {
            datetime: lambda v: v.isoformat()
        }
    
    @field_validator('height', mode='before')
    @classmethod
    def parse_height(cls, v):
        """Parse height from various formats"""
        if v is None:
            return None
        return HealthDataParser.parse_height(v)
    
    @field_validator('weight', mode='before')
    @classmethod
    def parse_weight(cls, v):
        """Parse weight from various formats"""
        if v is None:
            return None
        return HealthDataParser.parse_weight(v)
    
    @field_validator('age', mode='before')
    @classmethod
    def parse_age(cls, v):
        """Parse age from various formats"""
        if v is None:
            return None
        return HealthDataParser.parse_age(v)
    
    @model_validator(mode='after')
    def calculate_bmi_if_missing(self):
        """Auto-calculate BMI if weight and height are provided"""
        if self.bmi is None and self.weight and self.height:
            self.bmi = HealthDataValidator.calculate_bmi(self.weight, self.height)
        return self
    
    def get_bmi_category(self) -> str:
        """Get BMI category"""
        return HealthDataValidator.get_bmi_category(self.bmi)
    
    def is_complete(self) -> bool:
        """Check if profile has all essential data"""
        return all([
            self.age is not None,
            self.gender is not None,
            self.weight is not None,
            self.height is not None
        ])
    
    def get_missing_fields(self) -> List[str]:
        """Get list of missing essential fields"""
        missing = []
        if self.age is None:
            missing.append('age')
        if self.gender is None:
            missing.append('gender')
        if self.weight is None:
            missing.append('weight')
        if self.height is None:
            missing.append('height')
        return missing


class NutritionRecord(HealthRecord):
    """Nutrition-specific health record"""
    
    record_type: RecordType = Field(default=RecordType.NUTRITION, frozen=True)
    calories: Optional[float] = Field(None, ge=0, le=10000)
    protein: Optional[float] = Field(None, ge=0, le=500)
    carbs: Optional[float] = Field(None, ge=0, le=1000)
    fat: Optional[float] = Field(None, ge=0, le=500)
    meal_type: Optional[str] = None  # breakfast/lunch/dinner/snack


class ExerciseRecord(HealthRecord):
    """Exercise-specific health record"""
    
    record_type: RecordType = Field(default=RecordType.EXERCISE, frozen=True)
    exercise_type: Optional[str] = None  # cardio/strength/flexibility/sports
    duration_minutes: Optional[int] = Field(None, ge=0, le=600)
    intensity: Optional[str] = None  # low/medium/high
    calories_burned: Optional[float] = Field(None, ge=0, le=5000)


class SymptomRecord(HealthRecord):
    """Symptom-specific health record"""
    
    record_type: RecordType = Field(default=RecordType.SYMPTOM, frozen=True)
    symptoms: List[str] = Field(default_factory=list)
    severity: Optional[int] = Field(None, ge=1, le=10)
    duration_days: Optional[int] = Field(None, ge=0, le=365)
    body_part: Optional[str] = None


class MentalHealthRecord(HealthRecord):
    """Mental health-specific health record"""
    
    record_type: RecordType = Field(default=RecordType.MENTAL_HEALTH, frozen=True)
    mood: Optional[str] = None
    stress_level: Optional[int] = Field(None, ge=1, le=10)
    sleep_hours: Optional[float] = Field(None, ge=0, le=24)
    sleep_quality: Optional[int] = Field(None, ge=1, le=10)