"""Configuration helpers for the OpenMetadata SDK."""
from __future__ import annotations

import os
from typing import Optional, Union

from metadata.generated.schema.security.client.openMetadataJWTClientConfig import (
    OpenMetadataJWTClientConfig,
)
from metadata.generated.schema.security.ssl import (
    validateSSLClientConfig,
    verifySSLConfig,
)

from metadata.sdk.heimdall import HeimdallConfig


class OpenMetadataConfig:
    """Configuration for OpenMetadata SDK.

    Supports two authentication modes:

    **Heimdall (recommended)**:
        Use ``heimdall_configuration`` (mirrors hera/config/config.yaml
        authenticationConfiguration.heimdallConfiguration) or set
        ``heimdall_url`` / ``HEIMDALL_BASE_URL``.

    **Direct JWT (legacy)**:
        Uses an OpenMetadata-issued JWT token when Heimdall is not configured.
    """

    server_url: str
    jwt_token: Optional[str]
    api_key: Optional[str]
    verify_ssl: bool
    ca_bundle: Optional[str]
    client_timeout: int

    _heimdall_configuration: Optional[HeimdallConfig]
    heimdall_url: Optional[str]
    heimdall_timeout: int
    heimdall_trust_all: bool

    def __init__(
        self,
        server_url: str,
        jwt_token: Optional[str] = None,
        api_key: Optional[str] = None,
        verify_ssl: bool = False,
        ca_bundle: Optional[str] = None,
        client_timeout: int = 30,
        heimdall_url: Optional[str] = None,
        heimdall_timeout: int = 10,
        heimdall_trust_all: bool = True,
        heimdall_configuration: Optional[Union[HeimdallConfig, dict]] = None,
    ):
        self.server_url = server_url.rstrip("/")
        self.jwt_token = jwt_token or api_key
        self.api_key = api_key
        self.verify_ssl = verify_ssl
        self.ca_bundle = ca_bundle
        self.client_timeout = client_timeout

        if heimdall_configuration is not None:
            self._heimdall_configuration = (
                HeimdallConfig.from_config_dict(heimdall_configuration)
                if isinstance(heimdall_configuration, dict)
                else heimdall_configuration
            )
            self.heimdall_url = (
                self._heimdall_configuration.base_url
                if self._heimdall_configuration.enabled
                else None
            )
            self.heimdall_timeout = self._heimdall_configuration.timeout
            self.heimdall_trust_all = self._heimdall_configuration.trust_all
        else:
            self._heimdall_configuration = None
            self.heimdall_url = heimdall_url
            self.heimdall_timeout = heimdall_timeout
            self.heimdall_trust_all = heimdall_trust_all

    @property
    def use_heimdall(self) -> bool:
        """Whether Heimdall authentication is enabled."""
        if self._heimdall_configuration is not None:
            return self._heimdall_configuration.enabled
        return self.heimdall_url is not None

    @property
    def heimdall_configuration(self) -> Optional[HeimdallConfig]:
        """Heimdall config when set (mirrors config.yaml heimdallConfiguration)."""
        return self._heimdall_configuration

    @classmethod
    def builder(cls) -> "OpenMetadataConfigBuilder":
        """Create a configuration builder."""
        return OpenMetadataConfigBuilder()

    @classmethod
    def from_env(cls) -> "OpenMetadataConfig":
        """Create configuration from environment variables.

        Reads from:
        - OPENMETADATA_HOST or OPENMETADATA_SERVER_URL: Server URL
        - OPENMETADATA_JWT_TOKEN or OPENMETADATA_API_KEY: Authentication token
        - OPENMETADATA_VERIFY_SSL: SSL verification (default: false)
        - OPENMETADATA_CA_BUNDLE: CA bundle path
        - OPENMETADATA_CLIENT_TIMEOUT: Client timeout in seconds (default: 30)
        - HEIMDALL_BASE_URL: Heimdall service URL (enables Heimdall auth)
        - HEIMDALL_TIMEOUT: Heimdall request timeout (default: 10)
        - HEIMDALL_TRUST_ALL: Trust all SSL certs for Heimdall (default: true)
        """
        server_url = os.environ.get("OPENMETADATA_HOST") or os.environ.get(
            "OPENMETADATA_SERVER_URL"
        )
        if not server_url:
            raise ValueError(
                "Server URL must be provided via 'OPENMETADATA_HOST' or "
                + "'OPENMETADATA_SERVER_URL' environment variable"
            )

        jwt_token = os.environ.get("OPENMETADATA_JWT_TOKEN") or os.environ.get(
            "OPENMETADATA_API_KEY"
        )
        verify_ssl = (
            os.environ.get("OPENMETADATA_VERIFY_SSL", "false").lower() == "true"
        )
        ca_bundle = os.environ.get("OPENMETADATA_CA_BUNDLE")
        client_timeout = int(os.environ.get("OPENMETADATA_CLIENT_TIMEOUT", "30"))

        heimdall_url = os.environ.get("HEIMDALL_BASE_URL")
        heimdall_timeout = int(os.environ.get("HEIMDALL_TIMEOUT", "10"))
        heimdall_trust_all = (
            os.environ.get("HEIMDALL_TRUST_ALL", "true").lower() == "true"
        )

        return cls(
            server_url=server_url,
            jwt_token=jwt_token,
            verify_ssl=verify_ssl,
            ca_bundle=ca_bundle,
            client_timeout=client_timeout,
            heimdall_url=heimdall_url,
            heimdall_timeout=heimdall_timeout,
            heimdall_trust_all=heimdall_trust_all,
        )

    def to_ometa_config(self) -> OpenMetadataJWTClientConfig:
        """Translate the SDK config into the ingestion client's config model."""
        token = self.jwt_token or self.api_key
        if token is None:
            raise ValueError("JWT token or API key is required to authenticate")

        return OpenMetadataJWTClientConfig.model_validate({"jwtToken": token})

    def to_ssl_config(self) -> verifySSLConfig.SslConfig | None:
        """Build an optional SSL configuration block."""
        if not self.ca_bundle:
            return None
        ssl_payload = validateSSLClientConfig.ValidateSslClientConfig.model_validate(
            {
                "caCertificate": self.ca_bundle,
                "sslCertificate": None,
                "sslKey": None,
            }
        )
        return verifySSLConfig.SslConfig(ssl_payload)


class OpenMetadataConfigBuilder:
    """Builder for :class:`OpenMetadataConfig`."""

    def __init__(self) -> None:
        self._server_url: Optional[str] = None
        self._jwt_token: Optional[str] = None
        self._api_key: Optional[str] = None
        self._verify_ssl: bool = False
        self._ca_bundle: Optional[str] = None
        self._client_timeout: int = 30
        self._heimdall_configuration: Optional[HeimdallConfig] = None
        self._heimdall_url: Optional[str] = None
        self._heimdall_timeout: int = 10
        self._heimdall_trust_all: bool = True

    def server_url(self, url: str) -> "OpenMetadataConfigBuilder":
        """Set server URL."""
        self._server_url = url
        return self

    def jwt_token(self, token: str) -> "OpenMetadataConfigBuilder":
        """Set JWT token."""
        self._jwt_token = token
        return self

    def api_key(self, key: str) -> "OpenMetadataConfigBuilder":
        """Set API key (alias for ``jwt_token``)."""
        self._api_key = key
        return self

    def verify_ssl(self, verify: bool) -> "OpenMetadataConfigBuilder":
        """Configure SSL verification."""
        self._verify_ssl = verify
        return self

    def ca_bundle(self, bundle: str) -> "OpenMetadataConfigBuilder":
        """Set CA bundle path."""
        self._ca_bundle = bundle
        return self

    def client_timeout(self, timeout: int) -> "OpenMetadataConfigBuilder":
        """Set client timeout in seconds."""
        self._client_timeout = timeout
        return self

    def heimdall_configuration(
        self,
        config: Union[HeimdallConfig, dict],
    ) -> "OpenMetadataConfigBuilder":
        """Set Heimdall auth from heimdallConfiguration (mirrors config.yaml)."""
        self._heimdall_configuration = (
            HeimdallConfig.from_config_dict(config)
            if isinstance(config, dict)
            else config
        )
        return self

    def heimdall_url(self, url: str) -> "OpenMetadataConfigBuilder":
        """Set the Heimdall base URL to enable Heimdall authentication."""
        self._heimdall_url = url
        return self

    def heimdall_timeout(self, timeout: int) -> "OpenMetadataConfigBuilder":
        """Set the Heimdall request timeout in seconds."""
        self._heimdall_timeout = timeout
        return self

    def heimdall_trust_all(self, trust: bool) -> "OpenMetadataConfigBuilder":
        """Configure whether to trust all SSL certificates for Heimdall."""
        self._heimdall_trust_all = trust
        return self

    def build(self) -> OpenMetadataConfig:
        """Build configuration."""
        if not self._server_url:
            raise ValueError("Server URL is required")

        return OpenMetadataConfig(
            server_url=self._server_url,
            jwt_token=self._jwt_token,
            api_key=self._api_key,
            verify_ssl=self._verify_ssl,
            ca_bundle=self._ca_bundle,
            client_timeout=self._client_timeout,
            heimdall_url=self._heimdall_url,
            heimdall_timeout=self._heimdall_timeout,
            heimdall_trust_all=self._heimdall_trust_all,
            heimdall_configuration=self._heimdall_configuration,
        )
