# 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 collections.abc import Iterable
from dataclasses import dataclass, field, replace
from datetime import datetime, timedelta
from enum import Enum
from math import ceil

import voluptuous as v
from dateutil.relativedelta import relativedelta
from dateutil.tz import tzutc

from czc.date import iso_string_to_utc_datetime
from czc.type_util import deserializable


def _utc_now():
    return datetime.now(tz=tzutc())


def _monday(dt: datetime) -> datetime:
    return (dt - timedelta(days=dt.weekday())).replace(hour=0, minute=0, second=0, microsecond=0)


class GranularityType(str, Enum):
    hourly = 'hourly'
    daily = 'daily'
    weekly = 'weekly'
    thirty_days = '30days'
    monthly = 'monthly'
    yearly = 'yearly'

    def __str__(self):
        return self.value

    @property
    def delta(self) -> relativedelta:
        return {
            GranularityType.hourly: relativedelta(hours=1),
            GranularityType.daily: relativedelta(days=1),
            GranularityType.weekly: relativedelta(days=7),
            GranularityType.thirty_days: relativedelta(days=30),
            GranularityType.monthly: relativedelta(months=1),
            GranularityType.yearly: relativedelta(years=1),
        }[self]

    def duration_from(self, start: datetime) -> timedelta:
        return (start + self.delta) - start


@deserializable(
    v.Schema(
        {
            v.Required('start'): iso_string_to_utc_datetime,
            v.Optional('end'): v.All(v.Maybe(iso_string_to_utc_datetime), v.DefaultTo(_utc_now)),
            v.Optional('_end_inclusive'): bool,
        },
        extra=v.REMOVE_EXTRA,
    )
)
@dataclass(frozen=True)
class DateRange:
    start: datetime
    end: datetime = field(default_factory=_utc_now)
    _end_inclusive: bool = False

    @property
    def duration(self) -> timedelta:
        return (self.end - self.start) + (timedelta(microseconds=1) if self._end_inclusive else timedelta())

    @property
    def duration_in_hours(self) -> float:
        return self.duration.days * 24 + self.duration.seconds / (60 * 60)

    def _get_prior_period(self, offset):
        return replace(self, start=self.start - offset, end=self.end - offset)

    def get_previous_period_for_granularity(self, granularity: GranularityType) -> 'DateRange':
        if granularity in (GranularityType.monthly, GranularityType.yearly):
            count = len(list(self.time_periods_for(granularity)))
            return self._get_prior_period(granularity.delta * count)
        return self._get_prior_period(self.duration)

    def time_periods_for(self, granularity: GranularityType) -> Iterable[datetime]:
        curr = self.start
        while curr < self.start + self.duration:
            yield curr
            curr += granularity.delta

    def aligned_to_granularity(self, granularity: GranularityType, end_inclusive: bool = False) -> 'DateRange':
        def adjusted_end(dt: datetime) -> datetime:
            return dt + _get_alignment_boundary(granularity).delta - timedelta(microseconds=1) if end_inclusive else dt

        new_daterange = DateRange(
            start=_align_datetime_to_boundary(self.start, granularity),
            end=adjusted_end(_align_datetime_to_boundary(self.end, granularity)),
            _end_inclusive=end_inclusive,
        )
        if granularity != _get_alignment_boundary(granularity):
            alignment_timespan = granularity.duration_from(new_daterange.start + new_daterange.duration).total_seconds()
            aligned_duration = ceil(new_daterange.duration.total_seconds() / alignment_timespan) * alignment_timespan
            new_start = new_daterange.start + new_daterange.duration - timedelta(seconds=aligned_duration)
            return replace(new_daterange, start=new_start)
        return new_daterange

    @classmethod
    def from_month(cls, month_str: str) -> 'DateRange':
        year, month = month_str.split('-')
        month_start = datetime(year=int(year), month=int(month), day=1, tzinfo=tzutc())
        next_month_start = month_start + relativedelta(months=1)

        return DateRange(start=month_start, end=next_month_start)


class TimespanUnit(Enum):
    hours = 'hours'
    days = 'days'
    weeks = 'weeks'
    months = 'months'
    years = 'years'


@deserializable(
    schema=v.Schema(
        {
            v.Required('unit'): v.Coerce(TimespanUnit),
            v.Required('value'): v.All(v.Number(scale=0), v.Coerce(int), v.Range(min=1)),
        },
        extra=v.REMOVE_EXTRA,
    )
)
@dataclass(frozen=True)
class Timespan:
    unit: TimespanUnit
    value: int

    def to_date_range(self) -> DateRange:
        granularity = {
            TimespanUnit.hours: GranularityType.hourly,
            TimespanUnit.days: GranularityType.daily,
            TimespanUnit.weeks: GranularityType.weekly,
            TimespanUnit.months: GranularityType.monthly,
            TimespanUnit.years: GranularityType.yearly,
        }[self.unit]
        aligned_now = _align_datetime_to_boundary(_utc_now(), granularity)
        return DateRange(start=aligned_now - self.value * granularity.delta)


def _get_alignment_boundary(granularity: GranularityType):
    return {GranularityType.thirty_days: GranularityType.daily}.get(granularity, granularity)


def _align_datetime_to_boundary(date_time: datetime, granularity: GranularityType) -> datetime:
    hourly_aligned = date_time.replace(minute=0, second=0, microsecond=0)
    alignment_boundary = _get_alignment_boundary(granularity)
    return {
        GranularityType.hourly: hourly_aligned,
        GranularityType.daily: hourly_aligned.replace(hour=0),
        GranularityType.weekly: _monday(hourly_aligned),
        GranularityType.monthly: hourly_aligned.replace(day=1, hour=0),
        GranularityType.yearly: hourly_aligned.replace(month=1, day=1, hour=0),
    }[alignment_boundary]
