"""Unit tests for Databricks credential providers."""

import os
from unittest import mock

import pytest

from ray.data._internal.datasource.databricks_credentials import (
    DatabricksCredentialProvider,
    DatabricksTableCredentialConfig,
    EnvironmentCredentialProvider,
    StaticCredentialProvider,
    UnityCatalogCredentialConfig,
    resolve_credential_provider,
)

SAMPLE_TOKEN = "dapi_test_token_abc123"
SAMPLE_HOST = "https://my-workspace.cloud.databricks.com"
SAMPLE_URL = "https://uc-workspace.databricks.com"
ALT_TOKEN = "dapi_alt_token_xyz789"
ALT_HOST = "https://alt-workspace.databricks.com"


class TestDatabricksCredentialProvider:
    """Tests for the abstract DatabricksCredentialProvider base class."""

    def test_cannot_instantiate_abstract_class(self):
        """Verify DatabricksCredentialProvider cannot be instantiated directly."""
        with pytest.raises(TypeError, match="Can't instantiate abstract class"):
            DatabricksCredentialProvider()

    def test_abstract_methods_defined(self):
        """Verify all abstract methods are defined."""
        abstract_methods = DatabricksCredentialProvider.__abstractmethods__
        assert "get_token" in abstract_methods
        assert "get_host" in abstract_methods
        assert "invalidate" in abstract_methods


class TestStaticCredentialProvider:
    """Tests for StaticCredentialProvider."""

    def test_init_with_valid_token_and_host(self):
        """Test successful initialization with token and host."""
        provider = StaticCredentialProvider(token=SAMPLE_TOKEN, host=SAMPLE_HOST)
        assert provider.get_token() == SAMPLE_TOKEN
        assert provider.get_host() == SAMPLE_HOST

    @pytest.mark.parametrize(
        "token,host,expected_error",
        [
            ("", SAMPLE_HOST, "Token cannot be empty"),
            (None, SAMPLE_HOST, "Token cannot be empty"),
            (SAMPLE_TOKEN, "", "Host cannot be empty"),
            (SAMPLE_TOKEN, None, "Host cannot be empty"),
        ],
    )
    def test_init_with_invalid_inputs_raises_error(self, token, host, expected_error):
        """Test that invalid token or host raises ValueError."""
        with pytest.raises(ValueError, match=expected_error):
            StaticCredentialProvider(token=token, host=host)

    def test_invalidate_is_noop(self):
        """Test that invalidate doesn't affect the static token."""
        provider = StaticCredentialProvider(token=SAMPLE_TOKEN, host=SAMPLE_HOST)
        provider.invalidate()
        assert provider.get_token() == SAMPLE_TOKEN
        assert provider.get_host() == SAMPLE_HOST

    def test_get_token_returns_same_value(self):
        """Test that get_token always returns the same value."""
        provider = StaticCredentialProvider(token=SAMPLE_TOKEN, host=SAMPLE_HOST)
        assert provider.get_token() == SAMPLE_TOKEN
        assert provider.get_token() == SAMPLE_TOKEN


class TestEnvironmentCredentialProvider:
    """Tests for EnvironmentCredentialProvider."""

    def test_get_token_from_env(self):
        """Test get_token reads from environment variable."""
        with mock.patch.dict(
            os.environ,
            {"DATABRICKS_TOKEN": SAMPLE_TOKEN, "DATABRICKS_HOST": SAMPLE_HOST},
        ):
            provider = EnvironmentCredentialProvider()
            assert provider.get_token() == SAMPLE_TOKEN

    def test_get_host_from_env(self):
        """Test get_host reads from environment variable."""
        with mock.patch.dict(
            os.environ,
            {"DATABRICKS_TOKEN": SAMPLE_TOKEN, "DATABRICKS_HOST": SAMPLE_HOST},
        ):
            provider = EnvironmentCredentialProvider()
            assert provider.get_host() == SAMPLE_HOST

    @pytest.mark.parametrize(
        "env_vars,expected_error",
        [
            ({"DATABRICKS_HOST": SAMPLE_HOST}, "DATABRICKS_TOKEN.*not set"),
            (
                {"DATABRICKS_TOKEN": SAMPLE_TOKEN},
                "set environment variable.*DATABRICKS_HOST",
            ),
        ],
    )
    def test_init_raises_when_env_var_not_set(self, env_vars, expected_error):
        """Test __init__ raises ValueError when required env var is not set."""
        with mock.patch.dict(os.environ, env_vars, clear=True):
            with pytest.raises(ValueError, match=expected_error):
                EnvironmentCredentialProvider()

    def test_host_detected_from_databricks_runtime(self):
        """Test host is detected from Databricks runtime when env var not set."""
        detected_host = "detected-host.databricks.com"
        with (
            mock.patch.dict(os.environ, {"DATABRICKS_TOKEN": SAMPLE_TOKEN}, clear=True),
            mock.patch.object(
                EnvironmentCredentialProvider,
                "_detect_databricks_host",
                return_value=detected_host,
            ),
        ):
            provider = EnvironmentCredentialProvider()
            assert provider.get_host() == detected_host

    def test_custom_env_var_names(self):
        """Test using custom environment variable names."""
        with mock.patch.dict(
            os.environ, {"MY_TOKEN": SAMPLE_TOKEN, "MY_HOST": SAMPLE_HOST}
        ):
            provider = EnvironmentCredentialProvider(
                token_env_var="MY_TOKEN", host_env_var="MY_HOST"
            )
            assert provider.get_token() == SAMPLE_TOKEN
            assert provider.get_host() == SAMPLE_HOST

    def test_invalidate_refreshes_token_from_env(self):
        """Test that invalidate re-reads token from environment."""
        refreshed_token = "dapi_refreshed_token_456"
        with mock.patch.dict(
            os.environ,
            {"DATABRICKS_TOKEN": SAMPLE_TOKEN, "DATABRICKS_HOST": SAMPLE_HOST},
        ):
            provider = EnvironmentCredentialProvider()
            assert provider.get_token() == SAMPLE_TOKEN

            # Simulate external token refresh
            os.environ["DATABRICKS_TOKEN"] = refreshed_token
            provider.invalidate()
            assert provider.get_token() == refreshed_token

    def test_invalidate_keeps_token_if_env_unset(self):
        """Test that invalidate keeps existing token if env var is unset."""
        with mock.patch.dict(
            os.environ,
            {"DATABRICKS_TOKEN": SAMPLE_TOKEN, "DATABRICKS_HOST": SAMPLE_HOST},
        ):
            provider = EnvironmentCredentialProvider()

            # Remove env var after initialization
            del os.environ["DATABRICKS_TOKEN"]
            provider.invalidate()
            # Should keep the old token rather than failing
            assert provider.get_token() == SAMPLE_TOKEN


class TestDatabricksTableCredentialConfig:
    """Tests for DatabricksTableCredentialConfig and resolve_credential_provider."""

    def test_resolve_with_explicit_provider(self):
        """Test that explicit credential_provider is returned as-is."""
        provider = StaticCredentialProvider(token=SAMPLE_TOKEN, host=SAMPLE_HOST)
        config = DatabricksTableCredentialConfig(credential_provider=provider)
        result = resolve_credential_provider(config)
        assert result is provider

    @pytest.mark.parametrize("credential_provider_arg", [None, "no_arg"])
    def test_resolve_with_none_returns_environment_provider(
        self, credential_provider_arg
    ):
        """Test that EnvironmentCredentialProvider is returned when none provided."""
        with mock.patch.dict(
            os.environ,
            {"DATABRICKS_TOKEN": SAMPLE_TOKEN, "DATABRICKS_HOST": SAMPLE_HOST},
        ):
            if credential_provider_arg == "no_arg":
                config = DatabricksTableCredentialConfig()
            else:
                config = DatabricksTableCredentialConfig(
                    credential_provider=credential_provider_arg
                )
            result = resolve_credential_provider(config)
            assert isinstance(result, EnvironmentCredentialProvider)


class TestUnityCatalogCredentialConfig:
    """Tests for UnityCatalogCredentialConfig and resolve_credential_provider."""

    def test_resolve_with_explicit_provider(self):
        """Test that explicit credential_provider is returned as-is."""
        provider = StaticCredentialProvider(token=SAMPLE_TOKEN, host=SAMPLE_HOST)
        config = UnityCatalogCredentialConfig(credential_provider=provider)
        result = resolve_credential_provider(config)
        assert result is provider

    def test_resolve_with_explicit_provider_ignores_url_and_token(self):
        """Test that url/token are ignored when credential_provider is given."""
        provider = StaticCredentialProvider(token=SAMPLE_TOKEN, host=SAMPLE_HOST)
        config = UnityCatalogCredentialConfig(
            credential_provider=provider, url=ALT_HOST, token=ALT_TOKEN
        )
        result = resolve_credential_provider(config)
        assert result is provider

    def test_resolve_with_url_and_token(self):
        """Test that url and token create a StaticCredentialProvider."""
        config = UnityCatalogCredentialConfig(url=SAMPLE_URL, token=SAMPLE_TOKEN)
        result = resolve_credential_provider(config)
        assert isinstance(result, StaticCredentialProvider)
        assert result.get_token() == SAMPLE_TOKEN
        assert result.get_host() == SAMPLE_URL

    @pytest.mark.parametrize(
        "kwargs",
        [
            {},
            {"url": SAMPLE_URL},
            {"token": SAMPLE_TOKEN},
        ],
        ids=["no_args", "only_url", "only_token"],
    )
    def test_config_raises_with_incomplete_args(self, kwargs):
        """Test that ValueError is raised when args are missing or incomplete."""
        config = UnityCatalogCredentialConfig(**kwargs)
        with pytest.raises(ValueError, match="Either 'credential_provider' or both"):
            resolve_credential_provider(config)

    @pytest.mark.parametrize(
        "url,token",
        [
            ("", SAMPLE_TOKEN),
            (SAMPLE_URL, ""),
        ],
        ids=["empty_url", "empty_token"],
    )
    def test_resolve_with_empty_string_raises(self, url, token):
        """Test that empty strings for url or token raise ValueError."""
        config = UnityCatalogCredentialConfig(url=url, token=token)
        with pytest.raises(ValueError):
            resolve_credential_provider(config)


class TestCredentialProviderSerialization:
    """Tests for credential provider serialization (needed for Ray workers)."""

    @pytest.mark.parametrize(
        "provider_type,expected_token,expected_host",
        [
            ("static", SAMPLE_TOKEN, SAMPLE_HOST),
            ("environment", SAMPLE_TOKEN, SAMPLE_HOST),
        ],
    )
    def test_provider_is_picklable(self, provider_type, expected_token, expected_host):
        """Verify credential providers can be pickled and unpickled."""
        import pickle

        with mock.patch.dict(
            os.environ,
            {"DATABRICKS_TOKEN": expected_token, "DATABRICKS_HOST": expected_host},
        ):
            if provider_type == "static":
                provider = StaticCredentialProvider(
                    token=expected_token, host=expected_host
                )
            else:
                provider = EnvironmentCredentialProvider()

            pickled = pickle.dumps(provider)
            unpickled = pickle.loads(pickled)
            assert unpickled.get_token() == expected_token
            assert unpickled.get_host() == expected_host


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", __file__]))