# coding=utf-8
# --------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for license information.
# Code generated by Microsoft (R) AutoRest Code Generator.
# Changes may cause incorrect behavior and will be lost if the code is regenerated.
# --------------------------------------------------------------------------

from copy import deepcopy
from typing import Any, TYPE_CHECKING
from typing_extensions import Self

from azure.core.pipeline import policies
from azure.core.rest import HttpRequest, HttpResponse
from azure.mgmt.core import ARMPipelineClient
from azure.mgmt.core.policies import ARMAutoResourceProviderRegistrationPolicy

from . import models as _models
from ._configuration import CosmosDBManagementClientConfiguration
from ._serialization import Deserializer, Serializer
from .operations import (
    CassandraClustersOperations,
    CassandraDataCentersOperations,
    CassandraResourcesOperations,
    ChaosFaultOperations,
    CollectionOperations,
    CollectionPartitionOperations,
    CollectionPartitionRegionOperations,
    CollectionRegionOperations,
    DataTransferJobsOperations,
    DatabaseAccountRegionOperations,
    DatabaseAccountsOperations,
    DatabaseOperations,
    GraphResourcesOperations,
    GremlinResourcesOperations,
    LocationsOperations,
    MongoDBResourcesOperations,
    NetworkSecurityPerimeterConfigurationsOperations,
    NotebookWorkspacesOperations,
    Operations,
    PartitionKeyRangeIdOperations,
    PartitionKeyRangeIdRegionOperations,
    PercentileOperations,
    PercentileSourceTargetOperations,
    PercentileTargetOperations,
    PrivateEndpointConnectionsOperations,
    PrivateLinkResourcesOperations,
    RestorableDatabaseAccountsOperations,
    RestorableGremlinDatabasesOperations,
    RestorableGremlinGraphsOperations,
    RestorableGremlinResourcesOperations,
    RestorableMongodbCollectionsOperations,
    RestorableMongodbDatabasesOperations,
    RestorableMongodbResourcesOperations,
    RestorableSqlContainersOperations,
    RestorableSqlDatabasesOperations,
    RestorableSqlResourcesOperations,
    RestorableTableResourcesOperations,
    RestorableTablesOperations,
    ServiceOperations,
    SqlResourcesOperations,
    TableResourcesOperations,
    ThroughputPoolAccountOperations,
    ThroughputPoolAccountsOperations,
    ThroughputPoolOperations,
    ThroughputPoolsOperations,
)

if TYPE_CHECKING:
    from azure.core.credentials import TokenCredential


class CosmosDBManagementClient:  # pylint: disable=too-many-instance-attributes
    """Azure Cosmos DB Chaos Fault REST API.

    :ivar chaos_fault: ChaosFaultOperations operations
    :vartype chaos_fault: azure.mgmt.cosmosdb.operations.ChaosFaultOperations
    :ivar database_accounts: DatabaseAccountsOperations operations
    :vartype database_accounts: azure.mgmt.cosmosdb.operations.DatabaseAccountsOperations
    :ivar operations: Operations operations
    :vartype operations: azure.mgmt.cosmosdb.operations.Operations
    :ivar database: DatabaseOperations operations
    :vartype database: azure.mgmt.cosmosdb.operations.DatabaseOperations
    :ivar collection: CollectionOperations operations
    :vartype collection: azure.mgmt.cosmosdb.operations.CollectionOperations
    :ivar collection_region: CollectionRegionOperations operations
    :vartype collection_region: azure.mgmt.cosmosdb.operations.CollectionRegionOperations
    :ivar database_account_region: DatabaseAccountRegionOperations operations
    :vartype database_account_region:
     azure.mgmt.cosmosdb.operations.DatabaseAccountRegionOperations
    :ivar percentile_source_target: PercentileSourceTargetOperations operations
    :vartype percentile_source_target:
     azure.mgmt.cosmosdb.operations.PercentileSourceTargetOperations
    :ivar percentile_target: PercentileTargetOperations operations
    :vartype percentile_target: azure.mgmt.cosmosdb.operations.PercentileTargetOperations
    :ivar percentile: PercentileOperations operations
    :vartype percentile: azure.mgmt.cosmosdb.operations.PercentileOperations
    :ivar collection_partition_region: CollectionPartitionRegionOperations operations
    :vartype collection_partition_region:
     azure.mgmt.cosmosdb.operations.CollectionPartitionRegionOperations
    :ivar collection_partition: CollectionPartitionOperations operations
    :vartype collection_partition: azure.mgmt.cosmosdb.operations.CollectionPartitionOperations
    :ivar partition_key_range_id: PartitionKeyRangeIdOperations operations
    :vartype partition_key_range_id: azure.mgmt.cosmosdb.operations.PartitionKeyRangeIdOperations
    :ivar partition_key_range_id_region: PartitionKeyRangeIdRegionOperations operations
    :vartype partition_key_range_id_region:
     azure.mgmt.cosmosdb.operations.PartitionKeyRangeIdRegionOperations
    :ivar graph_resources: GraphResourcesOperations operations
    :vartype graph_resources: azure.mgmt.cosmosdb.operations.GraphResourcesOperations
    :ivar sql_resources: SqlResourcesOperations operations
    :vartype sql_resources: azure.mgmt.cosmosdb.operations.SqlResourcesOperations
    :ivar mongo_db_resources: MongoDBResourcesOperations operations
    :vartype mongo_db_resources: azure.mgmt.cosmosdb.operations.MongoDBResourcesOperations
    :ivar table_resources: TableResourcesOperations operations
    :vartype table_resources: azure.mgmt.cosmosdb.operations.TableResourcesOperations
    :ivar cassandra_resources: CassandraResourcesOperations operations
    :vartype cassandra_resources: azure.mgmt.cosmosdb.operations.CassandraResourcesOperations
    :ivar gremlin_resources: GremlinResourcesOperations operations
    :vartype gremlin_resources: azure.mgmt.cosmosdb.operations.GremlinResourcesOperations
    :ivar locations: LocationsOperations operations
    :vartype locations: azure.mgmt.cosmosdb.operations.LocationsOperations
    :ivar data_transfer_jobs: DataTransferJobsOperations operations
    :vartype data_transfer_jobs: azure.mgmt.cosmosdb.operations.DataTransferJobsOperations
    :ivar cassandra_clusters: CassandraClustersOperations operations
    :vartype cassandra_clusters: azure.mgmt.cosmosdb.operations.CassandraClustersOperations
    :ivar cassandra_data_centers: CassandraDataCentersOperations operations
    :vartype cassandra_data_centers: azure.mgmt.cosmosdb.operations.CassandraDataCentersOperations
    :ivar network_security_perimeter_configurations:
     NetworkSecurityPerimeterConfigurationsOperations operations
    :vartype network_security_perimeter_configurations:
     azure.mgmt.cosmosdb.operations.NetworkSecurityPerimeterConfigurationsOperations
    :ivar notebook_workspaces: NotebookWorkspacesOperations operations
    :vartype notebook_workspaces: azure.mgmt.cosmosdb.operations.NotebookWorkspacesOperations
    :ivar private_endpoint_connections: PrivateEndpointConnectionsOperations operations
    :vartype private_endpoint_connections:
     azure.mgmt.cosmosdb.operations.PrivateEndpointConnectionsOperations
    :ivar private_link_resources: PrivateLinkResourcesOperations operations
    :vartype private_link_resources: azure.mgmt.cosmosdb.operations.PrivateLinkResourcesOperations
    :ivar restorable_database_accounts: RestorableDatabaseAccountsOperations operations
    :vartype restorable_database_accounts:
     azure.mgmt.cosmosdb.operations.RestorableDatabaseAccountsOperations
    :ivar restorable_sql_databases: RestorableSqlDatabasesOperations operations
    :vartype restorable_sql_databases:
     azure.mgmt.cosmosdb.operations.RestorableSqlDatabasesOperations
    :ivar restorable_sql_containers: RestorableSqlContainersOperations operations
    :vartype restorable_sql_containers:
     azure.mgmt.cosmosdb.operations.RestorableSqlContainersOperations
    :ivar restorable_sql_resources: RestorableSqlResourcesOperations operations
    :vartype restorable_sql_resources:
     azure.mgmt.cosmosdb.operations.RestorableSqlResourcesOperations
    :ivar restorable_mongodb_databases: RestorableMongodbDatabasesOperations operations
    :vartype restorable_mongodb_databases:
     azure.mgmt.cosmosdb.operations.RestorableMongodbDatabasesOperations
    :ivar restorable_mongodb_collections: RestorableMongodbCollectionsOperations operations
    :vartype restorable_mongodb_collections:
     azure.mgmt.cosmosdb.operations.RestorableMongodbCollectionsOperations
    :ivar restorable_mongodb_resources: RestorableMongodbResourcesOperations operations
    :vartype restorable_mongodb_resources:
     azure.mgmt.cosmosdb.operations.RestorableMongodbResourcesOperations
    :ivar restorable_gremlin_databases: RestorableGremlinDatabasesOperations operations
    :vartype restorable_gremlin_databases:
     azure.mgmt.cosmosdb.operations.RestorableGremlinDatabasesOperations
    :ivar restorable_gremlin_graphs: RestorableGremlinGraphsOperations operations
    :vartype restorable_gremlin_graphs:
     azure.mgmt.cosmosdb.operations.RestorableGremlinGraphsOperations
    :ivar restorable_gremlin_resources: RestorableGremlinResourcesOperations operations
    :vartype restorable_gremlin_resources:
     azure.mgmt.cosmosdb.operations.RestorableGremlinResourcesOperations
    :ivar restorable_tables: RestorableTablesOperations operations
    :vartype restorable_tables: azure.mgmt.cosmosdb.operations.RestorableTablesOperations
    :ivar restorable_table_resources: RestorableTableResourcesOperations operations
    :vartype restorable_table_resources:
     azure.mgmt.cosmosdb.operations.RestorableTableResourcesOperations
    :ivar service: ServiceOperations operations
    :vartype service: azure.mgmt.cosmosdb.operations.ServiceOperations
    :ivar throughput_pools: ThroughputPoolsOperations operations
    :vartype throughput_pools: azure.mgmt.cosmosdb.operations.ThroughputPoolsOperations
    :ivar throughput_pool: ThroughputPoolOperations operations
    :vartype throughput_pool: azure.mgmt.cosmosdb.operations.ThroughputPoolOperations
    :ivar throughput_pool_accounts: ThroughputPoolAccountsOperations operations
    :vartype throughput_pool_accounts:
     azure.mgmt.cosmosdb.operations.ThroughputPoolAccountsOperations
    :ivar throughput_pool_account: ThroughputPoolAccountOperations operations
    :vartype throughput_pool_account:
     azure.mgmt.cosmosdb.operations.ThroughputPoolAccountOperations
    :param credential: Credential needed for the client to connect to Azure. Required.
    :type credential: ~azure.core.credentials.TokenCredential
    :param subscription_id: The ID of the target subscription. The value must be an UUID. Required.
    :type subscription_id: str
    :param base_url: Service URL. Default value is "https://management.azure.com".
    :type base_url: str
    :keyword api_version: Api Version. Default value is "2024-12-01-preview". Note that overriding
     this default value may result in unsupported behavior.
    :paramtype api_version: str
    :keyword int polling_interval: Default waiting time between two polls for LRO operations if no
     Retry-After header is present.
    """

    def __init__(
        self,
        credential: "TokenCredential",
        subscription_id: str,
        base_url: str = "https://management.azure.com",
        **kwargs: Any
    ) -> None:
        self._config = CosmosDBManagementClientConfiguration(
            credential=credential, subscription_id=subscription_id, **kwargs
        )
        _policies = kwargs.pop("policies", None)
        if _policies is None:
            _policies = [
                policies.RequestIdPolicy(**kwargs),
                self._config.headers_policy,
                self._config.user_agent_policy,
                self._config.proxy_policy,
                policies.ContentDecodePolicy(**kwargs),
                ARMAutoResourceProviderRegistrationPolicy(),
                self._config.redirect_policy,
                self._config.retry_policy,
                self._config.authentication_policy,
                self._config.custom_hook_policy,
                self._config.logging_policy,
                policies.DistributedTracingPolicy(**kwargs),
                policies.SensitiveHeaderCleanupPolicy(**kwargs) if self._config.redirect_policy else None,
                self._config.http_logging_policy,
            ]
        self._client: ARMPipelineClient = ARMPipelineClient(base_url=base_url, policies=_policies, **kwargs)

        client_models = {k: v for k, v in _models.__dict__.items() if isinstance(v, type)}
        self._serialize = Serializer(client_models)
        self._deserialize = Deserializer(client_models)
        self._serialize.client_side_validation = False
        self.chaos_fault = ChaosFaultOperations(self._client, self._config, self._serialize, self._deserialize)
        self.database_accounts = DatabaseAccountsOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.operations = Operations(self._client, self._config, self._serialize, self._deserialize)
        self.database = DatabaseOperations(self._client, self._config, self._serialize, self._deserialize)
        self.collection = CollectionOperations(self._client, self._config, self._serialize, self._deserialize)
        self.collection_region = CollectionRegionOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.database_account_region = DatabaseAccountRegionOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.percentile_source_target = PercentileSourceTargetOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.percentile_target = PercentileTargetOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.percentile = PercentileOperations(self._client, self._config, self._serialize, self._deserialize)
        self.collection_partition_region = CollectionPartitionRegionOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.collection_partition = CollectionPartitionOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.partition_key_range_id = PartitionKeyRangeIdOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.partition_key_range_id_region = PartitionKeyRangeIdRegionOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.graph_resources = GraphResourcesOperations(self._client, self._config, self._serialize, self._deserialize)
        self.sql_resources = SqlResourcesOperations(self._client, self._config, self._serialize, self._deserialize)
        self.mongo_db_resources = MongoDBResourcesOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.table_resources = TableResourcesOperations(self._client, self._config, self._serialize, self._deserialize)
        self.cassandra_resources = CassandraResourcesOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.gremlin_resources = GremlinResourcesOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.locations = LocationsOperations(self._client, self._config, self._serialize, self._deserialize)
        self.data_transfer_jobs = DataTransferJobsOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.cassandra_clusters = CassandraClustersOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.cassandra_data_centers = CassandraDataCentersOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.network_security_perimeter_configurations = NetworkSecurityPerimeterConfigurationsOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.notebook_workspaces = NotebookWorkspacesOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.private_endpoint_connections = PrivateEndpointConnectionsOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.private_link_resources = PrivateLinkResourcesOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.restorable_database_accounts = RestorableDatabaseAccountsOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.restorable_sql_databases = RestorableSqlDatabasesOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.restorable_sql_containers = RestorableSqlContainersOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.restorable_sql_resources = RestorableSqlResourcesOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.restorable_mongodb_databases = RestorableMongodbDatabasesOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.restorable_mongodb_collections = RestorableMongodbCollectionsOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.restorable_mongodb_resources = RestorableMongodbResourcesOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.restorable_gremlin_databases = RestorableGremlinDatabasesOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.restorable_gremlin_graphs = RestorableGremlinGraphsOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.restorable_gremlin_resources = RestorableGremlinResourcesOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.restorable_tables = RestorableTablesOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.restorable_table_resources = RestorableTableResourcesOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.service = ServiceOperations(self._client, self._config, self._serialize, self._deserialize)
        self.throughput_pools = ThroughputPoolsOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.throughput_pool = ThroughputPoolOperations(self._client, self._config, self._serialize, self._deserialize)
        self.throughput_pool_accounts = ThroughputPoolAccountsOperations(
            self._client, self._config, self._serialize, self._deserialize
        )
        self.throughput_pool_account = ThroughputPoolAccountOperations(
            self._client, self._config, self._serialize, self._deserialize
        )

    def _send_request(self, request: HttpRequest, *, stream: bool = False, **kwargs: Any) -> HttpResponse:
        """Runs the network request through the client's chained policies.

        >>> from azure.core.rest import HttpRequest
        >>> request = HttpRequest("GET", "https://www.example.org/")
        <HttpRequest [GET], url: 'https://www.example.org/'>
        >>> response = client._send_request(request)
        <HttpResponse: 200 OK>

        For more information on this code flow, see https://aka.ms/azsdk/dpcodegen/python/send_request

        :param request: The network request you want to make. Required.
        :type request: ~azure.core.rest.HttpRequest
        :keyword bool stream: Whether the response payload will be streamed. Defaults to False.
        :return: The response of your network call. Does not do error handling on your response.
        :rtype: ~azure.core.rest.HttpResponse
        """

        request_copy = deepcopy(request)
        request_copy.url = self._client.format_url(request_copy.url)
        return self._client.send_request(request_copy, stream=stream, **kwargs)  # type: ignore

    def close(self) -> None:
        self._client.close()

    def __enter__(self) -> Self:
        self._client.__enter__()
        return self

    def __exit__(self, *exc_details: Any) -> None:
        self._client.__exit__(*exc_details)
