from datetime import datetime
from zoneinfo import ZoneInfo

from timedeltacal import timedeltacal

# Some basic tests. These serve as a specification. All the tests use timezone
# aware datetimes in a timezone with DST, since this is the most complex case.
# Naive datetimes and dates should work analogously.

def test_add_same_day():
    t0 = datetime(2022, 4, 17,  10, 37, 21,  tzinfo=ZoneInfo("Europe/Vienna"))
    d = timedeltacal(hours=2, minutes=34, seconds=56, milliseconds=789)
    t1 = t0 + d
    assert t1.year == t0.year
    assert t1.month == t0.month
    assert t1.day == t0.day
    assert t1.hour == 13
    assert t1.minute == 12
    assert t1.second == 17
    assert t1.microsecond == 789000
    assert t1.tzinfo is t0.tzinfo


def test_add_next_day():
    t0 = datetime(2022, 4, 17,  10, 37, 21,  tzinfo=ZoneInfo("Europe/Vienna"))
    d = timedeltacal(hours=13, minutes=34, seconds=56, milliseconds=789)
    t1 = t0 + d
    assert t1.year == t0.year
    assert t1.month == t0.month
    assert t1.day == 18
    assert t1.hour == 0
    assert t1.minute == 12
    assert t1.second == 17
    assert t1.microsecond == 789000
    assert t1.tzinfo is t0.tzinfo

def test_add_days():
    t0 = datetime(2022, 4, 17,  10, 37, 21,  tzinfo=ZoneInfo("Europe/Vienna"))
    d = timedeltacal(days=1)
    t1 = t0 + d
    assert t1.year == t0.year
    assert t1.month == t0.month
    assert t1.day == 18
    assert t1.hour == t0.hour
    assert t1.minute == t0.minute
    assert t1.second == t0.second
    assert t1.microsecond == t0.microsecond
    assert t1.tzinfo is t0.tzinfo

# Add days -> next month
def test_add_days_to_next_month():
    t0 = datetime(2022, 4, 17,  10, 37, 21,  tzinfo=ZoneInfo("Europe/Vienna"))
    d = timedeltacal(days=20)
    t1 = t0 + d
    assert t1.year == t0.year
    assert t1.month == 5
    assert t1.day == 7
    assert t1.hour == t0.hour
    assert t1.minute == t0.minute
    assert t1.second == t0.second
    assert t1.microsecond == t0.microsecond
    assert t1.tzinfo is t0.tzinfo

# Add months
def test_add_months():
    t0 = datetime(2022, 4, 17,  10, 37, 21,  tzinfo=ZoneInfo("Europe/Vienna"))
    d = timedeltacal(months=1)
    t1 = t0 + d
    assert t1.year == t0.year
    assert t1.month == 5
    assert t1.day == 17
    assert t1.hour == t0.hour
    assert t1.minute == t0.minute
    assert t1.second == t0.second
    assert t1.microsecond == t0.microsecond
    assert t1.tzinfo is t0.tzinfo

# Add years
def test_add_years():
    t0 = datetime(2022, 4, 17,  10, 37, 21,  tzinfo=ZoneInfo("Europe/Vienna"))
    d = timedeltacal(years=1)
    t1 = t0 + d
    assert t1.year == 2023
    assert t1.month == 4
    assert t1.day == 17
    assert t1.hour == t0.hour
    assert t1.minute == t0.minute
    assert t1.second == t0.second
    assert t1.microsecond == t0.microsecond
    assert t1.tzinfo is t0.tzinfo


# Same tests, but crossing a DST boundary
def test_add_same_day_dst():
    t0 = datetime(2022, 3, 27,  1, 37, 21,  tzinfo=ZoneInfo("Europe/Vienna"))
    d = timedeltacal(hours=2, minutes=34, seconds=56, milliseconds=789)
    t1 = t0 + d
    assert t1.year == t0.year
    assert t1.month == t0.month
    assert t1.day == t0.day
    assert t1.hour == 5
    assert t1.minute == 12
    assert t1.second == 17
    assert t1.microsecond == 789000
    assert t1.tzinfo is t0.tzinfo

def test_add_next_day_dst():
    t0 = datetime(2022, 3, 26,  10, 37, 21,  tzinfo=ZoneInfo("Europe/Vienna"))
    d = timedeltacal(hours=23, minutes=34, seconds=56, milliseconds=789)
    t1 = t0 + d
    assert t1.year == t0.year
    assert t1.month == t0.month
    assert t1.day == 27
    assert t1.hour == 11
    assert t1.minute == 12
    assert t1.second == 17
    assert t1.microsecond == 789000
    assert t1.tzinfo is t0.tzinfo

def test_add_days_dst():
    t0 = datetime(2022, 3, 26,  10, 37, 21,  tzinfo=ZoneInfo("Europe/Vienna"))
    d = timedeltacal(days=1)
    t1 = t0 + d
    assert t1.year == t0.year
    assert t1.month == t0.month
    assert t1.day == 27
    assert t1.hour == t0.hour
    assert t1.minute == t0.minute
    assert t1.second == t0.second
    assert t1.microsecond == t0.microsecond
    assert t1.tzinfo is t0.tzinfo

# Add days -> next month
def test_add_days_to_next_month_dst():
    t0 = datetime(2022, 3, 26,  10, 37, 21,  tzinfo=ZoneInfo("Europe/Vienna"))
    d = timedeltacal(days=20)
    t1 = t0 + d
    assert t1.year == t0.year
    assert t1.month == 4
    assert t1.day == 15
    assert t1.hour == t0.hour
    assert t1.minute == t0.minute
    assert t1.second == t0.second
    assert t1.microsecond == t0.microsecond
    assert t1.tzinfo is t0.tzinfo

# Add months
def test_add_months_dst():
    t0 = datetime(2022, 3, 26,  10, 37, 21,  tzinfo=ZoneInfo("Europe/Vienna"))
    d = timedeltacal(months=1)
    t1 = t0 + d
    assert t1.year == t0.year
    assert t1.month == 4
    assert t1.day == t0.day
    assert t1.hour == t0.hour
    assert t1.minute == t0.minute
    assert t1.second == t0.second
    assert t1.microsecond == t0.microsecond
    assert t1.tzinfo is t0.tzinfo

# Add years
def test_add_years_dst():
    t0 = datetime(2022, 3, 26,  10, 37, 21,  tzinfo=ZoneInfo("Europe/Vienna"))
    d = timedeltacal(years=1)
    t1 = t0 + d
    assert t1.year == 2023
    assert t1.month == 3
    assert t1.day == 26
    assert t1.hour == t0.hour
    assert t1.minute == t0.minute
    assert t1.second == t0.second
    assert t1.microsecond == t0.microsecond
    assert t1.tzinfo is t0.tzinfo

# Month arithmetic is saturating: If the arithmetic result would result in a
# day of the month which doesn't exist, we choose the last day of the month
# instead.

def test_add_months_feb():
    t0 = datetime(2022, 1, 31, tzinfo=ZoneInfo("Europe/Vienna"))
    d = timedeltacal(months=1)
    t1 = t0 + d
    assert t1.year == t0.year
    assert t1.month == 2
    assert t1.day == 28
    assert t1.hour == t0.hour
    assert t1.minute == t0.minute
    assert t1.second == t0.second
    assert t1.microsecond == t0.microsecond
    assert t1.tzinfo is t0.tzinfo

def test_add_months_mar():
    t0 = datetime(2022, 1, 31, tzinfo=ZoneInfo("Europe/Vienna"))
    d = timedeltacal(months=2)
    t1 = t0 + d
    assert t1.year == t0.year
    assert t1.month == 3
    assert t1.day == 31
    assert t1.hour == t0.hour
    assert t1.minute == t0.minute
    assert t1.second == t0.second
    assert t1.microsecond == t0.microsecond
    assert t1.tzinfo is t0.tzinfo

def test_add_months_apr():
    t0 = datetime(2022, 1, 31, tzinfo=ZoneInfo("Europe/Vienna"))
    d = timedeltacal(months=3)
    t1 = t0 + d
    assert t1.year == t0.year
    assert t1.month == 4
    assert t1.day == 30
    assert t1.hour == t0.hour
    assert t1.minute == t0.minute
    assert t1.second == t0.second
    assert t1.microsecond == t0.microsecond
    assert t1.tzinfo is t0.tzinfo

def test_add_months_feb_leap():
    t0 = datetime(2022, 1, 31, tzinfo=ZoneInfo("Europe/Vienna"))
    d = timedeltacal(months=25)
    t1 = t0 + d
    assert t1.year == 2024
    assert t1.month == 2
    assert t1.day == 29
    assert t1.hour == t0.hour
    assert t1.minute == t0.minute
    assert t1.second == t0.second
    assert t1.microsecond == t0.microsecond
    assert t1.tzinfo is t0.tzinfo