import pytest

from huggingface_hub.utils import get_hf_hub_version, get_python_version
from huggingface_hub.utils._headers import _deduplicate_user_agent, _http_user_agent, build_hf_headers


# Only for tests that are not related to user agent
DEFAULT_USER_AGENT = _http_user_agent()

FAKE_TOKEN = "123456789"
FAKE_TOKEN_ORG = "api_org_123456789"
FAKE_TOKEN_HEADER = {
    "authorization": f"Bearer {FAKE_TOKEN}",
    "user-agent": DEFAULT_USER_AGENT,
}
NO_AUTH_HEADER = {"user-agent": DEFAULT_USER_AGENT}


class TestAuthHeadersUtil:
    def test_token_str(self) -> None:
        assert build_hf_headers(token=FAKE_TOKEN) == FAKE_TOKEN_HEADER

    def test_token_true_no_cached_token(self, mocker) -> None:
        mocker.patch("huggingface_hub.utils._headers.get_token", return_value=None)
        with pytest.raises(EnvironmentError):
            build_hf_headers(token=True)

    def test_token_true_has_cached_token(self, mocker) -> None:
        mocker.patch("huggingface_hub.utils._headers.get_token", return_value=FAKE_TOKEN)
        assert build_hf_headers(token=True) == FAKE_TOKEN_HEADER

    def test_token_false(self, mocker) -> None:
        mocker.patch("huggingface_hub.utils._headers.get_token", return_value=FAKE_TOKEN)
        assert build_hf_headers(token=False) == NO_AUTH_HEADER

    def test_token_none_no_cached_token(self, mocker) -> None:
        mocker.patch("huggingface_hub.utils._headers.get_token", return_value=None)
        assert build_hf_headers() == NO_AUTH_HEADER

    def test_token_none_has_cached_token(self, mocker) -> None:
        mocker.patch("huggingface_hub.utils._headers.get_token", return_value=FAKE_TOKEN)
        assert build_hf_headers() == FAKE_TOKEN_HEADER

    def test_implicit_use_disabled(self, mocker) -> None:
        mocker.patch("huggingface_hub.utils._headers.get_token", return_value=FAKE_TOKEN)
        mocker.patch("huggingface_hub.constants.HF_HUB_DISABLE_IMPLICIT_TOKEN", True)
        assert build_hf_headers() == NO_AUTH_HEADER  # token is not sent

    def test_implicit_use_disabled_but_explicit_use(self, mocker) -> None:
        mocker.patch("huggingface_hub.utils._headers.get_token", return_value=FAKE_TOKEN)
        mocker.patch("huggingface_hub.constants.HF_HUB_DISABLE_IMPLICIT_TOKEN", True)
        # This is not an implicit use so we still send it
        assert build_hf_headers(token=True) == FAKE_TOKEN_HEADER


class TestUserAgentHeadersUtil:
    def _get_user_agent(self, **kwargs) -> str:
        return build_hf_headers(**kwargs)["user-agent"]

    def test_default_user_agent(self, mocker) -> None:
        mocker.patch("huggingface_hub.utils._headers.get_torch_version", return_value="torch_version")
        mocker.patch("huggingface_hub.utils._headers.is_torch_available", return_value=True)
        assert (
            self._get_user_agent()
            == f"unknown/None; hf_hub/{get_hf_hub_version()}; python/{get_python_version()}; torch/torch_version"
        )

    def test_user_agent_with_library_name_multiple_missing(self, mocker) -> None:
        mocker.patch("huggingface_hub.utils._headers.is_torch_available", return_value=False)
        assert "torch" not in self._get_user_agent()

    def test_user_agent_with_library_name_and_version(self) -> None:
        assert self._get_user_agent(
            library_name="foo",
            library_version="bar",
        ).startswith("foo/bar;")

    def test_user_agent_with_library_name_no_version(self) -> None:
        assert self._get_user_agent(library_name="foo").startswith("foo/None;")

    def test_user_agent_with_custom_agent_string(self) -> None:
        assert self._get_user_agent(user_agent="this is a custom agent").endswith("this is a custom agent")

    def test_user_agent_with_custom_agent_dict(self) -> None:
        assert self._get_user_agent(user_agent={"a": "b", "c": "d"}).endswith("a/b; c/d")

    def test_user_agent_deduplicate(self) -> None:
        assert (
            _deduplicate_user_agent(
                "python/3.7; python/3.8; hf_hub/0.12; transformers/None; hf_hub/0.12; python/3.7; diffusers/0.12.1"
            )
            # 1. "python" is kept twice with different values
            # 2. "python/3.7" second occurrence is removed
            # 3. "hf_hub" second occurrence is removed
            # 4. order is preserved
            == "python/3.7; python/3.8; hf_hub/0.12; transformers/None; diffusers/0.12.1"
        )

    def test_user_agent_with_origin(self, mocker) -> None:
        mocker.patch("huggingface_hub.utils._telemetry.constants.HF_HUB_USER_AGENT_ORIGIN", "custom-origin")
        assert self._get_user_agent().endswith("origin/custom-origin")

    def test_user_agent_with_origin_and_user_agent(self, mocker) -> None:
        mocker.patch("huggingface_hub.utils._telemetry.constants.HF_HUB_USER_AGENT_ORIGIN", "custom-origin")
        assert self._get_user_agent(user_agent={"a": "b", "c": "d"}).endswith("a/b; c/d; origin/custom-origin")

    def test_user_agent_with_origin_and_user_agent_str(self, mocker) -> None:
        mocker.patch("huggingface_hub.utils._telemetry.constants.HF_HUB_USER_AGENT_ORIGIN", "custom-origin")
        assert self._get_user_agent(user_agent="a/b;c/d").endswith("a/b; c/d; origin/custom-origin")
