# 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
import collections.abc
from datetime import datetime
from enum import StrEnum
from typing import Any, Optional, Union, get_args

import pydantic
from pydantic import ConfigDict, field_validator
from toolz import keyfilter

from czc.type_util import deserializable, serializable

from .base import CustomBaseModel
from .cloud_providers.databricks import DatabricksDerivedAttributes, DatabricksProviderConfig
from .cloud_providers.datadog import DatadogDerivedAttributes, DatadogProviderConfig
from .cloud_providers.datadog_metric import DatadogMetricDerivedAttributes, DatadogMetricProviderConfig
from .cloud_providers.mongodb import MongoDBDerivedAttributes, MongoDBProviderConfig
from .cloud_providers.new_relic import NewRelicDerivedAttributes, NewRelicProviderConfig


class EmptyModel(CustomBaseModel):
    model_config = ConfigDict(extra='forbid')


ProviderConfigType = Union[
    DatadogProviderConfig,
    NewRelicProviderConfig,
    DatabricksProviderConfig,
    DatadogMetricProviderConfig,
    MongoDBProviderConfig,
    EmptyModel,
]

DerivedAttributesType = Union[
    DatadogDerivedAttributes,
    NewRelicDerivedAttributes,
    DatabricksDerivedAttributes,
    DatadogMetricDerivedAttributes,
    MongoDBDerivedAttributes,
    EmptyModel,
]


class CloudProvider(StrEnum):
    datadog = 'datadog'
    newrelic = 'newrelic'
    databricks = 'databricks'
    mongodb = 'mongodb'
    datadog_metric = 'datadog-metric'

    @property
    def name(self):
        return {
            CloudProvider.datadog: 'Datadog',
            CloudProvider.newrelic: 'New Relic',
            CloudProvider.databricks: 'Databricks',
            CloudProvider.mongodb: 'MongoDB',
            CloudProvider.datadog_metric: 'Datadog Metric',
        }.get(self)

    @property
    def provider_config_type(self) -> type[Any]:
        return {
            CloudProvider.datadog: DatadogProviderConfig,
            CloudProvider.newrelic: NewRelicProviderConfig,
            CloudProvider.databricks: DatabricksProviderConfig,
            CloudProvider.mongodb: MongoDBProviderConfig,
            CloudProvider.datadog_metric: DatadogMetricProviderConfig,
        }.get(self, EmptyModel)

    @property
    def derived_attributes_type(self) -> type[Any]:
        return {
            CloudProvider.datadog: DatadogDerivedAttributes,
            CloudProvider.newrelic: NewRelicDerivedAttributes,
            CloudProvider.databricks: DatabricksDerivedAttributes,
            CloudProvider.mongodb: MongoDBDerivedAttributes,
            CloudProvider.datadog_metric: DatadogMetricDerivedAttributes,
        }.get(self, EmptyModel)


@deserializable()
class AnyCostConnectionInfo(CustomBaseModel):
    connection_id: str
    organization_id: str
    name: str
    cloud_provider: CloudProvider
    created: Optional[datetime] = None
    last_updated: Optional[datetime] = None
    provider_config: ProviderConfigType
    derived_attributes: DerivedAttributesType
    is_enabled: bool = True
    is_deleted: bool = False
    capabilities: dict = {}  # noqa: RUF012
    billing_account_id: Optional[str] = None

    @field_validator('provider_config', mode='before')
    @classmethod
    def validate_provider_config(cls, value: Any, info: pydantic.ValidationInfo) -> ProviderConfigType:
        cloud_provider = CloudProvider(info.data['cloud_provider'])
        if isinstance(value, get_args(ProviderConfigType)):
            if type(value) is not (cloud_provider.provider_config_type or type(EmptyModel)):
                raise ValueError(
                    f"The provided 'provided_config' ({value}) does not match the expected schema for model: {cloud_provider.provider_config_type.__name__}."
                )
            return value

        if isinstance(value, collections.abc.Mapping):
            return cloud_provider.provider_config_type(**value) if value else EmptyModel()
        raise ValueError(f'The provided provider config is not valid: {value}')

    @field_validator('derived_attributes', mode='before')
    @classmethod
    def validate_derived_attributes(cls, value: Any, info: pydantic.ValidationInfo) -> DerivedAttributesType:
        cloud_provider = CloudProvider(info.data['cloud_provider'])
        if isinstance(value, get_args(DerivedAttributesType)):
            if type(value) is not cloud_provider.derived_attributes_type or type(EmptyModel):
                raise ValueError(
                    f"The provided 'derived_attributes': {value} does not match the expected schema for model: {cloud_provider.derived_attributes_type.__name__}."
                )
            return value
        if isinstance(value, collections.abc.Mapping):
            return cloud_provider.derived_attributes_type(**value) if value else EmptyModel()
        raise ValueError(f"The provided 'derived_attributes' is not valid: {value}")

    @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))
