#!/usr/bin/env python3
"""
Script to fix ObjectId serialization issues in server.py
"""

import re

# Read the server.py file
with open('/app/backend/server.py', 'r') as f:
    content = f.read()

# Define patterns and replacements for endpoints that return MongoDB documents
fixes = [
    # Diet plans endpoints
    (r'@api_router\.get\("/admin/diet-plans"\)\s*\nasync def get_all_diet_plans\(admin: dict = Depends\(get_admin_user\)\):\s*\n    plans = await db\.diet_plans\.find\(\)\.to_list\(1000\)\s*\n    return plans',
     '@api_router.get("/admin/diet-plans")\nasync def get_all_diet_plans(admin: dict = Depends(get_admin_user)):\n    plans = await db.diet_plans.find().to_list(1000)\n    return serialize_mongo_docs(plans)'),
    
    # Progress photos endpoints
    (r'@api_router\.get\("/user/progress-photos"\)\s*\nasync def get_progress_photos\(current_user: dict = Depends\(get_current_user\)\):\s*\n    photos = await db\.progress_photos\.find\({"user_id": current_user\["id"\]}\)\.sort\("created_at", -1\)\.to_list\(100\)\s*\n    return photos',
     '@api_router.get("/user/progress-photos")\nasync def get_progress_photos(current_user: dict = Depends(get_current_user)):\n    photos = await db.progress_photos.find({"user_id": current_user["id"]}).sort("created_at", -1).to_list(100)\n    return serialize_mongo_docs(photos)'),
    
    # Sessions endpoints
    (r'@api_router\.get\("/user/sessions"\)\s*\nasync def get_my_sessions\(current_user: dict = Depends\(get_current_user\)\):\s*\n    sessions = await db\.sessions\.find\({"user_id": current_user\["id"\]}\)\.sort\("date", -1\)\.to_list\(100\)\s*\n    return sessions',
     '@api_router.get("/user/sessions")\nasync def get_my_sessions(current_user: dict = Depends(get_current_user)):\n    sessions = await db.sessions.find({"user_id": current_user["id"]}).sort("date", -1).to_list(100)\n    return serialize_mongo_docs(sessions)'),
    
    # Supplement logs today
    (r'logs = await db\.supplement_logs\.find\({\s*"user_id": current_user\["id"\],\s*"created_at": {"\$gte": today}\s*}\)\.to_list\(100\)',
     'logs = await db.supplement_logs.find({\n        "user_id": current_user["id"],\n        "created_at": {"$gte": today}\n    }).to_list(100)\n    logs = serialize_mongo_docs(logs)'),
    
    (r'assignments = await db\.supplement_assignments\.find\({\s*"user_id": current_user\["id"\],\s*"active": True\s*}\)\.to_list\(100\)',
     'assignments = await db.supplement_assignments.find({\n        "user_id": current_user["id"],\n        "active": True\n    }).to_list(100)\n    assignments = serialize_mongo_docs(assignments)'),
]

# Apply fixes
for pattern, replacement in fixes:
    content = re.sub(pattern, replacement, content, flags=re.MULTILINE | re.DOTALL)

# Write back the fixed content
with open('/app/backend/server.py', 'w') as f:
    f.write(content)

print("Fixed ObjectId serialization issues in server.py")