# 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 os
from typing import Generic, Iterable, Type, TypeVar

import czc.internal_api.client as client
from czc import type_util
from czc.aws_environment import core_account_id, environment

from .types import (
    AnyCostConnectionInfo,
    DatabricksRateInformation,
    NewRelicRateInformation,
)

T = TypeVar("T")


def _iapi(api: str):
    return client.invoked(
        namespace=lambda _: os.environ["NAMESPACE"],
        feature="connections-anycost",
        account_id=lambda: core_account_id(environment(os.environ["NAMESPACE"])),
        api=api,
    )


@_iapi("get-all-connections")
def get_all_connections() -> Iterable[AnyCostConnectionInfo]:
    pass


@_iapi("get-organization-all-connections")
def get_organization_all_connections(
    organization_id: str,
) -> Iterable[AnyCostConnectionInfo]:
    pass


@_iapi("get-organization-provider-all-connections")
def get_organization_provider_all_connections(
    organization_id: str, provider: str
) -> Iterable[AnyCostConnectionInfo]:
    pass


@_iapi("get-connection")
def get_connection(organization_id: str, connection_id: str) -> AnyCostConnectionInfo:
    pass


@_iapi("get-rate-information")
def _get_rate_information(*, organization_id: str, connection_id: str):
    pass

@_iapi("update-org-connection")
def update_organization_connection(organization_id: str, connection_id: str, update_attributes: dict):
    pass


class _AdaptorClient(Generic[T]):
    def __init__(self, metadata_type: Type[T]):
        self._metadata_type = metadata_type

    def get_rate_information(
        self, organization_id: str, connection_id: str, **kwargs
    ) -> Iterable[T]:
        rate_info_list = _get_rate_information(
            organization_id=organization_id, connection_id=connection_id, **kwargs
        )
        for rate_info in rate_info_list:
            yield type_util.coerce(self._metadata_type, rate_info)


newrelic_client = _AdaptorClient(NewRelicRateInformation)
databricks_client = _AdaptorClient(DatabricksRateInformation)
