# Copyright (c) 2024 CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com
from datetime import datetime, timedelta
from math import floor

import hypothesis.strategies as st
import pytest
from dateutil.tz import tzutc
from freezegun import freeze_time
from hypothesis import HealthCheck, given, settings

from czc.date_range import DateRange, GranularityType, Timespan, TimespanUnit

settings.register_profile(__name__, suppress_health_check=[HealthCheck.function_scoped_fixture])
settings.load_profile(__name__)


@pytest.mark.soc1
@pytest.mark.unit
@pytest.mark.parametrize(
    'start,end,expected_days,expected_seconds',
    [
        (datetime(2019, 8, 14, 4, 23, 6), datetime(2020, 9, 20, 1, 7, 0), 402, 74634),
        (datetime(2019, 8, 14, 4, 23, 6), datetime(2019, 8, 14, 6, 7, 0), 0, 6234),
        (datetime(2024, 3, 17, 5, 5, 0), datetime(2024, 4, 2, 12, 23, 15), 16, 26295),
    ],
)
@pytest.mark.parametrize('end_inclusive', [True, False])
def test_date_range_should_report_expected_duration(start, end, expected_days, expected_seconds, end_inclusive):
    date_range = DateRange(start=start, end=end, _end_inclusive=end_inclusive)
    assert date_range.duration == timedelta(
        days=expected_days, seconds=expected_seconds, microseconds=(1 if end_inclusive else 0)
    )
    assert date_range.duration_in_hours == expected_days * 24 + expected_seconds / 3600


@pytest.mark.soc1
@pytest.mark.unit
@pytest.mark.parametrize(
    'start,end,granularity,expected_start,expected_end',
    [
        (
            datetime(2019, 8, 14, 4, 23, 6),
            datetime(2020, 9, 20, 1, 7, 0),
            GranularityType.monthly,
            datetime(2018, 6, 14, 4, 23, 6),
            datetime(2019, 7, 20, 1, 7),
        ),
        (
            datetime(2019, 8, 14, 4, 23, 6),
            datetime(2020, 9, 20, 1, 7, 0),
            GranularityType.daily,
            datetime(2018, 7, 7, 7, 39, 12),
            datetime(2019, 8, 14, 4, 23, 6),
        ),
        (
            datetime(2024, 3, 17, 5, 5, 0),
            datetime(2024, 4, 2, 12, 23, 15),
            GranularityType.hourly,
            datetime(2024, 2, 29, 21, 46, 45),
            datetime(2024, 3, 17, 5, 5, 0),
        ),
    ],
)
def test_date_range_should_report_expected_previous_period(start, end, granularity, expected_start, expected_end):
    previous_date_range = DateRange(start=start, end=end).get_previous_period_for_granularity(granularity)
    assert previous_date_range == DateRange(start=expected_start, end=expected_end)


@pytest.mark.soc1
@pytest.mark.unit
@pytest.mark.parametrize(
    'start,end,granularity,expected_periods',
    [
        (
            datetime(2019, 8, 14, 4, 23, 6),
            datetime(2020, 9, 20, 1, 7, 0),
            GranularityType.yearly,
            [datetime(2019, 8, 14, 4, 23, 6), datetime(2020, 8, 14, 4, 23, 6)],
        ),
        (
            datetime(2023, 8, 14, 4, 23, 6),
            datetime(2023, 8, 14, 7, 7, 0),
            GranularityType.hourly,
            [datetime(2023, 8, 14, 4, 23, 6), datetime(2023, 8, 14, 5, 23, 6), datetime(2023, 8, 14, 6, 23, 6)],
        ),
        (
            datetime(2024, 3, 17, 5, 5, 0),
            datetime(2024, 4, 2, 12, 23, 15),
            GranularityType.monthly,
            [datetime(2024, 3, 17, 5, 5, 0)],
        ),
        (
            datetime(2024, 3, 17, 5, 5, 0),
            datetime(2024, 4, 2, 12, 23, 15),
            GranularityType.weekly,
            [datetime(2024, 3, 17, 5, 5, 0), datetime(2024, 3, 24, 5, 5), datetime(2024, 3, 31, 5, 5)],
        ),
    ],
)
def test_date_range_should_report_expected_time_periods(start, end, granularity, expected_periods):
    periods = DateRange(start=start, end=end).time_periods_for(granularity)
    assert list(periods) == expected_periods


@pytest.mark.soc1
@pytest.mark.unit
@pytest.mark.parametrize(
    'start_date,end_date,granularity,end_inclusive,expected_aligned_start_date,expected_aligned_end_date',
    [
        (
            datetime(2019, 8, 14, 4, 23, 6),
            datetime(2020, 9, 20, 1, 7, 0),
            GranularityType.hourly,
            False,
            datetime(2019, 8, 14, 4, 0, 0),
            datetime(2020, 9, 20, 1, 0, 0),
        ),
        (
            datetime(2019, 8, 14, 4, 23, 6),
            datetime(2020, 9, 20, 1, 7, 0),
            GranularityType.daily,
            False,
            datetime(2019, 8, 14, 0, 0, 0),
            datetime(2020, 9, 20, 0, 0, 0),
        ),
        (
            datetime(2019, 8, 14, 4, 23, 6),
            datetime(2020, 9, 20, 1, 7, 0),
            GranularityType.weekly,
            False,
            datetime(2019, 8, 12, 0, 0, 0),
            datetime(2020, 9, 14, 0, 0, 0),
        ),
        (
            datetime(2020, 8, 14, 4, 23, 6),
            datetime(2020, 9, 20, 1, 7, 0),
            GranularityType.thirty_days,
            False,
            datetime(2020, 7, 22, 0, 0),
            datetime(2020, 9, 20, 0, 0, 0),
        ),
        (
            datetime(2019, 8, 14, 4, 23, 6),
            datetime(2020, 9, 20, 1, 7, 0),
            GranularityType.thirty_days,
            False,
            datetime(2019, 7, 28, 0, 0),
            datetime(2020, 9, 20, 0, 0, 0),
        ),
        (
            datetime(2019, 8, 14, 4, 23, 6),
            datetime(2020, 9, 20, 1, 7, 0),
            GranularityType.hourly,
            True,
            datetime(2019, 8, 14, 4, 0, 0),
            datetime(2020, 9, 20, 1, 59, 59, 999999),
        ),
        (
            datetime(2019, 8, 14, 4, 23, 6),
            datetime(2020, 9, 20, 1, 7, 0),
            GranularityType.daily,
            True,
            datetime(2019, 8, 14, 0, 0, 0),
            datetime(2020, 9, 20, 23, 59, 59, 999999),
        ),
        (
            datetime(2019, 8, 14, 4, 23, 6),
            datetime(2020, 9, 20, 1, 7, 0),
            GranularityType.weekly,
            True,
            datetime(2019, 8, 12, 0, 0, 0),
            datetime(2020, 9, 20, 23, 59, 59, 999999),
        ),
        (
            datetime(2020, 8, 14, 4, 23, 6),
            datetime(2020, 9, 20, 1, 7, 0),
            GranularityType.thirty_days,
            True,
            datetime(2020, 7, 23, 0, 0),
            datetime(2020, 9, 20, 23, 59, 59, 999999),
        ),
        (
            datetime(2019, 8, 14, 4, 23, 6),
            datetime(2020, 9, 20, 1, 7, 0),
            GranularityType.thirty_days,
            True,
            datetime(2019, 7, 29, 0, 0),
            datetime(2020, 9, 20, 23, 59, 59, 999999),
        ),
    ],
)
def test_aligned_to_granularity_should_align_dates_correctly(
    start_date, end_date, granularity, end_inclusive, expected_aligned_start_date, expected_aligned_end_date
):
    dr = DateRange(start=start_date, end=end_date)
    new_dr = dr.aligned_to_granularity(granularity, end_inclusive=end_inclusive)

    assert new_dr.start == expected_aligned_start_date
    assert new_dr.end == expected_aligned_end_date


def _assert_duration_is_integer_granularity_units(date_range: DateRange, granularity: GranularityType):
    if granularity == GranularityType.yearly:
        exclusive_end = date_range.start + date_range.duration
        adjusted_end = exclusive_end.replace(year=date_range.start.year)
        assert date_range.start == adjusted_end
    elif granularity == GranularityType.monthly:
        exclusive_end = date_range.start + date_range.duration
        adjusted_end = exclusive_end.replace(year=date_range.start.year, month=date_range.start.month)
        assert date_range.start == adjusted_end
    else:
        granularity_units = date_range.duration / granularity.duration_from(date_range.start)
        assert granularity_units == floor(granularity_units)


def _assert_aligned_start_rounded_down(date_range: DateRange, orig_start: datetime, granularity: GranularityType):
    delta = orig_start - date_range.start
    assert delta >= timedelta()
    assert delta < granularity.duration_from(date_range.start)


@pytest.mark.soc1
@pytest.mark.unit
@settings(max_examples=1000)
@given(
    start=st.datetimes(min_value=datetime(2017, 1, 1), max_value=datetime(2025, 1, 1)),
    duration=st.timedeltas(min_value=timedelta(), max_value=timedelta(weeks=200)),
    granularity=st.sampled_from(GranularityType),
    end_inclusive=st.booleans(),
)
def test_aligned_to_granularity_should_produce_valid_aligned_daterange(start, duration, granularity, end_inclusive):
    aligned_daterange = DateRange(start=start, end=start + duration).aligned_to_granularity(granularity, end_inclusive)
    _assert_duration_is_integer_granularity_units(aligned_daterange, granularity)
    _assert_aligned_start_rounded_down(aligned_daterange, start, granularity)


@pytest.mark.soc1
@pytest.mark.unit
@pytest.mark.parametrize(
    'unit,value,expected_start',
    [
        (TimespanUnit.years, 1, datetime(2023, 1, 1, tzinfo=tzutc())),
        (TimespanUnit.months, 2, datetime(2024, 2, 1, tzinfo=tzutc())),
        (TimespanUnit.weeks, 5, datetime(2024, 2, 26, tzinfo=tzutc())),
        (TimespanUnit.days, 17, datetime(2024, 3, 16, tzinfo=tzutc())),
        (TimespanUnit.hours, 4, datetime(2024, 4, 2, 3, tzinfo=tzutc())),
        (TimespanUnit.hours, 48, datetime(2024, 3, 31, 7, tzinfo=tzutc())),
    ],
)
def test_timespan_should_produce_expected_date_range(unit, value, expected_start):
    timespan = Timespan(unit=unit, value=value)
    now = datetime(2024, 4, 2, 7, 13, 47, tzinfo=tzutc())
    with freeze_time(now):
        date_range = timespan.to_date_range()
    assert date_range == DateRange(start=expected_start, end=now)
