"""Date calculation utilities for leave planning.

Pure functions with no I/O or state. Used by planning and analytics services.
"""

import datetime
from typing import Union


def add_months(date: datetime.date, n: int) -> datetime.date:
    """Add or subtract months from a date, handling month-end edge cases.

    Examples:
      add_months(datetime.date(2026, 1, 31), 1) -> 2026-02-28 (or 29 in leap year)
      add_months(datetime.date(2026, 3, 31), -1) -> 2026-02-28
    """
    month = date.month + n
    year = date.year

    while month > 12:
        month -= 12
        year += 1

    while month < 1:
        month += 12
        year -= 1

    # Handle month-end: if original day doesn't exist in target month,
    # use last day of that month
    day = min(date.day, _days_in_month(year, month))

    return datetime.date(year, month, day)


def _days_in_month(year: int, month: int) -> int:
    """Return number of days in the given month."""
    if month == 2:
        return 29 if _is_leap_year(year) else 28
    if month in (4, 6, 9, 11):
        return 30
    return 31


def _is_leap_year(year: int) -> bool:
    """Check if a year is a leap year."""
    return year % 4 == 0 and (year % 100 != 0 or year % 400 == 0)


def age_in_months(birth_date: datetime.date, on: datetime.date) -> int:
    """Calculate exact age in months between two dates.

    If on < birth_date, returns negative value (person not yet born).
    """
    months = (on.year - birth_date.year) * 12 + (on.month - birth_date.month)

    # Adjust if birthday hasn't occurred this month yet
    if on.day < birth_date.day:
        months -= 1

    return months


def fourth_birthday(birth_date: datetime.date) -> datetime.date:
    """Return the date when a child turns 4 years old."""
    return add_months(birth_date, 48)