# Copyright (c) 2023 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
from decimal import Decimal
from enum import StrEnum
from functools import wraps
from types import MappingProxyType
from typing import Callable, Mapping, NamedTuple, Optional, TypeVar

from czc.type_util import deserializable, isoparse, serializable
from dateutil import parser
from toolz import keyfilter
from toolz.functoolz import curry
from voluptuous import (
    ALLOW_EXTRA,
    PREVENT_EXTRA,
    REMOVE_EXTRA,
    All,
    Any,
    Coerce,
    Datetime,
    DefaultTo,
    Maybe,
    Number,
    Range,
    Required,
    Schema,
    SetTo,
    message,
)
from voluptuous import (
    Optional as OptionalParam,
)

T = TypeVar("T")


class CloudProvider(StrEnum):
    datadog = "datadog"
    newrelic = "newrelic"
    databricks = "databricks"

    @property
    def name(self):
        return {
            CloudProvider.datadog: "Datadog",
            CloudProvider.newrelic: "New Relic",
            CloudProvider.databricks: "Databricks",
        }.get(self)


@deserializable(
    schema=Schema(
        {
            Required("connection_id"): str,
            Required("organization_id"): str,
            Required("name"): str,
            Required("cloud_provider"): Coerce(CloudProvider),
            Required("created"): Maybe(isoparse),
            Required("last_updated"): Maybe(isoparse),
            "provider_config": All(dict, DefaultTo({})),
            "derived_attributes": All(dict, DefaultTo({})),
            "is_enabled": bool,
            "is_deleted": bool,
            "capabilities": All(dict, DefaultTo({})),
            "billing_account_id": Maybe(str),
        },
        extra=REMOVE_EXTRA,
    )
)
class AnyCostConnectionInfo(NamedTuple):
    connection_id: str
    organization_id: str
    name: str
    cloud_provider: CloudProvider
    created: datetime
    last_updated: datetime
    provider_config: Mapping = MappingProxyType({})
    derived_attributes: Mapping = MappingProxyType({})
    is_enabled: bool = True
    is_deleted: bool = False
    capabilities: Mapping = MappingProxyType({})
    billing_account_id: Optional[str] = None

    @staticmethod
    def _is_public_field(field_name: str) -> bool:
        return field_name != "is_deleted"

    def to_public(self) -> dict:
        return keyfilter(self._is_public_field, serializable(self))


@curry
def custom_validator(msg: str, predicate: Callable[[T], bool]) -> Callable[[T], T]:
    @message(msg)
    @wraps(predicate)
    def _validate(obj: T) -> T:
        try:
            if predicate(obj):
                return obj
        except (TypeError, ValueError, KeyError):
            pass
        raise ValueError

    return _validate()


TIERED_RATE = Schema(
    {
        Required("threshold"): All(Coerce(int), Range(min=0, min_included=True)),
        Required("rate"): Number(yield_decimal=True),
    },
    extra=PREVENT_EXTRA,
)


@deserializable(schema=TIERED_RATE)
class TieredRate(NamedTuple):
    threshold: int
    rate: float


@custom_validator("Tier rates must specify a tier with a threshold of zero.")
def contains_zero_tier(tiers: list[TieredRate]) -> bool:
    return any(tier.threshold == 0 for tier in tiers)


@custom_validator("Tier rates must have only one occurrence of each threshold.")
def contains_unique_tiers(tiers: list[TieredRate]) -> bool:
    return len({t.threshold for t in tiers}) == len(tiers)


TIERED_RATE_LIST = All([Coerce(TieredRate)], contains_zero_tier, contains_unique_tiers)

NEW_RELIC_RATE_INFORMATION = Schema(
    {
        Required("cost_per_gb"): TIERED_RATE_LIST,
        Required("cost_per_basic_user"): TIERED_RATE_LIST,
        Required("cost_per_core_user"): TIERED_RATE_LIST,
        Required("cost_per_full_platform_user"): TIERED_RATE_LIST,
        Required("cost_per_synthetic_check"): TIERED_RATE_LIST,
        OptionalParam("subscription_term_start_date", default="0001-01-01T00:00:01.000Z"): Datetime(),
    },
    extra=PREVENT_EXTRA,
)


@deserializable(schema=NEW_RELIC_RATE_INFORMATION)
class NewRelicRateInformation(NamedTuple):
    cost_per_gb: list[TieredRate]
    cost_per_basic_user: list[TieredRate]
    cost_per_core_user: list[TieredRate]
    cost_per_full_platform_user: list[TieredRate]
    cost_per_synthetic_check: list[TieredRate]
    subscription_term_start_date: str

    @property
    def term_start_datetime(self) -> datetime:
        return parser.isoparse(self.subscription_term_start_date)


DATABRICKS_PRICING_SCHEMA = Schema(
    {
        Required("default"): Coerce(Decimal),
    },
    extra=ALLOW_EXTRA,
)

DATABRICKS_RATE_INFORMATION = Schema(
    {
        Required("price_start_time"): Maybe(isoparse),
        Required("price_end_time"): Any(Maybe(isoparse), SetTo(datetime.max)),
        Required("sku_name"): str,
        Required("cloud"): str,
        Required("currency_code"): str,
        Required("usage_unit"): str,
        Required("pricing"): DATABRICKS_PRICING_SCHEMA,
    },
    extra=PREVENT_EXTRA,
)


@deserializable(schema=DATABRICKS_RATE_INFORMATION)
class DatabricksRateInformation(NamedTuple):
    price_start_time: datetime
    price_end_time: datetime
    sku_name: str
    cloud: str
    currency_code: str
    usage_unit: str
    pricing: dict


DATABRICKS_RATE_INFORMATION_LIST = All([Coerce(DatabricksRateInformation)])
