from __future__ import annotations

from pathlib import Path

from cognite.client import ClientConfig, CogniteClient, data_modeling as dm
from cognite.client.credentials import OAuthClientCredentials
{% for class_ in api_classes %}
from ._api.{{ class_.file_name }} import {{ class_.name }}{% endfor %}

{% if not is_single_model %}{% for multi_api_class, views in zip(multi_api_classes, view_sets) %}
class {{ multi_api_class.name }}:
    """
    {{ multi_api_class.name }}

    Data Model:
        space: {{ multi_api_class.model_id.space }}
        externalId: {{ multi_api_class.model_id.external_id }}
        version: {{ multi_api_class.model_id.version }}

    """

    def __init__(self, client: CogniteClient):{% for sub, view in zip(multi_api_class.sub_apis, views) %}
        self.{{ sub.client_attribute }} = {{ sub.name }}(client, dm.ViewId("{{ view.space }}", "{{ view.external_id }}", "{{ view.version }}")){% endfor %}

{% endfor %}{% endif %}
class {{ client_name }}:
    """
    {{ client_name }}

    Generated with:
        pygen = {{ pygen_version }}
        cognite-sdk = {{ cognite_sdk_version }}
        pydantic = {{ pydantic_version }}
{% if is_single_model %}
    Data Model:
        space: {{ data_model.space }}
        externalId: {{ data_model.external_id }}
        version: {{ data_model.version }}{% endif %}
    """

    def __init__(self, config_or_client: CogniteClient | ClientConfig):
        if isinstance(config_or_client, CogniteClient):
            client = config_or_client
        elif isinstance(config_or_client, ClientConfig):
            client = CogniteClient(config_or_client)
        else:
            raise ValueError(f"Expected CogniteClient or ClientConfig, got {type(config_or_client)}"){% if is_single_model %}{% for api_class, view in zip(api_classes, views) %}
        self.{{ api_class.client_attribute }} = {{ api_class.name }}(client, dm.ViewId("{{ view.space }}", "{{ view.external_id }}", "{{ view.version }}")){% endfor %}{% else %}{% for multi_api_class in multi_api_classes %}
        self.{{ multi_api_class.client_attribute }} = {{ multi_api_class.name }}(client){% endfor %}{% endif %}

    @classmethod
    def azure_project(
        cls, tenant_id: str, client_id: str, client_secret: str, cdf_cluster: str, project: str
    ) -> {{ client_name }}:
        credentials = OAuthClientCredentials.default_for_azure_ad(tenant_id, client_id, client_secret, cdf_cluster)
        config = ClientConfig.default(project, cdf_cluster, credentials)

        return cls(config)

    @classmethod
    def from_toml(cls, file_path: Path | str, section: str | None = "cognite") -> {{ client_name }}:
        import toml

        toml_content = toml.load(file_path)
        if section is not None:
            try:
                toml_content = toml_content[section]
            except KeyError as e:
                raise ValueError(f"Could not find section '{section}' in {file_path}") from e

        return cls.azure_project(**toml_content)
