api/models.py

"""
API Models and Schemas for HBV AI Assistant
"""
from pydantic import BaseModel, Field, validator
from typing import Optional, List, Dict, Any
from enum import Enum


class QueryType(str, Enum):
    """Types of medical queries supported"""
    GENERAL = "general"
    DRUG_INTERACTION = "drug_interaction"
    SIDE_EFFECTS = "side_effects"
    GUIDELINES = "guidelines"
    COMPARISON = "comparison"


class QueryRequest(BaseModel):
    """Request model for medical queries"""
    query: str = Field(..., description="Medical question or query", min_length=1)
    query_type: Optional[QueryType] = Field(None, description="Type of medical query")
    context: Optional[str] = Field(None, description="Additional context for the query")
    patient_info: Optional[Dict[str, Any]] = Field(None, description="Patient information if relevant")


class QueryResponse(BaseModel):
    """Response model for medical queries"""
    response: str = Field(..., description="AI-generated medical response")
    sources: Optional[List[str]] = Field(None, description="Sources used for the response")
    confidence: Optional[float] = Field(None, description="Confidence score of the response")
    query_type: Optional[QueryType] = Field(None, description="Detected or specified query type")
    processing_time: Optional[float] = Field(None, description="Time taken to process the query")


class StreamChunk(BaseModel):
    """Model for streaming response chunks"""
    chunk: str = Field(..., description="Chunk of the streaming response")
    is_final: bool = Field(False, description="Whether this is the final chunk")


class SideEffectReport(BaseModel):
    """Model for side effect reporting"""
    drug_name: str = Field(..., description="Name of the drug")
    side_effects: str = Field(..., description="Reported side effects")
    patient_age: Optional[int] = Field(None, description="Patient age")
    patient_gender: Optional[str] = Field(None, description="Patient gender")
    dosage: Optional[str] = Field(None, description="Drug dosage")
    duration: Optional[str] = Field(None, description="Duration of treatment")
    severity: Optional[str] = Field(None, description="Severity of side effects")
    outcome: Optional[str] = Field(None, description="Outcome of the side effects")
    additional_details: Optional[str] = Field(None, description="Additional clinical details")


class SideEffectResponse(BaseModel):
    """Response model for side effect reporting"""
    report_id: str = Field(..., description="Unique identifier for the report")
    status: str = Field(..., description="Status of the report submission")
    message: str = Field(..., description="Confirmation message")
    recommendations: Optional[List[str]] = Field(None, description="Clinical recommendations")


class ComparisonRequest(BaseModel):
    """Request model for provider/treatment comparisons"""
    providers: List[str] = Field(..., description="List of providers to compare", min_items=2)
    criteria: Optional[List[str]] = Field(None, description="Specific criteria for comparison")


class ComparisonResponse(BaseModel):
    """Response model for provider/treatment comparisons"""
    comparison: str = Field(..., description="Detailed comparison analysis")
    summary: Dict[str, Any] = Field(..., description="Summary of key differences")
    recommendations: Optional[str] = Field(None, description="Recommendations based on comparison")


class InitializationStatus(BaseModel):
    """Initialization status response model"""
    is_complete: bool = Field(..., description="Whether initialization is complete")
    status_message: str = Field(..., description="Current initialization status")
    is_successful: bool = Field(..., description="Whether initialization was successful")
    error: Optional[str] = Field(None, description="Initialization error if any")


class HealthStatus(BaseModel):
    """Health check response model"""
    status: str = Field(..., description="API health status")
    version: str = Field(..., description="API version")
    timestamp: str = Field(..., description="Current timestamp")
    components: Dict[str, str] = Field(..., description="Status of system components")
    initialization: Optional[InitializationStatus] = Field(None, description="Background initialization status")


class ErrorResponse(BaseModel):
    """Error response model"""
    error: str = Field(..., description="Error type")
    message: str = Field(..., description="Error message")
    details: Optional[Dict[str, Any]] = Field(None, description="Additional error details")
    timestamp: str = Field(..., description="Error timestamp")


# ============================================================================
# HBV PATIENT ASSESSMENT MODELS
# ============================================================================

class HBVPatientInput(BaseModel):
    """Input model for HBV patient assessment"""
    sex: str = Field(..., description="Patient sex: Male / Female")
    age: int = Field(..., description="Patient age in years", ge=0, le=120)
    pregnancy_status: str = Field(..., description="Pregnancy status: Not pregnant / Pregnant")
    hbsag_status: str = Field(..., description="HBsAg status: Positive / Negative")
    duration_hbsag_months: int = Field(..., description="Duration of HBsAg positivity in months", ge=0)
    hbv_dna_level: float = Field(..., description="HBV DNA level in IU/mL", ge=0)
    hbeag_status: str = Field(..., description="HBeAg status: Positive / Negative")
    alt_level: float = Field(..., description="ALT level in U/L", ge=0)
    fibrosis_stage: str = Field(..., description="Fibrosis/Cirrhosis stage: F0-F1 / F2-F3 / F4")
    necroinflammatory_activity: str = Field(..., description="Necroinflammatory activity: A0 / A1 / A2 / A3")
    extrahepatic_manifestations: bool = Field(..., description="Presence of extrahepatic manifestations")
    immunosuppression_status: Optional[str] = Field(None, description="Immunosuppression status: None / Chemotherapy / Other")
    coinfections: List[str] = Field(default_factory=list, description="Coinfections: HIV, HCV, HDV")
    family_history_cirrhosis_hcc: bool = Field(..., description="Family history of Cirrhosis or HCC (first-degree relative)")
    other_comorbidities: Optional[List[str]] = Field(None, description="Other comorbidities")
    
    @validator('sex')
    def validate_sex(cls, v):
        value = str(v).strip().lower()
        if value in ['male', 'm', 'man', 'boy']:
            return 'Male'
        if value in ['female', 'f', 'woman', 'girl']:
            return 'Female'
        raise ValueError('Sex must be either Male or Female')
    
    @validator('pregnancy_status')
    def validate_pregnancy(cls, v):
        value = str(v).strip().lower()
        if value in ['pregnant', 'yes', 'y']:
            return 'Pregnant'
        if value in ['not pregnant', 'non-pregnant', 'non pregnant', 'no', 'n', 'none']:
            return 'Not pregnant'
        raise ValueError('Pregnancy status must be either "Not pregnant" or "Pregnant"')
    
    @validator('hbsag_status')
    def validate_hbsag(cls, v):
        value = str(v).strip().lower()
        if value in ['positive', 'pos', '+', 'reactive']:
            return 'Positive'
        if value in ['negative', 'neg', '-', 'non-reactive', 'nonreactive']:
            return 'Negative'
        raise ValueError('HBsAg status must be either Positive or Negative')
    
    @validator('hbeag_status')
    def validate_hbeag(cls, v):
        value = str(v).strip().lower()
        if value in ['positive', 'pos', '+', 'reactive']:
            return 'Positive'
        if value in ['negative', 'neg', '-', 'non-reactive', 'nonreactive']:
            return 'Negative'
        raise ValueError('HBeAg status must be either Positive or Negative')
    
    @validator('fibrosis_stage')
    def validate_fibrosis(cls, v):
        value = str(v).strip().upper().replace(" ", "")
        if value in ['F0-F1', 'F2-F3', 'F4']:
            return value
        if value in ['F0', 'F1']:
            return 'F0-F1'
        if value in ['F2', 'F3']:
            return 'F2-F3'
        if value == 'F4':
            return 'F4'
        raise ValueError('Fibrosis stage must be F0-F1, F2-F3, or F4 (or F0, F1, F2, F3, F4 which will be mapped to these categories)')
    
    @validator('necroinflammatory_activity')
    def validate_necroinflammatory(cls, v):
        value = str(v).strip().upper().replace(" ", "")
        if value in ['A0', 'A1', 'A2', 'A3']:
            return value
        if value in ['0', '1', '2', '3']:
            return 'A' + value
        raise ValueError('Necroinflammatory activity must be A0, A1, A2, or A3')


class HBVAssessmentResponse(BaseModel):
    """Response model for HBV patient assessment"""
    eligible: bool = Field(..., description="Whether patient is eligible for treatment")
    recommendations: str = Field(..., description="Detailed recommendations with citations from SASLT 2021 guidelines including page numbers")


class TextAssessmentInput(BaseModel):
    """Input model for text-based HBV patient assessment"""
    text_input: str = Field(..., description="Free-form text containing patient data", min_length=10)
    
    class Config:
        json_schema_extra = {
            "example": {
                "text_input": "45-year-old male patient\\nHBsAg: Positive for 12 months\\nHBV DNA: 5000 IU/mL\\nHBeAg: Positive\\nALT: 80 U/L\\nFibrosis stage: F2-F3\\nNecroinflammatory activity: A2\\nNo extrahepatic manifestations\\nNo immunosuppression\\nNo coinfections (HIV, HCV, HDV)\\nNo family history of cirrhosis or HCC"
            }
        }


class ChatRequest(BaseModel):
    """Request model for chat interactions"""
    query: str = Field(..., description="Doctor's question about HBV guidelines", min_length=1)
    session_id: str = Field(default="default", description="Session identifier for conversation continuity")
    patient_context: Optional[HBVPatientInput] = Field(None, description="Optional patient context from assessment")
    assessment_result: Optional[HBVAssessmentResponse] = Field(None, description="Optional assessment result for context")


class ChatResponse(BaseModel):
    """Response model for chat interactions"""
    response: str = Field(..., description="AI response to the doctor's question")
    session_id: str = Field(..., description="Session identifier")