#!/usr/bin/env python3
"""
Test script for the updated Lifestyle Profile Updater with next_check_in functionality
"""

import json
from datetime import datetime, timedelta
from dataclasses import dataclass
from typing import List, Dict

@dataclass
class MockLifestyleProfile:
    patient_name: str = "Serhii"
    patient_age: str = "52"
    conditions: List[str] = None
    primary_goal: str = "Improve exercise tolerance safely"
    exercise_preferences: List[str] = None
    exercise_limitations: List[str] = None
    dietary_notes: List[str] = None
    personal_preferences: List[str] = None
    last_session_summary: str = ""
    progress_metrics: Dict = None
    
    def __post_init__(self):
        if self.conditions is None:
            self.conditions = ["Type 2 diabetes", "Hypertension"]
        if self.exercise_preferences is None:
            self.exercise_preferences = ["upper body exercises", "seated exercises"]
        if self.exercise_limitations is None:
            self.exercise_limitations = ["Right below knee amputation"]
        if self.dietary_notes is None:
            self.dietary_notes = ["Diabetic diet", "Low sodium"]
        if self.personal_preferences is None:
            self.personal_preferences = ["prefers gradual changes"]
        if self.progress_metrics is None:
            self.progress_metrics = {"baseline_bp": "148/98"}

class MockAPI:
    def generate_response(self, system_prompt: str, user_prompt: str, temperature: float = 0.3, call_type: str = "") -> str:
        """Mock response for profile updater"""
        
        # Simulate different scenarios based on session content
        if "new patient" in user_prompt.lower() or "first session" in user_prompt.lower():
            # New patient scenario - needs immediate follow-up
            return json.dumps({
                "updates_needed": True,
                "reasoning": "First lifestyle session completed. Patient shows motivation but needs close monitoring due to complex medical conditions.",
                "updated_fields": {
                    "exercise_preferences": ["upper body exercises", "seated exercises", "adaptive equipment training"],
                    "exercise_limitations": ["Right below knee amputation", "Monitor blood glucose before/after exercise"],
                    "dietary_notes": ["Diabetic diet", "Low sodium", "Discussed meal timing with exercise"],
                    "personal_preferences": ["prefers gradual changes", "wants medical supervision initially"],
                    "primary_goal": "Improve exercise tolerance safely with medical supervision",
                    "progress_metrics": {"baseline_bp": "148/98", "initial_motivation_level": "high"},
                    "session_summary": "Initial lifestyle assessment completed. Patient motivated to start adapted exercise program.",
                    "next_check_in": "2025-09-08"
                },
                "session_insights": "Patient demonstrates high motivation despite physical limitations. Requires careful medical supervision.",
                "next_session_rationale": "New patient with complex conditions needs immediate follow-up in 3 days to ensure safe program initiation."
            })
        
        elif "progress" in user_prompt.lower() or "week" in user_prompt.lower():
            # Ongoing coaching scenario - regular follow-up
            return json.dumps({
                "updates_needed": True,
                "reasoning": "Patient showing good progress with exercise program. Ready for program advancement.",
                "updated_fields": {
                    "exercise_preferences": ["upper body exercises", "seated exercises", "resistance band training"],
                    "progress_metrics": {"baseline_bp": "148/98", "week_2_bp": "142/92", "exercise_frequency": "3 times/week"},
                    "session_summary": "Good progress with exercise program. Patient comfortable with current routine.",
                    "next_check_in": "2025-09-19"
                },
                "session_insights": "Patient adapting well to exercise routine. Blood pressure showing improvement.",
                "next_session_rationale": "Stable progress allows for 2-week follow-up to monitor continued improvement."
            })
        
        elif "maintenance" in user_prompt.lower() or "stable" in user_prompt.lower():
            # Maintenance phase scenario - long-term follow-up
            return json.dumps({
                "updates_needed": False,
                "reasoning": "Patient in maintenance phase with stable progress and established routine.",
                "updated_fields": {
                    "session_summary": "Maintenance check-in. Patient continuing established routine successfully.",
                    "next_check_in": "2025-10-05"
                },
                "session_insights": "Patient has established sustainable lifestyle habits. Minimal intervention needed.",
                "next_session_rationale": "Maintenance phase patient can be followed up monthly to ensure continued adherence."
            })
        
        else:
            # Default scenario
            return json.dumps({
                "updates_needed": True,
                "reasoning": "Standard lifestyle coaching session completed.",
                "updated_fields": {
                    "session_summary": "Regular lifestyle coaching session completed.",
                    "next_check_in": "2025-09-12"
                },
                "session_insights": "Patient engaged in lifestyle coaching process.",
                "next_session_rationale": "Regular follow-up in 1 week for active coaching phase."
            })

def test_profile_updater_scenarios():
    """Test different scenarios for next_check_in planning"""
    
    print("🧪 Testing Lifestyle Profile Updater with Next Check-in Planning\n")
    
    api = MockAPI()
    profile = MockLifestyleProfile()
    
    # Test scenarios
    scenarios = [
        {
            "name": "New Patient - First Session",
            "session_context": "First lifestyle coaching session with new patient",
            "messages": [
                {"role": "user", "message": "I'm ready to start exercising but worried about my amputation"},
                {"role": "user", "message": "What exercises can I do safely?"}
            ]
        },
        {
            "name": "Active Coaching - Progress Check",
            "session_context": "Week 2 progress check - patient showing improvement",
            "messages": [
                {"role": "user", "message": "I've been doing the exercises 3 times this week"},
                {"role": "user", "message": "My blood pressure seems better"}
            ]
        },
        {
            "name": "Maintenance Phase - Stable Patient",
            "session_context": "Monthly maintenance check for stable patient",
            "messages": [
                {"role": "user", "message": "Everything is going well with my routine"},
                {"role": "user", "message": "I'm maintaining my exercise schedule"}
            ]
        }
    ]
    
    for scenario in scenarios:
        print(f"📋 **{scenario['name']}**")
        print(f"   Context: {scenario['session_context']}")
        
        # Simulate the prompt (simplified)
        user_prompt = f"""
        SESSION CONTEXT: {scenario['session_context']}
        PATIENT MESSAGES: {[msg['message'] for msg in scenario['messages']]}
        """
        
        try:
            response = api.generate_response("", user_prompt)
            result = json.loads(response)
            
            print(f"   ✅ Updates needed: {result.get('updates_needed')}")
            print(f"   📅 Next check-in: {result.get('updated_fields', {}).get('next_check_in', 'Not set')}")
            print(f"   💭 Rationale: {result.get('next_session_rationale', 'Not provided')}")
            print(f"   📝 Session summary: {result.get('updated_fields', {}).get('session_summary', 'Not provided')}")
            print()
            
        except Exception as e:
            print(f"   ❌ Error: {e}")
            print()

def test_next_checkin_date_formats():
    """Test different date format scenarios"""
    
    print("📅 Testing Next Check-in Date Formats\n")
    
    # Test different date scenarios
    today = datetime.now()
    
    date_scenarios = [
        ("Immediate follow-up", today + timedelta(days=2)),
        ("Short-term follow-up", today + timedelta(weeks=1)),
        ("Regular follow-up", today + timedelta(weeks=2)),
        ("Long-term follow-up", today + timedelta(weeks=4))
    ]
    
    for scenario_name, target_date in date_scenarios:
        formatted_date = target_date.strftime("%Y-%m-%d")
        print(f"   {scenario_name}: {formatted_date}")
    
    print("\n✅ Date format examples generated successfully")

if __name__ == "__main__":
    test_profile_updater_scenarios()
    test_next_checkin_date_formats()
    
    print("\n📋 **Summary of Next Check-in Feature:**")
    print("   • New patients: 1-3 days follow-up")
    print("   • Active coaching: 1 week follow-up")
    print("   • Stable progress: 2-3 weeks follow-up")
    print("   • Maintenance phase: 1 month+ follow-up")
    print("   • Date format: YYYY-MM-DD")
    print("   • Includes rationale for timing decision")
    print("\n✅ Profile updater enhanced with next session planning!")