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

import voluptuous
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(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(voluptuous.Schema({
    voluptuous.Required('start'): iso_string_to_utc_datetime,
    voluptuous.Optional('end'): iso_string_to_utc_datetime,
    voluptuous.Optional('_end_inclusive'): bool
}, extra=voluptuous.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)
        else:
            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


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


@deserializable()
@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]
