import os
from contextlib import contextmanager
from pathlib import Path
from unittest.mock import DEFAULT, MagicMock, Mock, patch

import pytest

from huggingface_hub import snapshot_download
from huggingface_hub.file_download import (
    HfFileMetadata,
    get_hf_file_metadata,
    hf_hub_download,
    hf_hub_url,
    try_to_load_from_cache,
    xet_get,
)
from huggingface_hub.utils import XetFileData

from .testing_constants import (
    DUMMY_XET_FILE,
    DUMMY_XET_MODEL_ID,
)


pytestmark = pytest.mark.xet


@pytest.mark.production
class TestXetFileDownload:
    @contextmanager
    def _patch_xet_file_metadata(self, with_xet_data: bool):
        patcher = patch("huggingface_hub.file_download.get_hf_file_metadata")
        mock_metadata = patcher.start()
        mock_metadata.return_value = HfFileMetadata(
            commit_hash="mock_commit",
            etag="mock_etag",
            location="mock_location",
            size=1024,
            xet_file_data=XetFileData(file_hash="mock_hash", refresh_route="mock/route") if with_xet_data else None,
        )
        try:
            yield mock_metadata
        finally:
            patcher.stop()

    def test_xet_get_called_when_xet_metadata_present(self, tmp_path):
        """Test that xet_get is called when xet metadata is present."""
        with self._patch_xet_file_metadata(with_xet_data=True) as mock_file_metadata:
            with patch("huggingface_hub.file_download.xet_get") as mock_xet_get:
                with patch("huggingface_hub.file_download._create_symlink"):
                    hf_hub_download(
                        DUMMY_XET_MODEL_ID,
                        filename=DUMMY_XET_FILE,
                        cache_dir=tmp_path,
                        force_download=True,
                    )

                    # Verify xet_get was called with correct parameters
                    mock_xet_get.assert_called_once()
                    _, kwargs = mock_xet_get.call_args
                    assert "xet_file_data" in kwargs
                    assert kwargs["xet_file_data"] == mock_file_metadata.return_value.xet_file_data

    def test_backward_compatibility_no_xet_metadata(self, tmp_path):
        """Test backward compatibility when response has no xet metadata."""
        with self._patch_xet_file_metadata(with_xet_data=False):
            with patch("huggingface_hub.file_download.http_get") as mock_http_get:
                with patch("huggingface_hub.file_download._create_symlink"):
                    hf_hub_download(
                        DUMMY_XET_MODEL_ID,
                        filename=DUMMY_XET_FILE,
                        cache_dir=tmp_path,
                        force_download=True,
                    )

                    # Verify http_get was called
                    mock_http_get.assert_called_once()

    def test_get_xet_file_metadata_basic(self) -> None:
        """Test getting metadata from a file on the Hub."""
        url = hf_hub_url(
            repo_id=DUMMY_XET_MODEL_ID,
            filename=DUMMY_XET_FILE,
        )
        metadata = get_hf_file_metadata(url)
        assert metadata.xet_file_data is not None
        assert metadata.xet_file_data.file_hash is not None
        assert metadata.xet_file_data.refresh_route is not None

    def test_basic_download(self, tmp_path):
        # Make sure that xet_get is called
        with patch("huggingface_hub.file_download.xet_get", wraps=xet_get) as _xet_get:
            filepath = hf_hub_download(
                DUMMY_XET_MODEL_ID,
                filename=DUMMY_XET_FILE,
                cache_dir=tmp_path,
            )

            assert os.path.exists(filepath)
            assert os.path.getsize(filepath) > 0

            _xet_get.assert_called_once()

    def test_try_to_load_from_cache(self, tmp_path):
        cached_path = try_to_load_from_cache(
            DUMMY_XET_MODEL_ID,
            filename=DUMMY_XET_FILE,
            cache_dir=tmp_path,
        )
        assert cached_path is None

        downloaded_path = hf_hub_download(
            DUMMY_XET_MODEL_ID,
            filename=DUMMY_XET_FILE,
            cache_dir=tmp_path,
        )

        # Now should find it in cache
        cached_path = try_to_load_from_cache(
            DUMMY_XET_MODEL_ID,
            filename=DUMMY_XET_FILE,
            cache_dir=tmp_path,
        )
        assert cached_path == downloaded_path

    def test_cache_reuse(self, tmp_path):
        path1 = hf_hub_download(
            DUMMY_XET_MODEL_ID,
            filename=DUMMY_XET_FILE,
            cache_dir=tmp_path,
        )

        assert os.path.exists(path1)

        with patch("huggingface_hub.file_download._download_to_tmp_and_move") as mock:
            # Second download should use cache
            path2 = hf_hub_download(
                DUMMY_XET_MODEL_ID,
                filename=DUMMY_XET_FILE,
                cache_dir=tmp_path,
            )

            assert path1 == path2
            mock.assert_not_called()

    def test_download_to_local_dir(self, tmp_path):
        local_dir = tmp_path / "local_dir"
        local_dir.mkdir(exist_ok=True, parents=True)

        cache_dir = tmp_path / "cache"
        cache_dir.mkdir(exist_ok=True, parents=True)
        # Download to local dir
        returned_path = hf_hub_download(
            DUMMY_XET_MODEL_ID,
            filename=DUMMY_XET_FILE,
            local_dir=local_dir,
            cache_dir=cache_dir,
        )
        assert local_dir in Path(returned_path).parents

        for path in cache_dir.glob("**/blobs/**"):
            assert not path.is_file()
        for path in cache_dir.glob("**/snapshots/**"):
            assert not path.is_file()

    def test_force_download(self, tmp_path):
        # First download
        path1 = hf_hub_download(
            DUMMY_XET_MODEL_ID,
            filename=DUMMY_XET_FILE,
            cache_dir=tmp_path,
        )

        # Force download should re-download even if in cache
        with patch("huggingface_hub.file_download.xet_get") as mock_xet_get:
            path2 = hf_hub_download(
                DUMMY_XET_MODEL_ID,
                filename=DUMMY_XET_FILE,
                cache_dir=tmp_path,
                force_download=True,
            )

            assert path1 == path2
            mock_xet_get.assert_called_once()

    def test_fallback_to_http_when_xet_not_available(self, tmp_path):
        """Test that http_get is used when hf_xet is not available."""
        with self._patch_xet_file_metadata(with_xet_data=True):
            # Mock is_xet_available to return False
            with patch.multiple(
                "huggingface_hub.file_download",
                is_xet_available=Mock(return_value=False),
                http_get=DEFAULT,
                xet_get=DEFAULT,
                _create_symlink=DEFAULT,
            ) as mocks:
                hf_hub_download(
                    DUMMY_XET_MODEL_ID,
                    filename=DUMMY_XET_FILE,
                    cache_dir=tmp_path,
                    force_download=True,
                )

                # Verify http_get was called and xet_get was not
                mocks["http_get"].assert_called_once()
                mocks["xet_get"].assert_not_called()

    def test_use_xet_when_available(self, tmp_path):
        """Test that xet_get is used when hf_xet is available."""
        with self._patch_xet_file_metadata(with_xet_data=True):
            with patch.multiple(
                "huggingface_hub.file_download",
                is_xet_available=Mock(return_value=True),
                http_get=DEFAULT,
                xet_get=DEFAULT,
                _create_symlink=DEFAULT,
            ) as mocks:
                hf_hub_download(
                    DUMMY_XET_MODEL_ID,
                    filename=DUMMY_XET_FILE,
                    cache_dir=tmp_path,
                    force_download=True,
                )

                # Verify xet_get was called and http_get was not
                mocks["xet_get"].assert_called_once()
                mocks["http_get"].assert_not_called()

    def test_request_headers_passed_to_download_files(self, tmp_path):
        """Test that headers (minus authorization) are passed as custom_headers to new_file_download_group."""
        headers = {
            "authorization": "Bearer my_token",
            "x-custom-header": "custom_value",
            "user-agent": "test-agent",
        }

        mock_group = MagicMock()
        mock_session = MagicMock()
        mock_session.new_file_download_group.return_value = mock_group

        with self._patch_xet_file_metadata(with_xet_data=True):
            with patch("huggingface_hub.utils._xet.get_xet_session", return_value=mock_session):
                with patch("huggingface_hub.file_download._create_symlink"):
                    hf_hub_download(
                        DUMMY_XET_MODEL_ID,
                        filename=DUMMY_XET_FILE,
                        cache_dir=tmp_path,
                        force_download=True,
                        headers=headers,
                    )

        mock_group.__enter__.return_value.start_download_file.assert_called_once()
        kwargs = mock_session.new_file_download_group.call_args.kwargs
        # Verify custom_headers excludes authorization
        assert kwargs["custom_headers"].get("x-custom-header") == "custom_value"
        assert kwargs["custom_headers"].get("user-agent") == "test-agent"
        assert "authorization" not in kwargs["custom_headers"]
        # Verify full headers (incl. auth) passed as token_refresh_headers
        assert "authorization" in kwargs["token_refresh_headers"]


@pytest.mark.production
class TestXetSnapshotDownload:
    def test_download_model(self, tmp_path):
        """Test that snapshot_download works with Xet storage."""
        storage_folder = snapshot_download(
            DUMMY_XET_MODEL_ID,
            cache_dir=tmp_path,
        )

        assert os.path.exists(storage_folder)
        assert os.path.isdir(storage_folder)
        assert os.path.exists(os.path.join(storage_folder, DUMMY_XET_FILE))

        with open(os.path.join(storage_folder, DUMMY_XET_FILE), "rb") as f:
            content = f.read()
            assert len(content) > 0

    def test_snapshot_download_cache_reuse(self, tmp_path):
        """Test that snapshot_download reuses cached files."""
        # First download
        storage_folder1 = snapshot_download(
            DUMMY_XET_MODEL_ID,
            cache_dir=tmp_path,
        )

        with patch("huggingface_hub.file_download.xet_get") as mock_xet_get:
            # Second download should use cache
            storage_folder2 = snapshot_download(
                DUMMY_XET_MODEL_ID,
                cache_dir=tmp_path,
            )

            # Verify same folder is returned
            assert storage_folder1 == storage_folder2

            # Verify xet_get was not called (files were cached)
            mock_xet_get.assert_not_called()
