# Copyright (c) 2022 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 typing import Type

import os
from datetime import datetime
from dateutil.tz import tzutc
from decimal import Decimal

from czc.date import get_current_timestamp


class TimestampContext():
    cached: bool
    return_type: Type
    reset_on_exit: bool

    def __init__(self, cached: bool, return_type: Type) -> None:
        self.cached = cached
        self.return_type = return_type
        self.reset_on_exit = False

    def __enter__(self):
        value = os.environ.get('TIMESTAMP') if self.cached else get_current_timestamp()

        if value is None:
            value = get_current_timestamp()
            self.reset_on_exit = True
            os.environ['TIMESTAMP'] = str(value)

        if issubclass(self.return_type, datetime):
            return datetime.fromtimestamp(float(value), tz=tzutc())
        elif issubclass(self.return_type, str):
            return datetime.fromtimestamp(float(value), tz=tzutc()).isoformat()

        return self.return_type(value)

    def __exit__(self, *args, **kwargs):
        if not self.reset_on_exit:
            return

        os.environ.pop('TIMESTAMP')


def timestamped(cached: bool = False, return_type: Type = Decimal):
    """ Decorator that will get a timestamp and add it to
    to kwargs as 'TIMESTAMP'. If enabled, it can get a timestamp from a cached
    source so things within a given handler can get the
    same timestamp

    Args:
        cached (bool): If True (default) then it will use a cached one (if available).
                       If False, then it will always generate a new timestamp.

    Returns:
        function (func): a function that is timestamped
    """

    def timestamped_handler(handler):
        def function_wrapper(*args, **kwargs):
            with TimestampContext(cached, return_type) as timestamp:
                kwargs['TIMESTAMP'] = timestamp
                return handler(*args, **kwargs)

        return function_wrapper

    return timestamped_handler
