"""OpenMetadata SDK Client - Main client class."""
from __future__ import annotations

import logging
from typing import ClassVar, Optional, cast

from metadata.generated.schema.entity.services.connections.metadata.openMetadataConnection import (
    AuthProvider,
    OpenMetadataConnection,
    OpenmetadataType,
)
from metadata.generated.schema.security.ssl.verifySSLConfig import VerifySSL
from metadata.ingestion.ometa.ometa_api import OpenMetadata as OMeta
from metadata.sdk.config import OpenMetadataConfig
from metadata.sdk.types import OMetaClient

logger = logging.getLogger("metadata.sdk.client")


class OpenMetadata:
    """Main SDK client for OpenMetadata."""

    _instance: ClassVar[Optional["OpenMetadata"]] = None
    _default_client: ClassVar[Optional[OMetaClient]] = None

    def __init__(self, config: OpenMetadataConfig):
        """Initialize OpenMetadata client.

        When ``config.heimdall_url`` is set, tokens are pre-validated
        through the Heimdall authorization service before being used
        for API calls.
        """
        self.config: OpenMetadataConfig = config

        if not config.verify_ssl:
            verify_ssl = VerifySSL.no_ssl
        elif config.ca_bundle:
            verify_ssl = VerifySSL.validate
        else:
            verify_ssl = VerifySSL.ignore

        ssl_config = config.to_ssl_config()

        om_connection = OpenMetadataConnection.model_construct(
            hostPort=config.server_url,
            authProvider=AuthProvider.openmetadata,
            securityConfig=config.to_ometa_config(),
            verifySSL=verify_ssl,
            sslConfig=ssl_config,
            type=OpenmetadataType.OpenMetadata,
            clusterName="openmetadata",
        )

        ometa = OMeta(config=om_connection)

        if config.use_heimdall:
            logger.info(
                "Heimdall auth enabled — validating token via %s",
                config.heimdall_url,
            )
            self._patch_heimdall_auth(ometa, config)

        self._ometa = cast(OMetaClient, ometa)

    @staticmethod
    def _patch_heimdall_auth(ometa: OMeta, config: OpenMetadataConfig) -> None:
        """Replace the default auth provider with a Heimdall-backed one.

        The HeimdallAuthenticationProvider pre-validates tokens against
        the Heimdall ``/api/v1/authorize`` endpoint, providing early and
        clear errors when tokens are invalid or expired.
        """
        from metadata.ingestion.ometa.auth_provider import (
            HeimdallAuthenticationProvider,
        )

        heimdall_provider = HeimdallAuthenticationProvider.create(
            ometa.config,
            heimdall_url=config.heimdall_url,
            heimdall_timeout=config.heimdall_timeout,
            heimdall_trust_all=config.heimdall_trust_all,
        )

        ometa._auth_provider = heimdall_provider
        ometa.client._auth_token = heimdall_provider.get_access_token
        ometa.client.config.access_token = None
        ometa.client.config.expires_in = None

    @classmethod
    def initialize(cls, config: OpenMetadataConfig) -> "OpenMetadata":
        """Initialize the default client instance."""
        cls._instance = cls(config)
        cls._default_client = cls._instance.ometa
        return cls._instance

    @classmethod
    def get_instance(cls) -> "OpenMetadata":
        """Get the default client instance."""
        if cls._instance is None:
            raise RuntimeError(
                "OpenMetadata client not initialized. Call initialize() first"
            )
        return cls._instance

    @classmethod
    def get_default_client(cls) -> OMetaClient:
        """Get the default OMeta client for internal use."""
        if cls._default_client is None:
            raise RuntimeError(
                "OpenMetadata client not initialized. Call initialize() first"
            )
        return cls._default_client

    @property
    def ometa(self) -> OMetaClient:
        """Get the underlying OMeta client."""
        return self._ometa

    def close(self):
        """Close the client connection."""
        if hasattr(self._ometa, "close"):
            self._ometa.close()

    @classmethod
    def reset(cls) -> None:
        """Clear the singleton instance and default client."""
        if cls._instance is not None:
            cls._instance.close()
        cls._instance = None
        cls._default_client = None
