from datetime import datetime, timezone
from typing import List, Optional, Union
from sqlmodel import select, func
from sqlmodel.ext.asyncio.session import AsyncSession
from app.models import SleepSession

def ensure_aware(dt: datetime) -> datetime:
    if dt.tzinfo is None:
        return dt.replace(tzinfo=timezone.utc)
    return dt

async def get_sessions(
    session: AsyncSession,
    from_date: Optional[datetime] = None,
    to_date: Optional[datetime] = None,
    limit: int = 200,
) -> tuple[List[SleepSession], int]:
    if to_date and to_date.hour == 0 and to_date.minute == 0 and to_date.second == 0 and to_date.microsecond == 0:
        to_date = to_date.replace(hour=23, minute=59, second=59, microsecond=999999)

    # If we have an aware to_date, but the DB is naive, 
    # we should make it naive for comparison to avoid errors in some drivers.
    query_to_date = to_date.replace(tzinfo=None) if to_date and to_date.tzinfo else to_date
    query_from_date = from_date.replace(tzinfo=None) if from_date and from_date.tzinfo else from_date

    count_statement = select(func.count()).select_from(SleepSession)
    if query_from_date:
        count_statement = count_statement.where(SleepSession.start_time >= query_from_date)
    if query_to_date:
        count_statement = count_statement.where(SleepSession.start_time <= query_to_date)
    count_result = await session.exec(count_statement)
    total = count_result.one() or 0

    statement = select(SleepSession).order_by(SleepSession.start_time.desc())
    if query_from_date:
        statement = statement.where(SleepSession.start_time >= query_from_date)
    if query_to_date:
        statement = statement.where(SleepSession.start_time <= query_to_date)
    statement = statement.limit(limit)

    result = await session.exec(statement)
    sessions = result.all()
    return list(sessions), total

async def get_active_session(session: AsyncSession) -> Optional[SleepSession]:
    statement = select(SleepSession).where(SleepSession.end_time == None)
    result = await session.exec(statement)
    return result.first()

async def create_sleep_session(session: AsyncSession, start_time: datetime) -> SleepSession:
    new_session = SleepSession(start_time=start_time)
    session.add(new_session)
    await session.commit()
    await session.refresh(new_session)
    return new_session

async def end_sleep_session(
    session: AsyncSession, session_id: int, end_time: datetime
) -> Union[SleepSession, bool, str, None]:
    statement = select(SleepSession).where(SleepSession.id == session_id)
    result = await session.exec(statement)
    sleep_session = result.first()

    if not sleep_session:
        return None  # Not found

    if sleep_session.end_time is not None:
        return False  # Already ended

    # Normalize to aware for comparison
    end_time_aware = end_time.replace(tzinfo=timezone.utc) if end_time.tzinfo is None else end_time
    start_time_aware = sleep_session.start_time.replace(tzinfo=timezone.utc) if sleep_session.start_time.tzinfo is None else sleep_session.start_time

    if end_time_aware < start_time_aware:
        return "invalid"  # End time before start time

    # To store back to naive DB, we strip tzinfo if it was aware
    sleep_session.end_time = end_time.replace(tzinfo=None) if end_time.tzinfo else end_time
    session.add(sleep_session)
    await session.commit()
    await session.refresh(sleep_session)
    return sleep_session

async def delete_session(session: AsyncSession, session_id: int) -> bool:
    statement = select(SleepSession).where(SleepSession.id == session_id)
    result = await session.exec(statement)
    sleep_session = result.first()

    if not sleep_session:
        return False

    await session.delete(sleep_session)
    await session.commit()
    return True

async def update_session(
    session: AsyncSession, session_id: int, update_data: dict
) -> Optional[SleepSession]:
    statement = select(SleepSession).where(SleepSession.id == session_id)
    result = await session.exec(statement)
    sleep_session = result.first()

    if not sleep_session:
        return None

    for key, value in update_data.items():
        if hasattr(sleep_session, key):
            setattr(sleep_session, key, value)

    session.add(sleep_session)
    await session.commit()
    await session.refresh(sleep_session)
    return sleep_session