import asyncio
import random
from datetime import datetime, timedelta, timezone
from sqlmodel import SQLModel
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession

from app.database import engine, init_db
from app.models import SleepSession

async def populate_data():
    # Ensure tables are created
    await init_db()

    # Today's date (as per environment) is April 24, 2026
    # We use UTC to be consistent with the model's default_factory
    now = datetime(2026, 4, 24, tzinfo=timezone.utc)
    
    # We want 2 weeks of data (14 days)
    # To ensure the oldest date gets the lowest ID, we iterate from the oldest to the newest.
    # i=13 is 13 days ago (oldest). i=0 is today.
    for day_offset in range(13, -1, -1):
        target_date = now - timedelta(days=day_offset)
        
        # Normalize target date to midnight UTC for consistency in base calculations
        base_date = datetime(target_date.year, target_date.month, target_date.day, tzinfo=timezone.utc)

        async with AsyncSession(engine) as session:
            # 1. Morning Nap (around 10:00 - 11:30)
            morning_start = base_date + timedelta(hours=10, minutes=0) + timedelta(minutes=random.randint(-15, 30))
            morning_end = base_date + timedelta(hours=11, minutes=30) + timedelta(minutes=random.randint(-30, 30))
            if morning_end <= morning_start:
                morning_end = morning_start + timedelta(hours=1, minutes=random.randint(0, 30))
            
            # 2. Afternoon Nap (around 15:00 - 16:30)
            afternoon_start = base_date + timedelta(hours=15, minutes=0) + timedelta(minutes=random.randint(-30, 30))
            afternoon_end = base_date + timedelta(hours=16, minutes=30) + timedelta(minutes=random.randint(-30, 30))
            if afternoon_end <= afternoon_start:
                afternoon_end = afternoon_start + timedelta(hours=1, minutes=random.randint(0, 30))

            # Add the naps
            session.add(SleepSession(start_time=morning_start, end_time=morning_end, notes="Morning nap"))
            session.add(SleepSession(start_time=afternoon_start, end_time=afternoon_end, notes="Afternoon nap"))

            # 3. Night Sleep (around 20:00 to 06:00 the following day)
            # ONLY add if it's not today (day_offset 0), to avoid adding future sleep.
            if day_offset > 0:
                night_start = base_date + timedelta(hours=20, minutes=0) + timedelta(minutes=random.randint(-30, 30))
                night_end = (base_date + timedelta(days=1)).replace(hour=6, minute=0) + timedelta(minutes=random.randint(-30, 30))
                session.add(SleepSession(start_time=night_start, end_time=night_end, notes="Night sleep"))
            
            await session.commit()
            print(f"Added data for {base_date.strftime('%Y-%m-%d')}")

    print("Successfully populated the database with 2 weeks of mock sleep data.")

if __name__ == "__main__":
    asyncio.run(populate_data())