"""Tests for sharing.py — two-tier Parquet index + safetensors tensor store."""

import json
import shutil
from pathlib import Path
from unittest.mock import MagicMock, patch

import pytest
import torch

import lmprobe.cache as cache_mod
from lmprobe.cache import (
    CachedPromptInfo,
    discover_cached,
    load_prompt_activations,
    load_prompt_pooled_activations,
    save_prompt_activations,
    save_prompt_logits,
    save_prompt_pooled_activations,
)
from lmprobe.sharing import (
    _STAGE_SENTINEL,
    FORMAT_VERSION,
    INFO_FILENAME,
    PARQUET_PATH,
    _build_lmprobe_info,
    _build_readme,
    _check_shards_on_remote,
    _compute_shard_boundaries_variable,
    _compute_shard_plan,
    _compute_tensor_intersection,
    _consolidate_and_shard,
    _discover_prompts,
    _enumerate_shard_files,
    _filter_tensor_types,
    _load_manifest,
    _new_manifest,
    _reconstruct_plan_from_cached,
    _save_manifest,
    _staging_dir_path,
    _write_parquet_index,
    load_activation_dataset,
    pull_dataset,
    push_dataset,
    upgrade_dataset_format,
)

_has_pyarrow = True
try:
    import pyarrow  # noqa: F401
except ImportError:
    _has_pyarrow = False

requires_pyarrow = pytest.mark.skipif(
    not _has_pyarrow, reason="pyarrow not installed"
)

TEST_MODEL = "stas/tiny-random-llama-2"
HIDDEN_DIM = 32
SEQ_LEN = 5


@pytest.fixture
def cache_dir(tmp_path, monkeypatch):
    """Set up a temporary cache directory."""
    monkeypatch.setenv("LMPROBE_CACHE_DIR", str(tmp_path / "cache"))
    return tmp_path / "cache"


@pytest.fixture
def populated_cache(cache_dir):
    """Populate cache with a small dataset (pooled activations)."""
    prompts = [
        "Who wants to go for a walk?",
        "Fetch the ball!",
        "Purring and scratching",
    ]
    for prompt in prompts:
        for layer in [0, 1]:
            pooled = torch.randn(1, HIDDEN_DIM)
            save_prompt_pooled_activations(
                TEST_MODEL, prompt, [layer], pooled, "last_token"
            )
    return prompts


@pytest.fixture
def populated_raw_cache(cache_dir):
    """Populate cache with raw (full-sequence) activations.

    Variable seq_lens: 3, 5, 4 tokens.
    """
    prompts = [
        "Who wants to go for a walk?",
        "Fetch the ball!",
        "Purring and scratching",
    ]
    seq_lens = [3, 5, 4]
    layers = [0, 1]
    original_raw = {}

    for prompt, sl in zip(prompts, seq_lens):
        # activations: (1, seq_len, hidden_dim * n_layers)
        act = torch.randn(1, sl, HIDDEN_DIM * len(layers))
        mask = torch.ones(1, sl, dtype=torch.long)
        save_prompt_activations(TEST_MODEL, prompt, layers, act, mask)
        original_raw[prompt] = (act, mask, sl)

    return prompts, seq_lens, layers, original_raw


class TestComputeTensorIntersection:
    def test_single_info(self):
        info = CachedPromptInfo(
            raw_layers=[0, 1],
            pooled={"last_token": [0, 1]},
            has_logits=False,
            logits_top_k=None,
            has_perplexity=True, has_token_perplexity=False,
            num_tokens=5,
        )
        result = _compute_tensor_intersection([info])
        assert result["raw_layers"] == [0, 1]
        assert result["pooled"] == {"last_token": [0, 1]}
        assert result["has_logits"] is False
        assert result["has_perplexity"] is True

    def test_intersection_of_layers(self):
        info1 = CachedPromptInfo(
            raw_layers=[0, 1, 2], pooled={}, has_logits=False,
            logits_top_k=None, has_perplexity=False, has_token_perplexity=False, num_tokens=5,
        )
        info2 = CachedPromptInfo(
            raw_layers=[1, 2, 3], pooled={}, has_logits=False,
            logits_top_k=None, has_perplexity=False, has_token_perplexity=False, num_tokens=5,
        )
        result = _compute_tensor_intersection([info1, info2])
        assert result["raw_layers"] == [1, 2]

    def test_intersection_of_pooled_strategies(self):
        info1 = CachedPromptInfo(
            raw_layers=[], pooled={"last_token": [0], "mean": [0]},
            has_logits=False, logits_top_k=None, has_perplexity=False, has_token_perplexity=False,
            num_tokens=5,
        )
        info2 = CachedPromptInfo(
            raw_layers=[], pooled={"last_token": [0]},
            has_logits=False, logits_top_k=None, has_perplexity=False, has_token_perplexity=False,
            num_tokens=5,
        )
        result = _compute_tensor_intersection([info1, info2])
        assert "last_token" in result["pooled"]
        assert "mean" not in result["pooled"]

    def test_logits_topk_consistent(self):
        info1 = CachedPromptInfo(
            raw_layers=[], pooled={}, has_logits=False,
            logits_top_k=100, has_perplexity=False, has_token_perplexity=False, num_tokens=5,
        )
        info2 = CachedPromptInfo(
            raw_layers=[], pooled={}, has_logits=False,
            logits_top_k=100, has_perplexity=False, has_token_perplexity=False, num_tokens=5,
        )
        result = _compute_tensor_intersection([info1, info2])
        assert result["has_logits"] is True
        assert result["logits_top_k"] == 100

    def test_logits_topk_inconsistent(self):
        info1 = CachedPromptInfo(
            raw_layers=[], pooled={}, has_logits=False,
            logits_top_k=100, has_perplexity=False, has_token_perplexity=False, num_tokens=5,
        )
        info2 = CachedPromptInfo(
            raw_layers=[], pooled={}, has_logits=False,
            logits_top_k=50, has_perplexity=False, has_token_perplexity=False, num_tokens=5,
        )
        result = _compute_tensor_intersection([info1, info2])
        assert result["has_logits"] is False

    def test_topk_preferred_over_full(self):
        info = CachedPromptInfo(
            raw_layers=[], pooled={}, has_logits=True,
            logits_top_k=100, has_perplexity=False, has_token_perplexity=False, num_tokens=5,
        )
        result = _compute_tensor_intersection([info, info])
        assert result["has_logits"] is True
        assert result["logits_top_k"] == 100

    def test_perplexity_intersection(self):
        info1 = CachedPromptInfo(
            raw_layers=[], pooled={}, has_logits=False,
            logits_top_k=None, has_perplexity=True, has_token_perplexity=False, num_tokens=5,
        )
        info2 = CachedPromptInfo(
            raw_layers=[], pooled={}, has_logits=False,
            logits_top_k=None, has_perplexity=False, has_token_perplexity=False, num_tokens=5,
        )
        result = _compute_tensor_intersection([info1, info2])
        assert result["has_perplexity"] is False


class TestFilterTensorTypes:
    def test_no_filter_returns_all(self):
        available = {
            "raw_layers": [0, 1],
            "pooled": {"last_token": [0]},
            "has_logits": True,
            "logits_top_k": 100,
            "has_perplexity": True,
        }
        result = _filter_tensor_types(available, None)
        assert result == available

    def test_filter_hidden_layers(self):
        available = {
            "raw_layers": [0, 1, 2],
            "pooled": {"last_token": [0, 1]},
            "has_logits": True,
            "logits_top_k": 100,
            "has_perplexity": False,
        }
        result = _filter_tensor_types(available, ["hidden_layers"])
        assert result["pooled"] == {"last_token": [0, 1]}
        assert result["raw_layers"] == [0, 1, 2]
        assert result["has_logits"] is False

    def test_filter_logits(self):
        available = {
            "raw_layers": [0],
            "pooled": {"last_token": [0]},
            "has_logits": True,
            "logits_top_k": 100,
            "has_perplexity": False,
        }
        result = _filter_tensor_types(available, ["logits_topk"])
        assert result["has_logits"] is True
        assert result["pooled"] == {}
        assert result["raw_layers"] == []


class TestDiscoverPrompts:
    def test_empty_prompts_raises(self, cache_dir):
        with pytest.raises(ValueError, match="No prompts"):
            _discover_prompts(TEST_MODEL, [])

    def test_all_missing_raises(self, cache_dir):
        with pytest.raises(ValueError, match="No prompts have cached data"):
            _discover_prompts(
                TEST_MODEL, ["missing1", "missing2"], skip_missing=True
            )

    def test_skip_missing_false_raises(self, cache_dir):
        with pytest.raises(FileNotFoundError):
            _discover_prompts(
                TEST_MODEL, ["missing"], skip_missing=False
            )

    def test_finds_cached_prompts(self, populated_cache):
        kept, infos = _discover_prompts(
            TEST_MODEL, populated_cache, skip_missing=True
        )
        assert len(kept) == 3
        assert len(infos) == 3
        assert all(isinstance(i, CachedPromptInfo) for i in infos)

    def test_partial_cache(self, populated_cache):
        prompts = populated_cache + ["uncached prompt"]
        kept, infos = _discover_prompts(
            TEST_MODEL, prompts, skip_missing=True
        )
        assert len(kept) == 3
        assert 3 not in kept


class TestConsolidateAndShard:
    def test_basic_consolidation(self, populated_cache):
        tensor_types = {
            "raw_layers": [],
            "pooled": {"last_token": [0, 1]},
            "has_logits": False,
            "logits_top_k": None,
            "has_perplexity": False,
        }
        tmpdir, tensor_descriptors, prompt_metadata = (
            _consolidate_and_shard(
                model_name=TEST_MODEL,
                prompts=populated_cache,
                kept_indices=[0, 1, 2],
                tensor_types=tensor_types,
                labels=None,
                shard_max_bytes=1_000_000_000,
                repo_id="user/test",
            )
        )
        assert len(prompt_metadata) == 3
        assert "hidden_layers" in tensor_descriptors
        desc = tensor_descriptors["hidden_layers"]
        assert desc["layers"] == [0, 1]
        assert desc["type"] == "hidden"
        assert desc["layout"] == "per_layer"

        # Verify per-layer shard files exist
        for shard_idx, shard in enumerate(desc["shards"]):
            for layer in desc["layers"]:
                fname = f"tensors/hidden_layer{layer:03d}_shard{shard_idx:03d}.safetensors"
                assert (tmpdir / fname).exists()

        shutil.rmtree(tmpdir, ignore_errors=True)

    def test_consolidation_with_labels(self, populated_cache):
        tensor_types = {
            "raw_layers": [],
            "pooled": {"last_token": [0]},
            "has_logits": False,
            "logits_top_k": None,
            "has_perplexity": False,
        }
        labels = [1, 1, 0]
        _, _, prompt_metadata = _consolidate_and_shard(
            model_name=TEST_MODEL,
            prompts=populated_cache,
            kept_indices=[0, 1, 2],
            tensor_types=tensor_types,
            labels=labels,
            shard_max_bytes=1_000_000_000,
            repo_id="user/test",
        )
        # Labels are present (order may be shuffled)
        label_set = {p["label"] for p in prompt_metadata}
        assert 1 in label_set
        assert 0 in label_set

    def test_shuffle_false_preserves_order(self, populated_cache):
        """shuffle=False preserves input prompt order in output metadata."""
        tensor_types = {
            "raw_layers": [],
            "pooled": {"last_token": [0]},
            "has_logits": False,
            "logits_top_k": None,
            "has_perplexity": False,
        }
        labels = [1, 1, 0]
        _, _, prompt_metadata = _consolidate_and_shard(
            model_name=TEST_MODEL,
            prompts=populated_cache,
            kept_indices=[0, 1, 2],
            tensor_types=tensor_types,
            labels=labels,
            shard_max_bytes=1_000_000_000,
            repo_id="user/test",
            shuffle=False,
        )
        output_prompts = [p["text"] for p in prompt_metadata]
        assert output_prompts == populated_cache
        output_labels = [p["label"] for p in prompt_metadata]
        assert output_labels == labels

    def test_shuffle_true_is_default(self, populated_cache):
        """Default shuffle=True produces a different order (for this seed)."""
        tensor_types = {
            "raw_layers": [],
            "pooled": {"last_token": [0]},
            "has_logits": False,
            "logits_top_k": None,
            "has_perplexity": False,
        }
        _, _, meta_shuffled = _consolidate_and_shard(
            model_name=TEST_MODEL,
            prompts=populated_cache,
            kept_indices=[0, 1, 2],
            tensor_types=tensor_types,
            labels=None,
            shard_max_bytes=1_000_000_000,
            repo_id="user/test",
        )
        _, _, meta_unshuffled = _consolidate_and_shard(
            model_name=TEST_MODEL,
            prompts=populated_cache,
            kept_indices=[0, 1, 2],
            tensor_types=tensor_types,
            labels=None,
            shard_max_bytes=1_000_000_000,
            repo_id="user/test",
            shuffle=False,
        )
        shuffled_prompts = [p["text"] for p in meta_shuffled]
        unshuffled_prompts = [p["text"] for p in meta_unshuffled]
        # Unshuffled should match input order
        assert unshuffled_prompts == populated_cache
        # Both contain same prompts
        assert set(shuffled_prompts) == set(unshuffled_prompts)

    def test_small_shard_limit(self, populated_cache):
        """Test that small shard limits create multiple shards."""
        tensor_types = {
            "raw_layers": [],
            "pooled": {"last_token": [0]},
            "has_logits": False,
            "logits_top_k": None,
            "has_perplexity": False,
        }
        tmpdir, tensor_descriptors, _ = _consolidate_and_shard(
            model_name=TEST_MODEL,
            prompts=populated_cache,
            kept_indices=[0, 1, 2],
            tensor_types=tensor_types,
            labels=None,
            shard_max_bytes=1,  # 1 byte — forces each prompt into own shard
            repo_id="user/test",
        )
        shards = tensor_descriptors["hidden_layers"]["shards"]
        assert len(shards) >= 2

        shutil.rmtree(tmpdir, ignore_errors=True)

    def test_per_layer_files(self, populated_cache):
        """Verify each layer gets its own shard file with a single key."""
        tensor_types = {
            "raw_layers": [],
            "pooled": {"last_token": [0, 1]},
            "has_logits": False,
            "logits_top_k": None,
            "has_perplexity": False,
        }
        tmpdir, tensor_descriptors, _ = _consolidate_and_shard(
            model_name=TEST_MODEL,
            prompts=populated_cache,
            kept_indices=[0, 1, 2],
            tensor_types=tensor_types,
            labels=None,
            shard_max_bytes=1_000_000_000,
            repo_id="user/test",
        )

        from safetensors import safe_open

        for layer in [0, 1]:
            fname = f"tensors/hidden_layer{layer:03d}_shard000.safetensors"
            shard_file = tmpdir / fname
            assert shard_file.exists()
            with safe_open(str(shard_file), framework="pt") as f:
                keys = list(f.keys())
            assert keys == [f"hidden.layer_{layer}"]

        shutil.rmtree(tmpdir, ignore_errors=True)

    def test_shard_metadata(self, populated_cache):
        """Verify prompt_metadata has shard_index and row_offset."""
        tensor_types = {
            "raw_layers": [],
            "pooled": {"last_token": [0]},
            "has_logits": False,
            "logits_top_k": None,
            "has_perplexity": False,
        }
        _, _, prompt_metadata = _consolidate_and_shard(
            model_name=TEST_MODEL,
            prompts=populated_cache,
            kept_indices=[0, 1, 2],
            tensor_types=tensor_types,
            labels=None,
            shard_max_bytes=1_000_000_000,
            repo_id="user/test",
        )

        for pm in prompt_metadata:
            assert "shard_index" in pm
            assert "row_offset" in pm


@requires_pyarrow
class TestParquetIndex:
    def test_write_and_read(self, tmp_path):
        import pyarrow.parquet as pq

        (tmp_path / "index").mkdir()
        prompt_metadata = [
            {"text": "hello", "label": 1, "num_tokens": 3,
             "shard_index": 0, "row_offset": 0},
            {"text": "world", "label": 0, "num_tokens": 4,
             "shard_index": 0, "row_offset": 1},
        ]
        _write_parquet_index(tmp_path, prompt_metadata)

        table = pq.read_table(str(tmp_path / PARQUET_PATH))

        assert table.num_rows == 2
        assert table.column_names == [
            "text", "label", "num_tokens", "shard_index", "row_offset",
        ]
        assert table.column("text").to_pylist() == ["hello", "world"]
        assert table.column("label").to_pylist() == [1, 0]
        assert table.column("shard_index").to_pylist() == [0, 0]
        assert table.column("row_offset").to_pylist() == [0, 1]

    def test_string_labels(self, tmp_path):
        import pyarrow.parquet as pq

        (tmp_path / "index").mkdir()
        prompt_metadata = [
            {"text": "a", "label": "positive", "num_tokens": 1,
             "shard_index": 0, "row_offset": 0},
            {"text": "b", "label": "negative", "num_tokens": 1,
             "shard_index": 0, "row_offset": 1},
        ]
        _write_parquet_index(tmp_path, prompt_metadata)

        table = pq.read_table(str(tmp_path / PARQUET_PATH))
        assert table.column("label").to_pylist() == ["positive", "negative"]


class TestBuildLmprobeInfo:
    @patch("huggingface_hub.model_info", side_effect=Exception("no network"))
    def test_basic_structure(self, mock_model_info):
        info = _build_lmprobe_info(TEST_MODEL, 100, {"hidden_layers": {}})
        assert info["format_version"] == FORMAT_VERSION
        assert info["model"]["name"] == TEST_MODEL
        assert info["num_prompts"] == 100
        assert info["prompt_ordering"] == "random"
        assert "tensors" in info
        assert "lmprobe_version" in info["provenance"]
        assert "torch_version" in info["provenance"]


class TestBuildReadme:
    def test_readme_has_yaml_frontmatter(self):
        lmprobe_info = {
            "format_version": "1.0",
            "model": {"name": TEST_MODEL, "revision": "abc123"},
            "tensors": {},
            "provenance": {
                "lmprobe_version": "0.7.3",
                "extraction_backend": "local",
                "created_at": "2026-01-01",
                "torch_version": "2.0",
                "transformers_version": "4.0",
            },
        }
        readme = _build_readme(
            TEST_MODEL, lmprobe_info, 0, "user/test",
        )
        assert readme.startswith("---\n")
        assert "tags:" in readme
        assert "- lmprobe" in readme
        assert "- activations" in readme
        assert "- interpretability" in readme
        assert "task_categories:" in readme
        assert "- feature-extraction" in readme
        assert "license: cc-by-4.0" in readme

    def test_readme_custom_license(self):
        lmprobe_info = {
            "format_version": "1.0",
            "model": {"name": TEST_MODEL, "revision": "abc123"},
            "tensors": {},
            "provenance": {},
        }
        readme = _build_readme(
            TEST_MODEL, lmprobe_info, 0, "user/test",
            license="apache-2.0",
        )
        assert "license: apache-2.0" in readme

    def test_readme_uses_new_api_names(self):
        lmprobe_info = {
            "format_version": "1.0",
            "model": {"name": TEST_MODEL, "revision": "abc123"},
            "tensors": {},
            "provenance": {},
        }
        readme = _build_readme(
            TEST_MODEL, lmprobe_info, 0, "user/test",
        )
        assert "load_activations" in readme
        assert "fit_from_activations" in readme
        # Old names should not appear
        assert "push_to_bucket" not in readme
        assert "pull_from_bucket" not in readme
        assert "load_from_bucket" not in readme

    def test_readme_contains_standalone_loading(self):
        lmprobe_info = {
            "format_version": "1.0",
            "model": {"name": TEST_MODEL, "revision": "abc123"},
            "tensors": {
                "hidden_layers": {
                    "type": "hidden",
                    "layers": [0, 1],
                    "dim": 32,
                    "pooling": "last_token",
                    "row_bytes": 256,
                    "shards": [
                        {"file": "tensors/hidden_layers_000.safetensors",
                         "num_prompts": 3},
                    ],
                },
            },
            "provenance": {},
        }
        readme = _build_readme(
            TEST_MODEL, lmprobe_info, 3, "user/test",
        )
        assert "Load without lmprobe" in readme
        assert "pyarrow.parquet" in readme
        assert "safe_open" in readme
        assert "row_offset" in readme
        assert "load_dataset" in readme

    def test_readme_description(self):
        lmprobe_info = {
            "format_version": "1.0",
            "model": {"name": TEST_MODEL, "revision": "abc123"},
            "tensors": {},
            "provenance": {},
        }
        readme = _build_readme(
            TEST_MODEL, lmprobe_info, 0, "user/test",
            description="Test dataset",
        )
        assert "Test dataset" in readme


@requires_pyarrow
class TestPushDataset:
    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    @patch("huggingface_hub.HfApi")
    def test_push_creates_correct_files(
        self, MockHfApi, mock_pyarrow, mock_deps, populated_cache,
    ):
        mock_api = MagicMock()
        MockHfApi.return_value = mock_api

        uploaded_files = {}

        def capture_upload(repo_id, folder_path, **kwargs):
            folder = Path(folder_path)
            for f in folder.rglob("*"):
                if f.is_file():
                    rel = f.relative_to(folder)
                    if f.suffix == ".json":
                        uploaded_files[str(rel)] = json.loads(f.read_text())
                    elif f.suffix == ".md":
                        uploaded_files[str(rel)] = f.read_text()
                    else:
                        uploaded_files[str(rel)] = f.read_bytes()

        mock_api.upload_large_folder.side_effect = capture_upload

        url = push_dataset(
            repo_id="user/test-dataset",
            model_name=TEST_MODEL,
            prompts=populated_cache,
            labels=[1, 1, 0],
            exist_ok=True,
        )

        assert "user/test-dataset" in url
        mock_api.create_repo.assert_called_once()
        mock_api.upload_large_folder.assert_called_once()

        # Check file structure
        assert INFO_FILENAME in uploaded_files
        assert PARQUET_PATH in uploaded_files
        assert "README.md" in uploaded_files

        info = uploaded_files[INFO_FILENAME]
        assert info["num_prompts"] == 3
        assert "hidden_layers" in info["tensors"]
        assert info["prompt_ordering"] == "random"
        assert info["tensors"]["hidden_layers"]["layout"] == "per_layer"

    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    @patch("huggingface_hub.HfApi")
    def test_push_always_uses_upload_large_folder(
        self, MockHfApi, mock_pyarrow, mock_deps, populated_cache,
    ):
        """upload_large_folder is always used regardless of dataset size."""
        mock_api = MagicMock()
        MockHfApi.return_value = mock_api

        push_dataset(
            repo_id="user/test-dataset",
            model_name=TEST_MODEL,
            prompts=populated_cache,
            labels=[1, 1, 0],
            exist_ok=True,
        )

        mock_api.upload_large_folder.assert_called_once()
        mock_api.upload_folder.assert_not_called()

    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    @patch("huggingface_hub.HfApi")
    def test_push_num_workers_passed_to_upload(
        self, MockHfApi, mock_pyarrow, mock_deps, populated_cache,
    ):
        """num_workers is forwarded to upload_large_folder."""
        mock_api = MagicMock()
        MockHfApi.return_value = mock_api

        push_dataset(
            repo_id="user/test-dataset",
            model_name=TEST_MODEL,
            prompts=populated_cache,
            labels=[1, 1, 0],
            exist_ok=True,
            num_workers=8,
        )

        mock_api.upload_large_folder.assert_called_once()
        call_kwargs = mock_api.upload_large_folder.call_args[1]
        assert call_kwargs["num_workers"] == 8

    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    @patch("huggingface_hub.HfApi")
    def test_push_commit_batch_size_overrides_scale(
        self, MockHfApi, mock_pyarrow, mock_deps, populated_cache,
    ):
        """commit_batch_size monkey-patches COMMIT_SIZE_SCALE during upload."""
        import huggingface_hub._upload_large_folder as _ulf

        original_scale = _ulf.COMMIT_SIZE_SCALE

        captured_scale = []

        def capture_scale(**kwargs):
            captured_scale.append(list(_ulf.COMMIT_SIZE_SCALE))

        mock_api = MagicMock()
        MockHfApi.return_value = mock_api
        mock_api.upload_large_folder.side_effect = capture_scale

        push_dataset(
            repo_id="user/test-dataset",
            model_name=TEST_MODEL,
            prompts=populated_cache,
            labels=[1, 1, 0],
            exist_ok=True,
            commit_batch_size=3,
        )

        # During upload, scale should have been overridden
        assert captured_scale == [[3]]
        # After upload, original scale should be restored
        assert _ulf.COMMIT_SIZE_SCALE is original_scale

    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    @patch("huggingface_hub.HfApi")
    def test_push_commit_batch_size_restored_on_error(
        self, MockHfApi, mock_pyarrow, mock_deps, populated_cache,
    ):
        """COMMIT_SIZE_SCALE is restored even if upload raises."""
        import huggingface_hub._upload_large_folder as _ulf

        original_scale = _ulf.COMMIT_SIZE_SCALE

        mock_api = MagicMock()
        MockHfApi.return_value = mock_api
        mock_api.upload_large_folder.side_effect = RuntimeError("connection lost")

        with pytest.raises(RuntimeError, match="connection lost"):
            push_dataset(
                repo_id="user/test-dataset",
                model_name=TEST_MODEL,
                prompts=populated_cache,
                labels=[1, 1, 0],
                exist_ok=True,
                commit_batch_size=1,
            )

        assert _ulf.COMMIT_SIZE_SCALE is original_scale

    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    def test_push_labels_length_mismatch(
        self, mock_pyarrow, mock_deps, populated_cache,
    ):
        with pytest.raises(ValueError, match="labels length"):
            push_dataset(
                repo_id="user/test",
                model_name=TEST_MODEL,
                prompts=populated_cache,
                labels=[1, 0],
            )

    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    def test_push_empty_prompts_raises(
        self, mock_pyarrow, mock_deps, cache_dir,
    ):
        with pytest.raises(ValueError, match="No prompts"):
            push_dataset(
                repo_id="user/empty",
                model_name=TEST_MODEL,
                prompts=[],
            )


class TestPushDatasetStreaming:
    @requires_pyarrow
    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    @patch("huggingface_hub.HfApi")
    def test_stream_uses_create_commit(
        self, MockHfApi, mock_pyarrow, mock_deps, populated_cache,
    ):
        """stream=True uses create_commit, not upload_large_folder."""
        mock_api = MagicMock()
        MockHfApi.return_value = mock_api

        url = push_dataset(
            repo_id="user/stream-test",
            model_name=TEST_MODEL,
            prompts=populated_cache,
            labels=[1, 1, 0],
            exist_ok=True,
            stream=True,
        )

        assert "user/stream-test" in url
        mock_api.create_repo.assert_called_once()
        mock_api.upload_large_folder.assert_not_called()
        mock_api.upload_file.assert_not_called()
        assert mock_api.create_commit.call_count > 0

    @requires_pyarrow
    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    @patch("huggingface_hub.HfApi")
    def test_stream_produces_correct_repo_paths(
        self, MockHfApi, mock_pyarrow, mock_deps, populated_cache,
    ):
        """Streaming uploads all expected tensor and metadata paths."""
        mock_api = MagicMock()
        MockHfApi.return_value = mock_api

        uploaded_paths = []

        def capture_commit(*, repo_id, operations, **kwargs):
            for op in operations:
                uploaded_paths.append(op.path_in_repo)

        mock_api.create_commit.side_effect = capture_commit

        push_dataset(
            repo_id="user/stream-paths",
            model_name=TEST_MODEL,
            prompts=populated_cache,
            labels=[1, 1, 0],
            exist_ok=True,
            stream=True,
        )

        # Should have tensor shards for layers 0 and 1
        tensor_paths = [p for p in uploaded_paths if p.startswith("tensors/")]
        assert len(tensor_paths) >= 2  # at least one shard per layer

        # Should have metadata files (v2: no lmprobe_info.json in streaming)
        assert PARQUET_PATH in uploaded_paths
        assert "README.md" in uploaded_paths
        assert INFO_FILENAME not in uploaded_paths

    @requires_pyarrow
    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    @patch("huggingface_hub.HfApi")
    def test_stream_manifest_tracks_progress(
        self, MockHfApi, mock_pyarrow, mock_deps, populated_cache,
    ):
        """Manifest tracks completed shards; survives partial upload failure."""
        mock_api = MagicMock()
        MockHfApi.return_value = mock_api

        call_count = 0

        def fail_on_second_commit(*, repo_id, operations, **kwargs):
            nonlocal call_count
            call_count += 1
            if call_count > 1:
                raise ConnectionError("simulated network failure")

        mock_api.create_commit.side_effect = fail_on_second_commit

        # Use stream_batch_size=1 so each shard gets its own commit
        staging_dir = _staging_dir_path(
            "user/manifest-test", TEST_MODEL, populated_cache,
            labels=[1, 1, 0], stream=True, stream_batch_size=1,
        )

        with pytest.raises(ConnectionError):
            push_dataset(
                repo_id="user/manifest-test",
                model_name=TEST_MODEL,
                prompts=populated_cache,
                labels=[1, 1, 0],
                exist_ok=True,
                stream=True,
                stream_batch_size=1,
            )

        # Manifest should exist with at least one completed shard
        manifest = _load_manifest(staging_dir)
        assert manifest is not None
        assert len(manifest["completed_shards"]) == 1
        assert manifest["metadata_uploaded"] is False

    @requires_pyarrow
    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    @patch("huggingface_hub.HfApi")
    def test_stream_resume_skips_completed(
        self, MockHfApi, mock_pyarrow, mock_deps, populated_cache,
    ):
        """Resume skips already-uploaded shards (verified via remote check)."""
        mock_api = MagicMock()
        MockHfApi.return_value = mock_api

        staging_dir = _staging_dir_path(
            "user/resume-test", TEST_MODEL, populated_cache,
            labels=[1, 1, 0], stream=True, stream_batch_size=10,
        )

        # Pre-populate manifest (simulating partial completion)
        staging_dir.mkdir(parents=True, exist_ok=True)
        (staging_dir / "tensors").mkdir(parents=True, exist_ok=True)
        (staging_dir / "index").mkdir(parents=True, exist_ok=True)
        shard_paths = [
            "tensors/hidden_layer000_shard000.safetensors",
            "tensors/hidden_layer001_shard000.safetensors",
        ]
        manifest = _new_manifest("user/resume-test")
        manifest["completed_shards"] = shard_paths
        _save_manifest(staging_dir, manifest)

        # Mock repo_info to report shards as existing on remote
        mock_sibling = type("Sibling", (), {})
        mock_repo_info = MagicMock()
        mock_repo_info.siblings = [
            mock_sibling() for _ in shard_paths
        ]
        for sib, path in zip(mock_repo_info.siblings, shard_paths):
            sib.rfilename = path
        mock_api.repo_info.return_value = mock_repo_info

        commit_paths = []

        def capture_commit(*, repo_id, operations, **kwargs):
            for op in operations:
                commit_paths.append(op.path_in_repo)

        mock_api.create_commit.side_effect = capture_commit

        push_dataset(
            repo_id="user/resume-test",
            model_name=TEST_MODEL,
            prompts=populated_cache,
            labels=[1, 1, 0],
            exist_ok=True,
            stream=True,
        )

        # Only metadata files should be uploaded (not tensor shards)
        for path in commit_paths:
            assert not path.startswith("tensors/"), (
                f"Shard {path} should have been skipped on resume"
            )
        # Should have uploaded metadata files only (v2: no JSON sidecar)
        assert PARQUET_PATH in commit_paths
        assert "README.md" in commit_paths
        assert INFO_FILENAME not in commit_paths

    @requires_pyarrow
    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    @patch("huggingface_hub.HfApi")
    def test_stream_cleanup_on_success(
        self, MockHfApi, mock_pyarrow, mock_deps, populated_cache,
    ):
        """Staging dir is removed after successful streaming upload."""
        mock_api = MagicMock()
        MockHfApi.return_value = mock_api

        staging_dir = _staging_dir_path(
            "user/cleanup-stream", TEST_MODEL, populated_cache,
            stream=True, stream_batch_size=10,
        )

        push_dataset(
            repo_id="user/cleanup-stream",
            model_name=TEST_MODEL,
            prompts=populated_cache,
            exist_ok=True,
            stream=True,
        )

        assert not staging_dir.exists()

    @requires_pyarrow
    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    @patch("huggingface_hub.HfApi")
    def test_stream_raw_cache(
        self, MockHfApi, mock_pyarrow, mock_deps, populated_raw_cache,
    ):
        """Streaming works with raw (full-sequence) activations."""
        mock_api = MagicMock()
        MockHfApi.return_value = mock_api

        prompts, seq_lens, layers, _ = populated_raw_cache
        uploaded_paths = []

        def capture_commit(*, repo_id, operations, **kwargs):
            for op in operations:
                uploaded_paths.append(op.path_in_repo)

        mock_api.create_commit.side_effect = capture_commit

        push_dataset(
            repo_id="user/stream-raw",
            model_name=TEST_MODEL,
            prompts=prompts,
            exist_ok=True,
            stream=True,
        )

        mock_api.upload_large_folder.assert_not_called()
        tensor_paths = [p for p in uploaded_paths if p.startswith("tensors/")]
        assert len(tensor_paths) >= 2  # at least one shard per layer

    @requires_pyarrow
    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    @patch("huggingface_hub.HfApi")
    def test_stream_warns_ignored_params(
        self, MockHfApi, mock_pyarrow, mock_deps, populated_cache, caplog,
    ):
        """num_workers and commit_batch_size produce warnings in stream mode."""
        import logging

        mock_api = MagicMock()
        MockHfApi.return_value = mock_api

        with caplog.at_level(logging.WARNING, logger="lmprobe.sharing"):
            push_dataset(
                repo_id="user/stream-warn",
                model_name=TEST_MODEL,
                prompts=populated_cache,
                exist_ok=True,
                stream=True,
                num_workers=8,
                commit_batch_size=5,
            )

        assert "num_workers is ignored" in caplog.text
        assert "commit_batch_size is ignored" in caplog.text

    @requires_pyarrow
    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    @patch("huggingface_hub.HfApi")
    def test_nonstream_still_uses_upload_large_folder(
        self, MockHfApi, mock_pyarrow, mock_deps, populated_cache,
    ):
        """stream=False (default) still uses upload_large_folder."""
        mock_api = MagicMock()
        MockHfApi.return_value = mock_api

        push_dataset(
            repo_id="user/no-stream",
            model_name=TEST_MODEL,
            prompts=populated_cache,
            exist_ok=True,
            stream=False,
        )

        mock_api.upload_large_folder.assert_called_once()
        mock_api.create_commit.assert_not_called()

    @requires_pyarrow
    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    @patch("huggingface_hub.HfApi")
    def test_stream_retries_pending_batch_on_resume(
        self, MockHfApi, mock_pyarrow, mock_deps, populated_cache,
    ):
        """Resume retries a pending batch from a prior failed create_commit."""
        from safetensors.torch import save_file

        mock_api = MagicMock()
        MockHfApi.return_value = mock_api

        staging_dir = _staging_dir_path(
            "user/retry-batch", TEST_MODEL, populated_cache,
            labels=[1, 1, 0], stream=True, stream_batch_size=10,
        )

        # Simulate a prior failed run: manifest has a pending_batch
        # with local files that still exist on disk.
        staging_dir.mkdir(parents=True, exist_ok=True)
        (staging_dir / "tensors").mkdir(parents=True, exist_ok=True)
        (staging_dir / "index").mkdir(parents=True, exist_ok=True)

        shard_file = staging_dir / "tensors" / "hidden_layer000_shard000.safetensors"
        save_file(
            {"hidden.layer_0": torch.randn(3, HIDDEN_DIM)},
            str(shard_file),
        )

        manifest = _new_manifest("user/retry-batch")
        manifest["pending_batch"] = [
            [str(shard_file), "tensors/hidden_layer000_shard000.safetensors"],
        ]
        _save_manifest(staging_dir, manifest)

        commit_messages = []

        def track_commits(*, repo_id, operations, commit_message="", **kw):
            commit_messages.append(commit_message)

        mock_api.create_commit.side_effect = track_commits

        push_dataset(
            repo_id="user/retry-batch",
            model_name=TEST_MODEL,
            prompts=populated_cache,
            labels=[1, 1, 0],
            exist_ok=True,
            stream=True,
        )

        # First commit should be the retry of the pending batch
        assert any("retry" in m.lower() for m in commit_messages), (
            f"Expected a retry commit, got: {commit_messages}"
        )

    @requires_pyarrow
    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    @patch("huggingface_hub.HfApi")
    def test_stream_batch_size_controls_commit_grouping(
        self, MockHfApi, mock_pyarrow, mock_deps, populated_cache,
    ):
        """stream_batch_size=1 produces one commit per shard."""
        mock_api = MagicMock()
        MockHfApi.return_value = mock_api

        shard_commit_count = 0

        def count_commits(*, repo_id, operations, commit_message="", **kw):
            nonlocal shard_commit_count
            if "shard" in commit_message.lower():
                shard_commit_count += 1

        mock_api.create_commit.side_effect = count_commits

        push_dataset(
            repo_id="user/batch1",
            model_name=TEST_MODEL,
            prompts=populated_cache,
            labels=[1, 1, 0],
            exist_ok=True,
            stream=True,
            stream_batch_size=1,
        )

        # With batch_size=1, each shard gets its own commit.
        # 2 layers × 1 shard each = 2 shard commits
        assert shard_commit_count >= 2

    @requires_pyarrow
    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    @patch("huggingface_hub.HfApi")
    def test_stream_large_batch_fewer_commits(
        self, MockHfApi, mock_pyarrow, mock_deps, populated_cache,
    ):
        """Large batch_size groups all shards into fewer commits."""
        mock_api = MagicMock()
        MockHfApi.return_value = mock_api

        shard_commits = []

        def track_commits(*, repo_id, operations, commit_message="", **kw):
            if "shard" in commit_message.lower():
                shard_commits.append(len(operations))

        mock_api.create_commit.side_effect = track_commits

        push_dataset(
            repo_id="user/batch-large",
            model_name=TEST_MODEL,
            prompts=populated_cache,
            labels=[1, 1, 0],
            exist_ok=True,
            stream=True,
            stream_batch_size=100,  # larger than total shards
        )

        # All shards fit in one batch → single shard commit
        assert len(shard_commits) == 1
        assert shard_commits[0] >= 2  # at least 2 shards in the batch


class TestDryConsolidation:
    """Tests for dry consolidation: shard plan, remote check, skip logic."""

    def test_compute_shard_plan_returns_metadata(self, populated_cache):
        """_compute_shard_plan returns prompt_metadata and tensor_descriptors."""
        tensor_types = {
            "pooled": {"last_token": [0, 1]},
            "has_logits": False,
            "logits_top_k": None,
            "raw_layers": [],
        }
        kept_indices = list(range(len(populated_cache)))

        plan = _compute_shard_plan(
            model_name=TEST_MODEL,
            prompts=populated_cache,
            kept_indices=kept_indices,
            tensor_types=tensor_types,
            labels=None,
            shard_max_bytes=1_073_741_824,
            repo_id="user/test-plan",
        )

        assert len(plan["prompt_metadata"]) == len(populated_cache)
        assert len(plan["valid_prompts"]) == len(populated_cache)
        assert "hidden_layers" in plan["tensor_descriptors"]
        assert plan["has_hidden"] is True
        assert plan["hidden_dim"] == HIDDEN_DIM
        assert len(plan["hidden_boundaries"]) > 0

    def test_enumerate_shard_files_hidden(self, populated_cache):
        """_enumerate_shard_files lists all expected hidden shard files."""
        plan = {
            "has_hidden": True,
            "hidden_layers": [0, 1],
            "hidden_boundaries": [3],  # one shard with 3 prompts
            "want_logits": False,
            "logits_boundaries": [],
        }
        files = _enumerate_shard_files(plan)
        assert files == [
            "tensors/hidden_layer000_shard000.safetensors",
            "tensors/hidden_layer001_shard000.safetensors",
        ]

    def test_enumerate_shard_files_with_logits(self):
        """_enumerate_shard_files includes logits shards."""
        plan = {
            "has_hidden": True,
            "hidden_layers": [0],
            "hidden_boundaries": [2, 1],  # 2 shards
            "want_logits": True,
            "logits_boundaries": [3],
        }
        files = _enumerate_shard_files(plan)
        assert "tensors/hidden_layer000_shard000.safetensors" in files
        assert "tensors/hidden_layer000_shard001.safetensors" in files
        assert "tensors/logits_topk_000.safetensors" in files

    def test_check_shards_on_remote_finds_existing(self):
        """_check_shards_on_remote returns shards found on remote."""
        mock_api = MagicMock()
        mock_sibling = type("Sibling", (), {})
        s1 = mock_sibling()
        s1.rfilename = "tensors/hidden_layer000_shard000.safetensors"
        s2 = mock_sibling()
        s2.rfilename = "README.md"
        mock_repo_info = MagicMock()
        mock_repo_info.siblings = [s1, s2]
        mock_api.repo_info.return_value = mock_repo_info

        shard_files = [
            "tensors/hidden_layer000_shard000.safetensors",
            "tensors/hidden_layer001_shard000.safetensors",
        ]
        existing = _check_shards_on_remote(mock_api, "user/test", shard_files)
        assert existing == {"tensors/hidden_layer000_shard000.safetensors"}

    def test_check_shards_on_remote_handles_error(self):
        """_check_shards_on_remote returns empty set on API error."""
        mock_api = MagicMock()
        mock_api.repo_info.side_effect = Exception("repo not found")

        existing = _check_shards_on_remote(
            mock_api, "user/nonexistent", ["tensors/shard.safetensors"]
        )
        assert existing == set()

    def test_reconstruct_plan_from_cached(self):
        """_reconstruct_plan_from_cached produces valid plan for enumeration."""
        cached = {
            "tensor_descriptors": {
                "hidden_layers": {
                    "layers": [0, 1, 2],
                    "shards": [
                        {"num_prompts": 100},
                        {"num_prompts": 50},
                    ],
                },
                "logits_topk": {
                    "shards": [{"num_prompts": 150}],
                },
            },
            "prompt_metadata": [],
        }
        plan = _reconstruct_plan_from_cached(cached)
        assert plan["has_hidden"] is True
        assert plan["hidden_layers"] == [0, 1, 2]
        assert plan["hidden_boundaries"] == [100, 50]
        assert plan["want_logits"] is True
        assert plan["logits_boundaries"] == [150]

        files = _enumerate_shard_files(plan)
        assert len(files) == 3 * 2 + 1  # 3 layers × 2 shards + 1 logits

    def test_consolidate_skip_shards(self, populated_cache):
        """_consolidate_and_shard with skip_shards skips tensor loading."""
        tensor_types = {
            "pooled": {"last_token": [0, 1]},
            "has_logits": False,
            "logits_top_k": None,
            "raw_layers": [],
        }
        kept_indices = list(range(len(populated_cache)))

        # First run: get the full plan to know shard names
        plan = _compute_shard_plan(
            model_name=TEST_MODEL,
            prompts=populated_cache,
            kept_indices=kept_indices,
            tensor_types=tensor_types,
            labels=None,
            shard_max_bytes=1_073_741_824,
            repo_id="user/skip-test",
        )
        all_files = set(_enumerate_shard_files(plan))

        written_files = []

        def track_written(local_path, repo_path):
            written_files.append(repo_path)

        # Run consolidation with all shards skipped
        _consolidate_and_shard(
            model_name=TEST_MODEL,
            prompts=populated_cache,
            kept_indices=kept_indices,
            tensor_types=tensor_types,
            labels=None,
            shard_max_bytes=1_073_741_824,
            repo_id="user/skip-test",
            on_shard_written=track_written,
            skip_shards=all_files,
        )

        # No shards should have been written
        assert written_files == []

    @requires_pyarrow
    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    @patch("huggingface_hub.HfApi")
    def test_stream_partial_resume(
        self, MockHfApi, mock_pyarrow, mock_deps, populated_cache,
    ):
        """Partial resume: only missing shards are consolidated."""
        mock_api = MagicMock()
        MockHfApi.return_value = mock_api

        # Mock remote has only layer 0 shard, not layer 1
        mock_sibling = type("Sibling", (), {})
        s1 = mock_sibling()
        s1.rfilename = "tensors/hidden_layer000_shard000.safetensors"
        mock_repo_info = MagicMock()
        mock_repo_info.siblings = [s1]
        mock_api.repo_info.return_value = mock_repo_info

        commit_paths = []

        def capture_commit(*, repo_id, operations, **kwargs):
            for op in operations:
                commit_paths.append(op.path_in_repo)

        mock_api.create_commit.side_effect = capture_commit

        push_dataset(
            repo_id="user/partial-resume",
            model_name=TEST_MODEL,
            prompts=populated_cache,
            labels=[1, 1, 0],
            exist_ok=True,
            stream=True,
        )

        # Layer 0 shard should NOT be in commits (already on remote)
        shard_paths = [
            p for p in commit_paths if p.startswith("tensors/")
        ]
        assert "tensors/hidden_layer000_shard000.safetensors" not in shard_paths
        # Layer 1 shard SHOULD be uploaded
        assert "tensors/hidden_layer001_shard000.safetensors" in shard_paths


@requires_pyarrow
class TestLoadActivationDataset:
    def _setup_remote_files(self, tmp_path):
        """Create remote files in v2 per-layer format with embedded metadata."""
        import pyarrow as pa
        import pyarrow.parquet as pq
        from safetensors.torch import save_file

        from lmprobe.sharing import _embed_metadata_in_schema

        (tmp_path / "tensors").mkdir(parents=True)
        (tmp_path / "index").mkdir(parents=True)

        tensor = torch.randn(3, HIDDEN_DIM)
        save_file(
            {"hidden.layer_0": tensor},
            str(tmp_path / "tensors" / "hidden_layer000_shard000.safetensors"),
        )

        lmprobe_info = {
            "format_version": "2.0",
            "model": {"name": TEST_MODEL, "revision": None},
            "num_prompts": 3,
            "prompt_ordering": "random",
            "tensors": {
                "hidden_layers": {
                    "type": "hidden",
                    "layers": [0],
                    "dim": HIDDEN_DIM,
                    "dtype": "float32",
                    "layout": "per_layer",
                    "pooling": "last_token",
                    "row_bytes": HIDDEN_DIM * 4,
                    "shards": [
                        {
                            "num_prompts": 3,
                        }
                    ],
                }
            },
            "provenance": {},
        }

        # Write Parquet index with embedded metadata (v2 format)
        table = pa.table({
            "text": pa.array(["p0", "p1", "p2"], type=pa.string()),
            "label": pa.array([None, None, None], type=pa.int32()),
            "num_tokens": pa.array([5, 5, 5], type=pa.int32()),
            "shard_index": pa.array([0, 0, 0], type=pa.int32()),
            "row_offset": pa.array([0, 1, 2], type=pa.int32()),
        })
        table = _embed_metadata_in_schema(table, lmprobe_info)
        pq.write_table(table, str(tmp_path / PARQUET_PATH))

        return tmp_path, tensor

    @patch("lmprobe.sharing._check_hub_deps")
    def test_load_activation_dataset(self, mock_deps, tmp_path):
        remote_dir, expected_tensor = self._setup_remote_files(tmp_path)

        def mock_download(repo_id, filename, **kwargs):
            return str(remote_dir / filename)

        with patch(
            "huggingface_hub.hf_hub_download", side_effect=mock_download,
        ):
            result, info = load_activation_dataset("user/test-dataset")

        assert "hidden.layer_0" in result
        assert result["hidden.layer_0"].shape == (3, HIDDEN_DIM)
        assert torch.allclose(result["hidden.layer_0"], expected_tensor)

    @patch("lmprobe.sharing._check_hub_deps")
    def test_load_selective(self, mock_deps, tmp_path):
        remote_dir, _ = self._setup_remote_files(tmp_path)

        def mock_download(repo_id, filename, **kwargs):
            return str(remote_dir / filename)

        with patch(
            "huggingface_hub.hf_hub_download", side_effect=mock_download,
        ):
            result, info = load_activation_dataset(
                "user/test", tensors=["hidden_layers"]
            )

        assert "hidden.layer_0" in result

    @patch("lmprobe.sharing._check_hub_deps")
    def test_version_mismatch_raises(self, mock_deps, tmp_path):
        import pyarrow as pa
        import pyarrow.parquet as pq

        from lmprobe.sharing import _embed_metadata_in_schema

        (tmp_path / "index").mkdir(parents=True)
        lmprobe_info = {
            "format_version": "3.0",
            "model": {"name": TEST_MODEL, "revision": None},
            "num_prompts": 1,
            "tensors": {},
            "provenance": {},
        }
        table = pa.table({
            "text": pa.array(["p0"], type=pa.string()),
            "label": pa.array([None], type=pa.int32()),
            "num_tokens": pa.array([5], type=pa.int32()),
            "shard_index": pa.array([0], type=pa.int32()),
            "row_offset": pa.array([0], type=pa.int32()),
        })
        table = _embed_metadata_in_schema(table, lmprobe_info)
        pq.write_table(table, str(tmp_path / PARQUET_PATH))

        def mock_download(repo_id, filename, **kwargs):
            return str(tmp_path / filename)

        with patch(
            "huggingface_hub.hf_hub_download", side_effect=mock_download,
        ):
            with pytest.raises(
                ValueError, match="Incompatible format version",
            ):
                load_activation_dataset("user/test")


@requires_pyarrow
class TestPullDataset:
    def _setup_remote_files(self, tmp_path):
        """Create remote files in v2 per-layer format with embedded metadata."""
        import pyarrow as pa
        import pyarrow.parquet as pq
        from safetensors.torch import save_file

        from lmprobe.sharing import _embed_metadata_in_schema

        (tmp_path / "tensors").mkdir(parents=True)
        (tmp_path / "index").mkdir(parents=True)

        prompts = ["prompt 0", "prompt 1", "prompt 2"]
        tensor = torch.randn(3, HIDDEN_DIM)
        save_file(
            {"hidden.layer_0": tensor},
            str(tmp_path / "tensors" / "hidden_layer000_shard000.safetensors"),
        )

        lmprobe_info = {
            "format_version": "2.0",
            "model": {"name": TEST_MODEL, "revision": None},
            "num_prompts": 3,
            "prompt_ordering": "random",
            "tensors": {
                "hidden_layers": {
                    "type": "hidden",
                    "layers": [0],
                    "dim": HIDDEN_DIM,
                    "dtype": "float32",
                    "layout": "per_layer",
                    "pooling": "last_token",
                    "row_bytes": HIDDEN_DIM * 4,
                    "shards": [
                        {
                            "num_prompts": 3,
                        }
                    ],
                }
            },
            "provenance": {},
        }

        # Write Parquet index with embedded metadata (v2)
        table = pa.table({
            "text": pa.array(prompts, type=pa.string()),
            "label": pa.array([None, None, None], type=pa.int32()),
            "num_tokens": pa.array([5, 5, 5], type=pa.int32()),
            "shard_index": pa.array([0, 0, 0], type=pa.int32()),
            "row_offset": pa.array([0, 1, 2], type=pa.int32()),
        })
        table = _embed_metadata_in_schema(table, lmprobe_info)
        pq.write_table(table, str(tmp_path / PARQUET_PATH))

        return tmp_path, prompts

    @patch("lmprobe.sharing._check_hub_deps")
    def test_pull_populates_cache(self, mock_deps, tmp_path, cache_dir):
        remote_dir, prompts = self._setup_remote_files(tmp_path / "remote")

        def mock_download(repo_id, filename, **kwargs):
            return str(remote_dir / filename)

        with patch(
            "huggingface_hub.hf_hub_download", side_effect=mock_download,
        ):
            count = pull_dataset("user/test", materialize=True)

        assert count == 3

        for prompt in prompts:
            info = discover_cached(TEST_MODEL, prompt)
            assert info is not None

    @patch("lmprobe.sharing._check_hub_deps")
    def test_pull_dedup(self, mock_deps, tmp_path, cache_dir):
        remote_dir, prompts = self._setup_remote_files(tmp_path / "remote")

        # Pre-cache one prompt
        pooled = torch.randn(1, HIDDEN_DIM)
        save_prompt_pooled_activations(
            TEST_MODEL, "prompt 0", [0], pooled, "last_token"
        )

        def mock_download(repo_id, filename, **kwargs):
            return str(remote_dir / filename)

        with patch(
            "huggingface_hub.hf_hub_download", side_effect=mock_download,
        ):
            count = pull_dataset("user/test", overwrite=False, materialize=True)

        assert count == 2


@requires_pyarrow
class TestLazyPullDataset:
    """Tests for lazy (default) pull_dataset behavior."""

    def _setup_remote_files(self, tmp_path):
        """Create remote files in v2 per-layer format with embedded metadata."""
        import pyarrow as pa
        import pyarrow.parquet as pq
        from safetensors.torch import save_file

        from lmprobe.sharing import _embed_metadata_in_schema

        (tmp_path / "tensors").mkdir(parents=True)
        (tmp_path / "index").mkdir(parents=True)

        prompts = ["prompt 0", "prompt 1", "prompt 2"]
        tensor = torch.randn(3, HIDDEN_DIM)
        save_file(
            {"hidden.layer_0": tensor},
            str(tmp_path / "tensors" / "hidden_layer000_shard000.safetensors"),
        )

        lmprobe_info = {
            "format_version": "2.0",
            "model": {"name": TEST_MODEL, "revision": None},
            "num_prompts": 3,
            "prompt_ordering": "random",
            "tensors": {
                "hidden_layers": {
                    "type": "hidden",
                    "layers": [0],
                    "dim": HIDDEN_DIM,
                    "dtype": "float32",
                    "layout": "per_layer",
                    "pooling": "last_token",
                    "row_bytes": HIDDEN_DIM * 4,
                    "shards": [
                        {
                            "num_prompts": 3,
                        }
                    ],
                }
            },
            "provenance": {},
        }

        table = pa.table({
            "text": pa.array(prompts, type=pa.string()),
            "label": pa.array([None, None, None], type=pa.int32()),
            "num_tokens": pa.array([5, 5, 5], type=pa.int32()),
            "shard_index": pa.array([0, 0, 0], type=pa.int32()),
            "row_offset": pa.array([0, 1, 2], type=pa.int32()),
        })
        table = _embed_metadata_in_schema(table, lmprobe_info)
        pq.write_table(table, str(tmp_path / PARQUET_PATH))

        return tmp_path, prompts, tensor

    @patch("lmprobe.sharing._check_hub_deps")
    def test_pull_lazy_default(self, mock_deps, tmp_path, cache_dir):
        """Default pull creates no per-prompt files, but manifest exists."""
        remote_dir, prompts, _ = self._setup_remote_files(tmp_path / "remote")

        def mock_download(repo_id, filename, **kwargs):
            return str(remote_dir / filename)

        with patch(
            "huggingface_hub.hf_hub_download", side_effect=mock_download,
        ):
            count = pull_dataset("user/test")

        assert count == 3

        # No per-prompt safetensors files should exist
        from lmprobe.cache import _hash_string, get_cache_dir
        model_hash = _hash_string(TEST_MODEL)
        model_dir = get_cache_dir() / model_hash
        per_prompt_files = list(model_dir.glob("*.safetensors"))
        assert len(per_prompt_files) == 0

        # But manifest and index should exist
        repo_hash = _hash_string("user/test")
        assert (model_dir / f"_shard_manifest_{repo_hash}.json").exists()
        assert (model_dir / "_shard_index.json").exists()

    @patch("lmprobe.sharing._check_hub_deps")
    def test_pull_lazy_discover_cached(self, mock_deps, tmp_path, cache_dir):
        """discover_cached returns correct info for lazy-pulled prompts."""
        remote_dir, prompts, _ = self._setup_remote_files(tmp_path / "remote")

        def mock_download(repo_id, filename, **kwargs):
            return str(remote_dir / filename)

        with patch(
            "huggingface_hub.hf_hub_download", side_effect=mock_download,
        ):
            pull_dataset("user/test")

        for prompt in prompts:
            info = discover_cached(TEST_MODEL, prompt)
            assert info is not None, f"Prompt not discoverable: {prompt}"
            assert "last_token" in info.pooled
            assert 0 in info.pooled["last_token"]

    @patch("lmprobe.sharing._check_hub_deps")
    def test_pull_lazy_load_pooled(self, mock_deps, tmp_path, cache_dir):
        """load_prompt_pooled_activations returns correct data from shard."""
        remote_dir, prompts, tensor = self._setup_remote_files(
            tmp_path / "remote"
        )

        def mock_download(repo_id, filename, **kwargs):
            return str(remote_dir / filename)

        with patch(
            "huggingface_hub.hf_hub_download", side_effect=mock_download,
        ):
            pull_dataset("user/test")

        for i, prompt in enumerate(prompts):
            pooled = load_prompt_pooled_activations(
                TEST_MODEL, prompt, [0], "last_token"
            )
            assert pooled.shape == (1, HIDDEN_DIM)
            assert torch.allclose(pooled, tensor[i : i + 1], atol=1e-6)

    @patch("lmprobe.sharing._check_hub_deps")
    def test_pull_lazy_dedup_prefers_per_prompt(
        self, mock_deps, tmp_path, cache_dir,
    ):
        """Per-prompt cache takes priority over shard data."""
        remote_dir, prompts, tensor = self._setup_remote_files(
            tmp_path / "remote"
        )

        # Pre-cache prompt 0 with different data
        custom_pooled = torch.randn(1, HIDDEN_DIM)
        save_prompt_pooled_activations(
            TEST_MODEL, "prompt 0", [0], custom_pooled, "last_token"
        )

        def mock_download(repo_id, filename, **kwargs):
            return str(remote_dir / filename)

        with patch(
            "huggingface_hub.hf_hub_download", side_effect=mock_download,
        ):
            pull_dataset("user/test")

        # prompt 0 should return the per-prompt data, not shard data
        loaded = load_prompt_pooled_activations(
            TEST_MODEL, "prompt 0", [0], "last_token"
        )
        assert torch.allclose(loaded, custom_pooled, atol=1e-6)

    @patch("lmprobe.sharing._check_hub_deps")
    def test_pull_materialize_true(self, mock_deps, tmp_path, cache_dir):
        """materialize=True creates per-prompt files (old behavior)."""
        remote_dir, prompts, _ = self._setup_remote_files(tmp_path / "remote")

        def mock_download(repo_id, filename, **kwargs):
            return str(remote_dir / filename)

        with patch(
            "huggingface_hub.hf_hub_download", side_effect=mock_download,
        ):
            count = pull_dataset("user/test", materialize=True)

        assert count == 3

        # Per-prompt safetensors files should exist
        from lmprobe.cache import _hash_string, get_cache_dir
        model_hash = _hash_string(TEST_MODEL)
        model_dir = get_cache_dir() / model_hash
        per_prompt_files = list(model_dir.glob("*.safetensors"))
        assert len(per_prompt_files) == 3

    @patch("lmprobe.sharing._check_hub_deps")
    def test_pull_materialize_parallel(self, mock_deps, tmp_path, cache_dir):
        """num_workers=2 produces same results as sequential."""
        remote_dir, prompts, tensor = self._setup_remote_files(
            tmp_path / "remote"
        )

        def mock_download(repo_id, filename, **kwargs):
            return str(remote_dir / filename)

        with patch(
            "huggingface_hub.hf_hub_download", side_effect=mock_download,
        ):
            count = pull_dataset(
                "user/test", materialize=True, num_workers=2,
            )

        assert count == 3
        for i, prompt in enumerate(prompts):
            pooled = load_prompt_pooled_activations(
                TEST_MODEL, prompt, [0], "last_token"
            )
            assert pooled.shape == (1, HIDDEN_DIM)
            assert torch.allclose(pooled, tensor[i : i + 1], atol=1e-6)

    def _setup_raw_remote_files(self, tmp_path):
        """Create remote files with full_sequence storage (v2 per-layer)."""
        import pyarrow as pa
        import pyarrow.parquet as pq
        from safetensors.torch import save_file

        from lmprobe.sharing import _embed_metadata_in_schema

        (tmp_path / "tensors").mkdir(parents=True)
        (tmp_path / "index").mkdir(parents=True)

        prompts = ["raw prompt 0", "raw prompt 1"]
        seq_lens = [3, 5]
        layers = [0, 1]

        # Build per-layer tensors: concat all prompts' tokens
        total_tokens = sum(seq_lens)
        layer_tensors = {}
        for layer in layers:
            t = torch.randn(total_tokens, HIDDEN_DIM)
            layer_tensors[f"hidden.layer_{layer}"] = t
            save_file(
                {f"hidden.layer_{layer}": t},
                str(
                    tmp_path / "tensors"
                    / f"hidden_layer{layer:03d}_shard000.safetensors"
                ),
            )

        lmprobe_info = {
            "format_version": "2.0",
            "model": {"name": TEST_MODEL, "revision": None},
            "num_prompts": 2,
            "prompt_ordering": "sequential",
            "tensors": {
                "hidden_layers": {
                    "type": "hidden",
                    "layers": layers,
                    "dim": HIDDEN_DIM,
                    "dtype": "float32",
                    "storage": "full_sequence",
                    "layout": "per_layer",
                    "pooling": "last_token",
                    "row_bytes": HIDDEN_DIM * 4,
                    "shards": [
                        {
                            "num_prompts": 2,
                        }
                    ],
                }
            },
            "provenance": {},
        }

        token_offsets = [0, seq_lens[0]]
        table = pa.table({
            "text": pa.array(prompts, type=pa.string()),
            "label": pa.array([None, None], type=pa.int32()),
            "num_tokens": pa.array(seq_lens, type=pa.int32()),
            "shard_index": pa.array([0, 0], type=pa.int32()),
            "row_offset": pa.array([0, 0], type=pa.int32()),
            "token_offset": pa.array(token_offsets, type=pa.int32()),
        })
        table = _embed_metadata_in_schema(table, lmprobe_info)
        pq.write_table(table, str(tmp_path / PARQUET_PATH))

        return tmp_path, prompts, seq_lens, layers, layer_tensors

    @patch("lmprobe.sharing._check_hub_deps")
    def test_pull_lazy_load_raw(self, mock_deps, tmp_path, cache_dir):
        """load_prompt_activations returns correct data from shard."""
        remote_dir, prompts, seq_lens, layers, layer_tensors = (
            self._setup_raw_remote_files(tmp_path / "remote")
        )

        def mock_download(repo_id, filename, **kwargs):
            return str(remote_dir / filename)

        with patch(
            "huggingface_hub.hf_hub_download", side_effect=mock_download,
        ):
            pull_dataset("user/test")

        token_offsets = [0, seq_lens[0]]
        for i, prompt in enumerate(prompts):
            act, mask = load_prompt_activations(TEST_MODEL, prompt, layers)
            assert act.shape == (1, seq_lens[i], HIDDEN_DIM * len(layers))
            assert mask.shape == (1, seq_lens[i])

            # Verify data matches
            tok_off = token_offsets[i]
            num_tok = seq_lens[i]
            expected_slices = []
            for layer in layers:
                expected_slices.append(
                    layer_tensors[f"hidden.layer_{layer}"][
                        tok_off : tok_off + num_tok
                    ]
                )
            expected = torch.cat(expected_slices, dim=-1).unsqueeze(0)
            assert torch.allclose(act, expected, atol=1e-6)


@requires_pyarrow
class TestRoundtrip:
    """Push from cache -> pull into fresh cache -> verify data matches."""

    def test_push_pull_roundtrip_pooled_activations(
        self, tmp_path, monkeypatch,
    ):
        """Roundtrip: pooled activations survive push -> pull."""
        # --- Phase 1: populate source cache ---
        src_cache = tmp_path / "src_cache"
        monkeypatch.setenv("LMPROBE_CACHE_DIR", str(src_cache))
        cache_mod._backend = None

        prompts = [
            "Who wants to go for a walk?",
            "Fetch the ball!",
            "Purring and scratching",
        ]
        labels = [1, 1, 0]

        original_pooled: dict[str, dict[int, torch.Tensor]] = {}
        for prompt in prompts:
            original_pooled[prompt] = {}
            for layer in [0, 1]:
                pooled = torch.randn(1, HIDDEN_DIM)
                save_prompt_pooled_activations(
                    TEST_MODEL, prompt, [layer], pooled, "last_token"
                )
                original_pooled[prompt][layer] = pooled

        # --- Phase 2: push (capture the uploaded folder) ---
        remote_dir = tmp_path / "remote"
        remote_dir.mkdir()

        mock_api = MagicMock()

        def capture_upload(repo_id, folder_path, **kwargs):
            src = Path(folder_path)
            for f in src.rglob("*"):
                if f.is_file():
                    rel = f.relative_to(src)
                    dest = remote_dir / rel
                    dest.parent.mkdir(parents=True, exist_ok=True)
                    shutil.copy2(str(f), str(dest))

        mock_api.upload_large_folder.side_effect = capture_upload

        with (
            patch("lmprobe.sharing._check_hub_deps"),
            patch("huggingface_hub.HfApi", return_value=mock_api),
        ):
            push_dataset(
                repo_id="user/roundtrip-test",
                model_name=TEST_MODEL,
                prompts=prompts,
                labels=labels,
                exist_ok=True,
            )

        # Verify remote files
        assert (remote_dir / INFO_FILENAME).exists()
        assert (remote_dir / PARQUET_PATH).exists()

        with open(remote_dir / INFO_FILENAME) as f:
            info = json.load(f)
        assert info["num_prompts"] == 3
        assert "hidden_layers" in info["tensors"]

        # --- Phase 3: pull into a fresh cache ---
        dst_cache = tmp_path / "dst_cache"
        monkeypatch.setenv("LMPROBE_CACHE_DIR", str(dst_cache))
        cache_mod._backend = None

        def mock_download(repo_id, filename, **kwargs):
            return str(remote_dir / filename)

        with (
            patch("lmprobe.sharing._check_hub_deps"),
            patch(
                "huggingface_hub.hf_hub_download",
                side_effect=mock_download,
            ),
        ):
            count = pull_dataset("user/roundtrip-test", materialize=True)

        assert count == 3

        # --- Phase 4: verify pulled data matches originals ---
        for prompt in prompts:
            info = discover_cached(TEST_MODEL, prompt)
            assert info is not None, (
                f"Prompt not cached after pull: {prompt}"
            )
            assert "last_token" in info.pooled

            for layer in [0, 1]:
                pulled = load_prompt_pooled_activations(
                    TEST_MODEL, prompt, [layer], "last_token"
                )
                orig = original_pooled[prompt][layer]

                assert pulled.shape == orig.shape, (
                    f"Shape mismatch for {prompt!r} layer {layer}"
                )
                assert torch.allclose(
                    pulled, orig, atol=1e-6
                ), f"Value mismatch for {prompt!r} layer {layer}"

    def test_push_load_roundtrip(self, tmp_path, monkeypatch):
        """Push -> load_activation_dataset returns concatenated tensors."""
        src_cache = tmp_path / "src_cache"
        monkeypatch.setenv("LMPROBE_CACHE_DIR", str(src_cache))
        cache_mod._backend = None

        prompts = ["alpha", "beta", "gamma"]
        all_pooled_layer0 = {}
        for prompt in prompts:
            pooled = torch.randn(1, HIDDEN_DIM)
            save_prompt_pooled_activations(
                TEST_MODEL, prompt, [0], pooled, "last_token"
            )
            all_pooled_layer0[prompt] = pooled

        # Push
        remote_dir = tmp_path / "remote"
        remote_dir.mkdir()
        mock_api = MagicMock()

        def capture_upload(repo_id, folder_path, **kwargs):
            for f in Path(folder_path).rglob("*"):
                if f.is_file():
                    rel = f.relative_to(folder_path)
                    dest = remote_dir / rel
                    dest.parent.mkdir(parents=True, exist_ok=True)
                    shutil.copy2(str(f), str(dest))

        mock_api.upload_large_folder.side_effect = capture_upload

        with (
            patch("lmprobe.sharing._check_hub_deps"),
            patch("huggingface_hub.HfApi", return_value=mock_api),
        ):
            push_dataset(
                repo_id="user/load-test",
                model_name=TEST_MODEL,
                prompts=prompts,
                exist_ok=True,
            )

        # Load directly
        def mock_download(repo_id, filename, **kwargs):
            return str(remote_dir / filename)

        with (
            patch("lmprobe.sharing._check_hub_deps"),
            patch(
                "huggingface_hub.hf_hub_download",
                side_effect=mock_download,
            ),
        ):
            result, info = load_activation_dataset("user/load-test")

        assert "hidden.layer_0" in result
        loaded = result["hidden.layer_0"]
        assert loaded.shape == (3, HIDDEN_DIM)

        # Verify all values are present (order may differ due to shuffle)
        for prompt in prompts:
            orig = all_pooled_layer0[prompt].squeeze(0)
            # Find this vector in loaded
            found = False
            for row_idx in range(loaded.shape[0]):
                if torch.allclose(loaded[row_idx], orig, atol=1e-6):
                    found = True
                    break
            assert found, f"Original tensor for {prompt!r} not found in loaded"

    def test_parquet_standalone_readable(self, tmp_path, monkeypatch):
        """Verify Parquet index is standalone-readable with pyarrow."""
        import pyarrow.parquet as pq

        src_cache = tmp_path / "src_cache"
        monkeypatch.setenv("LMPROBE_CACHE_DIR", str(src_cache))
        cache_mod._backend = None

        prompts = ["hello", "world"]
        for prompt in prompts:
            pooled = torch.randn(1, HIDDEN_DIM)
            save_prompt_pooled_activations(
                TEST_MODEL, prompt, [0], pooled, "last_token"
            )

        remote_dir = tmp_path / "remote"
        remote_dir.mkdir()
        mock_api = MagicMock()

        def capture_upload(repo_id, folder_path, **kwargs):
            for f in Path(folder_path).rglob("*"):
                if f.is_file():
                    rel = f.relative_to(folder_path)
                    dest = remote_dir / rel
                    dest.parent.mkdir(parents=True, exist_ok=True)
                    shutil.copy2(str(f), str(dest))

        mock_api.upload_large_folder.side_effect = capture_upload

        with (
            patch("lmprobe.sharing._check_hub_deps"),
            patch("huggingface_hub.HfApi", return_value=mock_api),
        ):
            push_dataset(
                repo_id="user/parquet-test",
                model_name=TEST_MODEL,
                prompts=prompts,
                labels=[1, 0],
                exist_ok=True,
            )

        # Read Parquet directly
        table = pq.read_table(str(remote_dir / PARQUET_PATH))

        assert table.num_rows == 2
        expected_cols = {
            "text", "label", "num_tokens", "shard_index", "row_offset",
        }
        assert set(table.column_names) == expected_cols
        assert set(table.column("text").to_pylist()) == set(prompts)
        assert set(table.column("label").to_pylist()) == {0, 1}

    def test_safetensors_standalone_readable(self, tmp_path, monkeypatch):
        """Verify per-layer safetensors shards have single key each."""
        from safetensors import safe_open

        src_cache = tmp_path / "src_cache"
        monkeypatch.setenv("LMPROBE_CACHE_DIR", str(src_cache))
        cache_mod._backend = None

        prompts = ["alpha", "beta"]
        for prompt in prompts:
            for layer in [0, 1]:
                pooled = torch.randn(1, HIDDEN_DIM)
                save_prompt_pooled_activations(
                    TEST_MODEL, prompt, [layer], pooled, "last_token"
                )

        remote_dir = tmp_path / "remote"
        remote_dir.mkdir()
        mock_api = MagicMock()

        def capture_upload(repo_id, folder_path, **kwargs):
            for f in Path(folder_path).rglob("*"):
                if f.is_file():
                    rel = f.relative_to(folder_path)
                    dest = remote_dir / rel
                    dest.parent.mkdir(parents=True, exist_ok=True)
                    shutil.copy2(str(f), str(dest))

        mock_api.upload_large_folder.side_effect = capture_upload

        with (
            patch("lmprobe.sharing._check_hub_deps"),
            patch("huggingface_hub.HfApi", return_value=mock_api),
        ):
            push_dataset(
                repo_id="user/sf-test",
                model_name=TEST_MODEL,
                prompts=prompts,
                exist_ok=True,
            )

        # Per-layer files: hidden_layer{L}_shard{S}.safetensors
        for layer in [0, 1]:
            shard_files = list(
                (remote_dir / "tensors").glob(
                    f"hidden_layer{layer:03d}_shard*.safetensors"
                )
            )
            assert len(shard_files) >= 1

            with safe_open(str(shard_files[0]), framework="pt") as f:
                keys = list(f.keys())
                assert keys == [f"hidden.layer_{layer}"]

                tensor = f.get_tensor(f"hidden.layer_{layer}")
                assert tensor.shape == (2, HIDDEN_DIM)


# =============================================================================
# Metadata tests
# =============================================================================


class TestConsolidationWithMetadata:
    def test_extra_keys_in_prompt_metadata(self, populated_cache):
        """Extra metadata keys appear in prompt_metadata dicts."""
        tensor_types = {
            "raw_layers": [],
            "pooled": {"last_token": [0]},
            "has_logits": False,
            "logits_top_k": None,
            "has_perplexity": False,
        }
        meta = [
            {"statement_id": "s0", "category": "cities"},
            {"statement_id": "s1", "category": "cities"},
            {"statement_id": "s2", "category": "animals"},
        ]
        _, _, prompt_metadata = _consolidate_and_shard(
            model_name=TEST_MODEL,
            prompts=populated_cache,
            kept_indices=[0, 1, 2],
            tensor_types=tensor_types,
            labels=None,
            shard_max_bytes=1_000_000_000,
            repo_id="user/test",
            metadata=meta,
        )
        # All prompts should have the extra keys (order may be shuffled)
        for pm in prompt_metadata:
            assert "statement_id" in pm
            assert "category" in pm
        # All original values should be present
        ids = {pm["statement_id"] for pm in prompt_metadata}
        assert ids == {"s0", "s1", "s2"}


@requires_pyarrow
class TestExtraMetadataColumns:
    def test_extra_columns_in_parquet(self, tmp_path):
        """Extra metadata keys become extra Parquet columns."""
        import pyarrow.parquet as pq

        (tmp_path / "index").mkdir()
        prompt_metadata = [
            {"text": "hello", "label": 1, "num_tokens": 3,
             "shard_index": 0, "row_offset": 0,
             "statement_id": "s0", "score": 0.95},
            {"text": "world", "label": 0, "num_tokens": 4,
             "shard_index": 0, "row_offset": 1,
             "statement_id": "s1", "score": 0.42},
        ]
        _write_parquet_index(tmp_path, prompt_metadata)

        table = pq.read_table(str(tmp_path / PARQUET_PATH))
        assert "statement_id" in table.column_names
        assert "score" in table.column_names
        assert table.column("statement_id").to_pylist() == ["s0", "s1"]
        assert table.column("score").to_pylist() == [
            pytest.approx(0.95), pytest.approx(0.42),
        ]

    def test_token_offset_column(self, tmp_path):
        """token_offset from full-sequence storage appears as extra column."""
        import pyarrow.parquet as pq

        (tmp_path / "index").mkdir()
        prompt_metadata = [
            {"text": "a", "label": 1, "num_tokens": 3,
             "shard_index": 0, "row_offset": 0, "token_offset": 0},
            {"text": "b", "label": 0, "num_tokens": 5,
             "shard_index": 0, "row_offset": 0, "token_offset": 3},
        ]
        _write_parquet_index(tmp_path, prompt_metadata)

        table = pq.read_table(str(tmp_path / PARQUET_PATH))
        assert "token_offset" in table.column_names
        assert table.column("token_offset").to_pylist() == [0, 3]

    def test_bool_metadata_columns(self, tmp_path):
        """Bool metadata columns should not crash the Parquet index write.

        Regression test for #155: isinstance(True, int) is True in Python,
        so bool values were incorrectly routed to pa.int32(), causing
        ArrowTypeError: Expected integer, got bool.
        """
        import pyarrow as pa
        import pyarrow.parquet as pq

        (tmp_path / "index").mkdir()
        prompt_metadata = [
            {"text": "hello", "label": 1, "num_tokens": 3,
             "shard_index": 0, "row_offset": 0, "negated": True},
            {"text": "world", "label": 0, "num_tokens": 4,
             "shard_index": 0, "row_offset": 1, "negated": False},
        ]
        _write_parquet_index(tmp_path, prompt_metadata)

        table = pq.read_table(str(tmp_path / PARQUET_PATH))
        assert "negated" in table.column_names
        assert table.column("negated").type == pa.bool_()
        assert table.column("negated").to_pylist() == [True, False]

    def test_bool_list_metadata_columns(self, tmp_path):
        """List-of-bool metadata columns should use pa.list_(pa.bool_())."""
        import pyarrow as pa
        import pyarrow.parquet as pq

        (tmp_path / "index").mkdir()
        prompt_metadata = [
            {"text": "a", "label": 1, "num_tokens": 3,
             "shard_index": 0, "row_offset": 0, "flags": [True, False]},
            {"text": "b", "label": 0, "num_tokens": 4,
             "shard_index": 0, "row_offset": 1, "flags": [False, True]},
        ]
        _write_parquet_index(tmp_path, prompt_metadata)

        table = pq.read_table(str(tmp_path / PARQUET_PATH))
        assert "flags" in table.column_names
        assert table.column("flags").type == pa.list_(pa.bool_())
        assert table.column("flags").to_pylist() == [[True, False], [False, True]]


@requires_pyarrow
class TestPushMetadataValidation:
    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    def test_metadata_length_mismatch(
        self, mock_pyarrow, mock_deps, populated_cache,
    ):
        with pytest.raises(ValueError, match="metadata length"):
            push_dataset(
                repo_id="user/test",
                model_name=TEST_MODEL,
                prompts=populated_cache,
                metadata=[{"a": 1}],  # length 1, prompts length 3
            )

    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    def test_metadata_inconsistent_keys(
        self, mock_pyarrow, mock_deps, populated_cache,
    ):
        with pytest.raises(ValueError, match="metadata"):
            push_dataset(
                repo_id="user/test",
                model_name=TEST_MODEL,
                prompts=populated_cache,
                metadata=[
                    {"a": 1},
                    {"a": 2},
                    {"b": 3},  # different keys
                ],
            )

    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    @patch("huggingface_hub.HfApi")
    def test_push_with_metadata_end_to_end(
        self, MockHfApi, mock_pyarrow, mock_deps, populated_cache,
    ):
        mock_api = MagicMock()
        MockHfApi.return_value = mock_api
        uploaded_files = {}

        def capture_upload(repo_id, folder_path, **kwargs):
            folder = Path(folder_path)
            for f in folder.rglob("*"):
                if f.is_file():
                    rel = f.relative_to(folder)
                    if f.suffix == ".json":
                        uploaded_files[str(rel)] = json.loads(f.read_text())
                    elif f.suffix == ".parquet":
                        uploaded_files[str(rel)] = f.read_bytes()
                    else:
                        uploaded_files[str(rel)] = f.read_bytes()

        mock_api.upload_large_folder.side_effect = capture_upload

        meta = [
            {"statement_id": "s0", "category": "cities"},
            {"statement_id": "s1", "category": "cities"},
            {"statement_id": "s2", "category": "animals"},
        ]

        push_dataset(
            repo_id="user/meta-test",
            model_name=TEST_MODEL,
            prompts=populated_cache,
            labels=[1, 1, 0],
            metadata=meta,
            exist_ok=True,
        )

        # Verify Parquet has extra columns
        import io

        import pyarrow.parquet as pq
        parquet_bytes = uploaded_files[PARQUET_PATH]
        table = pq.read_table(io.BytesIO(parquet_bytes))
        assert "statement_id" in table.column_names
        assert "category" in table.column_names


# =============================================================================
# Full-sequence activation tests
# =============================================================================


class TestComputeShardBoundariesVariable:
    def test_single_shard(self):
        boundaries = _compute_shard_boundaries_variable(
            [100, 200, 300], shard_max_bytes=1000
        )
        assert boundaries == [3]

    def test_multiple_shards(self):
        boundaries = _compute_shard_boundaries_variable(
            [100, 200, 300, 400], shard_max_bytes=350
        )
        # 100+200=300 fits, +300 would be 600 > 350 → split
        assert len(boundaries) >= 2
        assert sum(boundaries) == 4

    def test_each_prompt_own_shard(self):
        boundaries = _compute_shard_boundaries_variable(
            [100, 100, 100], shard_max_bytes=1
        )
        # Each prompt exceeds limit alone, so each goes in its own shard
        assert boundaries == [1, 1, 1]


class TestConsolidateRawActivations:
    def test_full_sequence_storage_flag(self, populated_raw_cache):
        """Consolidation with raw layers sets storage='full_sequence'."""
        prompts, seq_lens, layers, _ = populated_raw_cache
        tensor_types = {
            "raw_layers": layers,
            "pooled": {},
            "has_logits": False,
            "logits_top_k": None,
            "has_perplexity": False,
        }
        tmpdir, tensor_descriptors, prompt_metadata = _consolidate_and_shard(
            model_name=TEST_MODEL,
            prompts=prompts,
            kept_indices=[0, 1, 2],
            tensor_types=tensor_types,
            labels=None,
            shard_max_bytes=1_000_000_000,
            repo_id="user/raw-test",
        )

        desc = tensor_descriptors["hidden_layers"]
        assert desc["storage"] == "full_sequence"
        assert "pooling" not in desc
        assert "row_bytes" not in desc

        # Total tokens across shard should match sum of seq_lens
        total_shard_tokens = sum(
            s["num_tokens"] for s in desc["shards"]
        )
        assert total_shard_tokens == sum(seq_lens)

        shutil.rmtree(tmpdir, ignore_errors=True)

    def test_full_sequence_tensor_shape(self, populated_raw_cache):
        """Last-token shard has 1 token per prompt; rest shard has remainder."""
        from safetensors import safe_open

        prompts, seq_lens, layers, _ = populated_raw_cache
        tensor_types = {
            "raw_layers": layers,
            "pooled": {},
            "has_logits": False,
            "logits_top_k": None,
            "has_perplexity": False,
        }
        tmpdir, tensor_descriptors, _ = _consolidate_and_shard(
            model_name=TEST_MODEL,
            prompts=prompts,
            kept_indices=[0, 1, 2],
            tensor_types=tensor_types,
            labels=None,
            shard_max_bytes=1_000_000_000,
            repo_id="user/raw-test",
        )

        # Last-token shard (shard 0): one token per prompt
        lt_file = tmpdir / "tensors/hidden_layer000_shard000.safetensors"
        with safe_open(str(lt_file), framework="pt") as f:
            layer_0 = f.get_tensor("hidden.layer_0")
            assert layer_0.shape == (len(prompts), HIDDEN_DIM)

        # Total tokens across ALL shards equals sum(seq_lens)
        total_tokens = 0
        for si in range(len(tensor_descriptors["hidden_layers"]["shards"])):
            fpath = tmpdir / f"tensors/hidden_layer000_shard{si:03d}.safetensors"
            if fpath.exists():
                with safe_open(str(fpath), framework="pt") as f:
                    total_tokens += f.get_tensor("hidden.layer_0").shape[0]
        assert total_tokens == sum(seq_lens)

        shutil.rmtree(tmpdir, ignore_errors=True)

    def test_token_offset_in_metadata(self, populated_raw_cache):
        """prompt_metadata has token_offset instead of row_offset for raw."""
        prompts, seq_lens, layers, _ = populated_raw_cache
        tensor_types = {
            "raw_layers": layers,
            "pooled": {},
            "has_logits": False,
            "logits_top_k": None,
            "has_perplexity": False,
        }
        _, _, prompt_metadata = _consolidate_and_shard(
            model_name=TEST_MODEL,
            prompts=prompts,
            kept_indices=[0, 1, 2],
            tensor_types=tensor_types,
            labels=None,
            shard_max_bytes=1_000_000_000,
            repo_id="user/raw-test",
        )

        for pm in prompt_metadata:
            assert "token_offset" in pm
            assert "shard_index" in pm


@requires_pyarrow
class TestFullSequenceRoundtrip:
    """Push raw activations, pull into fresh cache, verify raw + pooled."""

    def test_push_pull_roundtrip_raw(self, tmp_path, monkeypatch):
        # --- Phase 1: populate source cache with raw activations ---
        src_cache = tmp_path / "src_cache"
        monkeypatch.setenv("LMPROBE_CACHE_DIR", str(src_cache))
        cache_mod._backend = None

        prompts = ["alpha", "beta", "gamma"]
        seq_lens = [3, 5, 4]
        layers = [0, 1]
        original_raw = {}

        for prompt, sl in zip(prompts, seq_lens):
            act = torch.randn(1, sl, HIDDEN_DIM * len(layers))
            mask = torch.ones(1, sl, dtype=torch.long)
            save_prompt_activations(TEST_MODEL, prompt, layers, act, mask)
            original_raw[prompt] = (act, mask, sl)

        # --- Phase 2: push ---
        remote_dir = tmp_path / "remote"
        remote_dir.mkdir()
        mock_api = MagicMock()

        def capture_upload(repo_id, folder_path, **kwargs):
            for f in Path(folder_path).rglob("*"):
                if f.is_file():
                    rel = f.relative_to(folder_path)
                    dest = remote_dir / rel
                    dest.parent.mkdir(parents=True, exist_ok=True)
                    shutil.copy2(str(f), str(dest))

        mock_api.upload_large_folder.side_effect = capture_upload

        with (
            patch("lmprobe.sharing._check_hub_deps"),
            patch("huggingface_hub.HfApi", return_value=mock_api),
        ):
            push_dataset(
                repo_id="user/raw-roundtrip",
                model_name=TEST_MODEL,
                prompts=prompts,
                labels=[1, 1, 0],
                exist_ok=True,
            )

        # Verify storage type in info
        with open(remote_dir / INFO_FILENAME) as f:
            info = json.load(f)
        assert info["tensors"]["hidden_layers"]["storage"] == "full_sequence"

        # --- Phase 3: pull into fresh cache ---
        dst_cache = tmp_path / "dst_cache"
        monkeypatch.setenv("LMPROBE_CACHE_DIR", str(dst_cache))
        cache_mod._backend = None

        def mock_download(repo_id, filename, **kwargs):
            return str(remote_dir / filename)

        with (
            patch("lmprobe.sharing._check_hub_deps"),
            patch(
                "huggingface_hub.hf_hub_download",
                side_effect=mock_download,
            ),
        ):
            count = pull_dataset("user/raw-roundtrip", materialize=True)

        assert count == 3

        # --- Phase 4: verify pulled data ---
        for prompt in prompts:
            cached = discover_cached(TEST_MODEL, prompt)
            assert cached is not None

            # Raw activations should be saved
            pulled_act, pulled_mask = load_prompt_activations(
                TEST_MODEL, prompt, layers,
            )
            orig_act, orig_mask, sl = original_raw[prompt]

            assert pulled_act.shape == orig_act.shape
            assert torch.allclose(pulled_act, orig_act, atol=1e-5)

            # Pooled (last_token) should also be saved
            pooled = load_prompt_pooled_activations(
                TEST_MODEL, prompt, [0], "last_token",
            )
            assert pooled.shape == (1, HIDDEN_DIM)

    def test_push_load_roundtrip_raw(self, tmp_path, monkeypatch):
        """Push raw -> load_activation_dataset returns correct shapes."""
        src_cache = tmp_path / "src_cache"
        monkeypatch.setenv("LMPROBE_CACHE_DIR", str(src_cache))
        cache_mod._backend = None

        prompts = ["alpha", "beta"]
        seq_lens = [3, 5]
        layers = [0, 1]

        for prompt, sl in zip(prompts, seq_lens):
            act = torch.randn(1, sl, HIDDEN_DIM * len(layers))
            mask = torch.ones(1, sl, dtype=torch.long)
            save_prompt_activations(TEST_MODEL, prompt, layers, act, mask)

        remote_dir = tmp_path / "remote"
        remote_dir.mkdir()
        mock_api = MagicMock()

        def capture_upload(repo_id, folder_path, **kwargs):
            for f in Path(folder_path).rglob("*"):
                if f.is_file():
                    rel = f.relative_to(folder_path)
                    dest = remote_dir / rel
                    dest.parent.mkdir(parents=True, exist_ok=True)
                    shutil.copy2(str(f), str(dest))

        mock_api.upload_large_folder.side_effect = capture_upload

        with (
            patch("lmprobe.sharing._check_hub_deps"),
            patch("huggingface_hub.HfApi", return_value=mock_api),
        ):
            push_dataset(
                repo_id="user/raw-load",
                model_name=TEST_MODEL,
                prompts=prompts,
                exist_ok=True,
            )

        def mock_download(repo_id, filename, **kwargs):
            return str(remote_dir / filename)

        with (
            patch("lmprobe.sharing._check_hub_deps"),
            patch(
                "huggingface_hub.hf_hub_download",
                side_effect=mock_download,
            ),
        ):
            result, info = load_activation_dataset("user/raw-load")

        # Total tokens = 3 + 5 = 8
        assert result["hidden.layer_0"].shape == (sum(seq_lens), HIDDEN_DIM)
        assert result["hidden.layer_1"].shape == (sum(seq_lens), HIDDEN_DIM)


@requires_pyarrow
class TestFullSequenceWithLogits:
    """Test combining full-sequence hidden layers + logits_topk."""

    def test_push_pull_raw_with_logits(self, tmp_path, monkeypatch):
        """Push raw hidden + logits, pull, verify both are correct."""
        src_cache = tmp_path / "src_cache"
        monkeypatch.setenv("LMPROBE_CACHE_DIR", str(src_cache))
        cache_mod._backend = None

        prompts = ["alpha", "beta", "gamma"]
        seq_lens = [3, 5, 4]
        layers = [0, 1]
        top_k = 10
        vocab_size = 50

        original_raw = {}

        for prompt, sl in zip(prompts, seq_lens):
            # Save raw activations
            act = torch.randn(1, sl, HIDDEN_DIM * len(layers))
            mask = torch.ones(1, sl, dtype=torch.long)
            save_prompt_activations(TEST_MODEL, prompt, layers, act, mask)
            original_raw[prompt] = act

            # Save top-k logits
            logits = torch.randn(1, sl, vocab_size)
            save_prompt_logits(
                TEST_MODEL, prompt, logits, mask,
                top_k=top_k, positions="last",
            )

        # Push
        remote_dir = tmp_path / "remote"
        remote_dir.mkdir()
        mock_api = MagicMock()

        def capture_upload(repo_id, folder_path, **kwargs):
            for f in Path(folder_path).rglob("*"):
                if f.is_file():
                    rel = f.relative_to(folder_path)
                    dest = remote_dir / rel
                    dest.parent.mkdir(parents=True, exist_ok=True)
                    shutil.copy2(str(f), str(dest))

        mock_api.upload_large_folder.side_effect = capture_upload

        with (
            patch("lmprobe.sharing._check_hub_deps"),
            patch("huggingface_hub.HfApi", return_value=mock_api),
        ):
            push_dataset(
                repo_id="user/raw-logits",
                model_name=TEST_MODEL,
                prompts=prompts,
                labels=[1, 1, 0],
                exist_ok=True,
            )

        # Verify both tensor types in info
        with open(remote_dir / INFO_FILENAME) as f:
            info = json.load(f)
        assert "hidden_layers" in info["tensors"]
        assert "logits_topk" in info["tensors"]
        assert info["tensors"]["hidden_layers"]["storage"] == "full_sequence"

        # Verify Parquet has both row_offset and token_offset
        import pyarrow.parquet as pq

        table = pq.read_table(str(remote_dir / PARQUET_PATH))
        assert "row_offset" in table.column_names
        assert "token_offset" in table.column_names

        # Pull into fresh cache
        dst_cache = tmp_path / "dst_cache"
        monkeypatch.setenv("LMPROBE_CACHE_DIR", str(dst_cache))
        cache_mod._backend = None

        def mock_download(repo_id, filename, **kwargs):
            return str(remote_dir / filename)

        with (
            patch("lmprobe.sharing._check_hub_deps"),
            patch(
                "huggingface_hub.hf_hub_download",
                side_effect=mock_download,
            ),
        ):
            count = pull_dataset("user/raw-logits", materialize=True)

        assert count == 3

        # Verify raw activations roundtrip
        for prompt in prompts:
            pulled_act, pulled_mask = load_prompt_activations(
                TEST_MODEL, prompt, layers,
            )
            orig_act = original_raw[prompt]
            assert pulled_act.shape == orig_act.shape
            assert torch.allclose(pulled_act, orig_act, atol=1e-5)


class TestPooledStorageUnchanged:
    """Existing pooled tests still pass — auto-detect falls back to pooled."""

    def test_pooled_storage_flag(self, populated_cache):
        tensor_types = {
            "raw_layers": [],
            "pooled": {"last_token": [0, 1]},
            "has_logits": False,
            "logits_top_k": None,
            "has_perplexity": False,
        }
        tmpdir, tensor_descriptors, prompt_metadata = _consolidate_and_shard(
            model_name=TEST_MODEL,
            prompts=populated_cache,
            kept_indices=[0, 1, 2],
            tensor_types=tensor_types,
            labels=None,
            shard_max_bytes=1_000_000_000,
            repo_id="user/pooled-test",
        )

        desc = tensor_descriptors["hidden_layers"]
        assert desc["storage"] == "pooled"
        assert desc["pooling"] == "last_token"
        assert "row_bytes" in desc

        # row_offset should be in metadata, not token_offset
        for pm in prompt_metadata:
            assert "row_offset" in pm

        shutil.rmtree(tmpdir, ignore_errors=True)


# =============================================================================
# v1.0 backward-compatibility tests
# =============================================================================


@requires_pyarrow
class TestLayersParam:
    """Test selective layer download with layers parameter."""

    @patch("lmprobe.sharing._check_hub_deps")
    def test_load_activation_dataset_layers_filter(
        self, mock_deps, tmp_path,
    ):
        """layers param downloads only requested layers."""
        import pyarrow as pa
        import pyarrow.parquet as pq
        from safetensors.torch import save_file

        from lmprobe.sharing import _embed_metadata_in_schema

        (tmp_path / "tensors").mkdir(parents=True)
        (tmp_path / "index").mkdir(parents=True)

        t0 = torch.randn(3, HIDDEN_DIM)
        t1 = torch.randn(3, HIDDEN_DIM)
        save_file(
            {"hidden.layer_0": t0},
            str(tmp_path / "tensors" / "hidden_layer000_shard000.safetensors"),
        )
        save_file(
            {"hidden.layer_1": t1},
            str(tmp_path / "tensors" / "hidden_layer001_shard000.safetensors"),
        )

        lmprobe_info = {
            "format_version": "2.0",
            "model": {"name": TEST_MODEL, "revision": None},
            "num_prompts": 3,
            "tensors": {
                "hidden_layers": {
                    "type": "hidden",
                    "layers": [0, 1],
                    "dim": HIDDEN_DIM,
                    "dtype": "float32",
                    "layout": "per_layer",
                    "pooling": "last_token",
                    "row_bytes": HIDDEN_DIM * 4,
                    "shards": [
                        {"num_prompts": 3},
                    ],
                }
            },
            "provenance": {},
        }
        table = pa.table({
            "text": pa.array(["a", "b", "c"], type=pa.string()),
            "label": pa.array([None, None, None], type=pa.int32()),
            "num_tokens": pa.array([5, 5, 5], type=pa.int32()),
            "shard_index": pa.array([0, 0, 0], type=pa.int32()),
            "row_offset": pa.array([0, 1, 2], type=pa.int32()),
        })
        table = _embed_metadata_in_schema(table, lmprobe_info)
        pq.write_table(table, str(tmp_path / PARQUET_PATH))

        downloaded_files = []

        def mock_download(repo_id, filename, **kwargs):
            downloaded_files.append(filename)
            return str(tmp_path / filename)

        with patch(
            "huggingface_hub.hf_hub_download", side_effect=mock_download,
        ):
            result, info = load_activation_dataset(
                "user/test", layers=[0],
            )

        assert "hidden.layer_0" in result
        assert "hidden.layer_1" not in result

        # Only layer 0 shard should be downloaded (plus info file)
        tensor_downloads = [
            f for f in downloaded_files if f.startswith("tensors/")
        ]
        assert len(tensor_downloads) == 1
        assert "hidden_layer000" in tensor_downloads[0]

    @patch("lmprobe.sharing._check_hub_deps")
    def test_pull_dataset_layers_filter(self, mock_deps, tmp_path, cache_dir):
        """pull_dataset with layers only downloads requested layers."""
        import pyarrow as pa
        import pyarrow.parquet as pq
        from safetensors.torch import save_file

        remote_dir = tmp_path / "remote"
        (remote_dir / "tensors").mkdir(parents=True)
        (remote_dir / "index").mkdir(parents=True)

        prompts = ["prompt 0", "prompt 1"]
        t0 = torch.randn(2, HIDDEN_DIM)
        t1 = torch.randn(2, HIDDEN_DIM)
        save_file(
            {"hidden.layer_0": t0},
            str(
                remote_dir / "tensors"
                / "hidden_layer000_shard000.safetensors"
            ),
        )
        save_file(
            {"hidden.layer_1": t1},
            str(
                remote_dir / "tensors"
                / "hidden_layer001_shard000.safetensors"
            ),
        )

        from lmprobe.sharing import _embed_metadata_in_schema

        lmprobe_info = {
            "format_version": "2.0",
            "model": {"name": TEST_MODEL, "revision": None},
            "num_prompts": 2,
            "tensors": {
                "hidden_layers": {
                    "type": "hidden",
                    "layers": [0, 1],
                    "dim": HIDDEN_DIM,
                    "dtype": "float32",
                    "layout": "per_layer",
                    "pooling": "last_token",
                    "row_bytes": HIDDEN_DIM * 4,
                    "shards": [
                        {"num_prompts": 2},
                    ],
                }
            },
            "provenance": {},
        }

        table = pa.table({
            "text": pa.array(prompts, type=pa.string()),
            "label": pa.array([None, None], type=pa.int32()),
            "num_tokens": pa.array([5, 5], type=pa.int32()),
            "shard_index": pa.array([0, 0], type=pa.int32()),
            "row_offset": pa.array([0, 1], type=pa.int32()),
        })
        table = _embed_metadata_in_schema(table, lmprobe_info)
        pq.write_table(table, str(remote_dir / PARQUET_PATH))

        downloaded_files = []

        def mock_download(repo_id, filename, **kwargs):
            downloaded_files.append(filename)
            return str(remote_dir / filename)

        with patch(
            "huggingface_hub.hf_hub_download", side_effect=mock_download,
        ):
            count = pull_dataset("user/test", layers=[0])

        assert count == 2

        # Only layer 0 shard was downloaded (plus info + parquet)
        tensor_downloads = [
            f for f in downloaded_files if f.startswith("tensors/")
        ]
        assert len(tensor_downloads) == 1
        assert "hidden_layer000" in tensor_downloads[0]


# =============================================================================
# v1.2 independent shard boundary tests
# =============================================================================


class TestIndependentShardBoundaries:
    """Test that hidden and logits get independent shard boundaries (v1.2)."""

    @requires_pyarrow
    def test_different_shard_counts_per_type(self, tmp_path, monkeypatch):
        """With small logits and large hidden, logits should use fewer shards."""
        src_cache = tmp_path / "src_cache"
        monkeypatch.setenv("LMPROBE_CACHE_DIR", str(src_cache))
        cache_mod._backend = None

        prompts = [f"prompt {i}" for i in range(10)]
        layers = [0, 1]
        top_k = 5
        vocab_size = 50

        for prompt in prompts:
            # Pooled hidden (fairly large per row)
            for layer in layers:
                pooled = torch.randn(1, HIDDEN_DIM)
                save_prompt_pooled_activations(
                    TEST_MODEL, prompt, [layer], pooled, "last_token",
                )
            # Logits (small per row)
            logits = torch.randn(1, 1, vocab_size)
            mask = torch.ones(1, 1, dtype=torch.long)
            save_prompt_logits(
                TEST_MODEL, prompt, logits, mask,
                top_k=top_k, positions="last",
            )

        # Use a tiny shard size so hidden gets multiple shards
        # hidden_row_bytes = 32 * 4 = 128 bytes per row
        # logits_row_bytes = 5 * 4 + 5 * 8 = 60 bytes per row
        # With shard_max_bytes=384, hidden fits 3 rows/shard (4 shards),
        # logits fits 6 rows/shard (2 shards)
        tmpdir, tensor_descriptors, prompt_metadata = _consolidate_and_shard(
            model_name=TEST_MODEL,
            prompts=prompts,
            kept_indices=list(range(10)),
            tensor_types={
                "raw_layers": [],
                "pooled": {"last_token": [0, 1]},
                "has_logits": True,
                "logits_top_k": top_k,
                "has_perplexity": False,
            },
            labels=None,
            shard_max_bytes=384,
            repo_id="user/independent-test",
        )

        hidden_desc = tensor_descriptors["hidden_layers"]
        logits_desc = tensor_descriptors["logits_topk"]

        hidden_shard_count = len(hidden_desc["shards"])
        logits_shard_count = len(logits_desc["shards"])

        # Logits should have fewer shards than hidden
        assert logits_shard_count < hidden_shard_count, (
            f"Expected logits ({logits_shard_count} shards) < "
            f"hidden ({hidden_shard_count} shards)"
        )

        # Total prompts across all shards should equal 10 for each type
        assert sum(s["num_prompts"] for s in hidden_desc["shards"]) == 10
        assert sum(s["num_prompts"] for s in logits_desc["shards"]) == 10

        # Check per-type columns in metadata
        for pm in prompt_metadata:
            # Hidden uses universal shard_index/row_offset
            assert "shard_index" in pm
            assert "row_offset" in pm
            # Logits keeps per-type columns
            assert "shard_index_logits" in pm
            assert "row_offset_logits" in pm

        shutil.rmtree(tmpdir, ignore_errors=True)

    @requires_pyarrow
    def test_per_type_columns_in_parquet(self, tmp_path, monkeypatch):
        """Parquet index includes per-type shard columns."""
        import pyarrow.parquet as pq

        src_cache = tmp_path / "src_cache"
        monkeypatch.setenv("LMPROBE_CACHE_DIR", str(src_cache))
        cache_mod._backend = None

        prompts = ["hello", "world"]
        for prompt in prompts:
            pooled = torch.randn(1, HIDDEN_DIM)
            save_prompt_pooled_activations(
                TEST_MODEL, prompt, [0], pooled, "last_token",
            )
            logits = torch.randn(1, 1, 50)
            mask = torch.ones(1, 1, dtype=torch.long)
            save_prompt_logits(
                TEST_MODEL, prompt, logits, mask,
                top_k=5, positions="last",
            )

        remote_dir = tmp_path / "remote"
        remote_dir.mkdir()
        mock_api = MagicMock()

        def capture_upload(repo_id, folder_path, **kwargs):
            for f in Path(folder_path).rglob("*"):
                if f.is_file():
                    rel = f.relative_to(folder_path)
                    dest = remote_dir / rel
                    dest.parent.mkdir(parents=True, exist_ok=True)
                    shutil.copy2(str(f), str(dest))

        mock_api.upload_large_folder.side_effect = capture_upload

        with (
            patch("lmprobe.sharing._check_hub_deps"),
            patch("huggingface_hub.HfApi", return_value=mock_api),
        ):
            push_dataset(
                repo_id="user/pertype-test",
                model_name=TEST_MODEL,
                prompts=prompts,
                labels=[1, 0],
                exist_ok=True,
            )

        table = pq.read_table(str(remote_dir / PARQUET_PATH))
        col_names = set(table.column_names)
        # Hidden uses universal shard_index/row_offset (no per-type hidden columns)
        assert "shard_index_hidden" not in col_names
        assert "row_offset_hidden" not in col_names
        # Logits keeps per-type columns
        assert "shard_index_logits" in col_names
        assert "row_offset_logits" in col_names
        # Universal columns present
        assert "shard_index" in col_names
        assert "row_offset" in col_names

    def test_format_version_is_2_0(self):
        """Format version is bumped to 2.0."""
        assert FORMAT_VERSION == "2.0"

    @requires_pyarrow
    def test_hidden_only_no_logits_columns(self, tmp_path, monkeypatch):
        """When only hidden is present, no logits columns appear."""
        src_cache = tmp_path / "src_cache"
        monkeypatch.setenv("LMPROBE_CACHE_DIR", str(src_cache))
        cache_mod._backend = None

        prompts = ["hello", "world"]
        for prompt in prompts:
            pooled = torch.randn(1, HIDDEN_DIM)
            save_prompt_pooled_activations(
                TEST_MODEL, prompt, [0], pooled, "last_token",
            )

        tmpdir, tensor_descriptors, prompt_metadata = _consolidate_and_shard(
            model_name=TEST_MODEL,
            prompts=prompts,
            kept_indices=[0, 1],
            tensor_types={
                "raw_layers": [],
                "pooled": {"last_token": [0]},
                "has_logits": False,
                "logits_top_k": None,
                "has_perplexity": False,
            },
            labels=None,
            shard_max_bytes=1_000_000_000,
            repo_id="user/hidden-only",
        )

        assert "hidden_layers" in tensor_descriptors
        assert "logits_topk" not in tensor_descriptors

        # No logits columns in metadata
        for pm in prompt_metadata:
            assert "shard_index_logits" not in pm
            assert "row_offset_logits" not in pm
            # Universal columns present (used for hidden)
            assert "shard_index" in pm
            assert "row_offset" in pm

        shutil.rmtree(tmpdir, ignore_errors=True)


class TestRawShardBoundaryFix107:
    """Test fix for #107: raw per_prompt_bytes should not multiply by layers."""

    def test_raw_shard_count_reasonable(self, tmp_path, monkeypatch):
        """Raw mode shard count should be based on per-layer size, not all layers."""
        src_cache = tmp_path / "src_cache"
        monkeypatch.setenv("LMPROBE_CACHE_DIR", str(src_cache))
        cache_mod._backend = None

        prompts = [f"prompt {i}" for i in range(5)]
        seq_lens = [10, 10, 10, 10, 10]
        layers = [0, 1, 2, 3]  # 4 layers

        for prompt, sl in zip(prompts, seq_lens):
            act = torch.randn(1, sl, HIDDEN_DIM * len(layers))
            mask = torch.ones(1, sl, dtype=torch.long)
            save_prompt_activations(TEST_MODEL, prompt, layers, act, mask)

        # Per-prompt bytes (single layer): 10 tokens * 32 dim * 4 bytes = 1280
        # With shard_max_bytes=4000, should fit 3 rows per shard -> 2 shards
        # Bug #107 would compute 10*32*4*4=5120 per prompt -> only 1 per shard -> 5 shards
        tmpdir, tensor_descriptors, prompt_metadata = _consolidate_and_shard(
            model_name=TEST_MODEL,
            prompts=prompts,
            kept_indices=list(range(5)),
            tensor_types={
                "raw_layers": [0, 1, 2, 3],
                "pooled": {},
                "has_logits": False,
                "logits_top_k": None,
                "has_perplexity": False,
            },
            labels=None,
            shard_max_bytes=4000,
            repo_id="user/raw-fix107",
        )

        hidden_desc = tensor_descriptors["hidden_layers"]
        shard_count = len(hidden_desc["shards"])
        lt_count = hidden_desc.get("last_token_shards", 0)
        rest_count = shard_count - lt_count

        # Last-token shards: 5 prompts * 32 dim * 4 bytes = 640 -> 1 shard
        # Rest-token shards (9 tokens/prompt): 9*32*4=1152 bytes/prompt,
        # 4000 byte shards -> 3 prompts/shard -> 2 shards
        # Total = 1 + 2 = 3 shards
        # Without fix #107: rest would be 9*32*4*4=4608 -> 1/shard -> 5 rest shards
        assert rest_count == 2, (
            f"Expected 2 rest shards (per-layer sizing), got {rest_count}"
        )
        assert lt_count == 1, f"Expected 1 last-token shard, got {lt_count}"

        shutil.rmtree(tmpdir, ignore_errors=True)


class TestBackwardCompatLegacySingleShardIndex:
    """Test that v1.1 datasets with single shard_index still work."""

    @requires_pyarrow
    @patch("lmprobe.sharing._check_hub_deps")
    def test_pull_v11_single_shard_index(self, mock_deps, tmp_path, cache_dir):
        """pull_dataset works with v1.1 datasets that only have shard_index."""
        import pyarrow as pa
        import pyarrow.parquet as pq
        from safetensors.torch import save_file

        remote_dir = tmp_path / "remote"
        (remote_dir / "tensors").mkdir(parents=True)
        (remote_dir / "index").mkdir(parents=True)

        prompts = ["prompt 0", "prompt 1"]
        t0 = torch.randn(2, HIDDEN_DIM)
        save_file(
            {"hidden.layer_0": t0},
            str(remote_dir / "tensors" / "hidden_layer000_shard000.safetensors"),
        )

        from lmprobe.sharing import _embed_metadata_in_schema

        lmprobe_info = {
            "format_version": "2.0",
            "model": {"name": TEST_MODEL, "revision": None},
            "num_prompts": 2,
            "tensors": {
                "hidden_layers": {
                    "type": "hidden",
                    "layers": [0],
                    "dim": HIDDEN_DIM,
                    "dtype": "float32",
                    "layout": "per_layer",
                    "pooling": "last_token",
                    "row_bytes": HIDDEN_DIM * 4,
                    "shards": [{"num_prompts": 2}],
                }
            },
            "provenance": {},
        }

        # v2 Parquet: only legacy columns (no per-type columns)
        table = pa.table({
            "text": pa.array(prompts, type=pa.string()),
            "label": pa.array([None, None], type=pa.int32()),
            "num_tokens": pa.array([5, 5], type=pa.int32()),
            "shard_index": pa.array([0, 0], type=pa.int32()),
            "row_offset": pa.array([0, 1], type=pa.int32()),
        })
        table = _embed_metadata_in_schema(table, lmprobe_info)
        pq.write_table(table, str(remote_dir / PARQUET_PATH))

        def mock_download(repo_id, filename, **kwargs):
            return str(remote_dir / filename)

        with patch(
            "huggingface_hub.hf_hub_download", side_effect=mock_download,
        ):
            count = pull_dataset("user/v11-test")

        assert count == 2


# =============================================================================
# Deterministic staging directory tests
# =============================================================================


class TestStagingDir:
    """Tests for deterministic staging directory and resume logic."""

    def test_staging_dir_deterministic(self, cache_dir):
        """Same arguments produce the same staging path."""
        prompts = ["hello", "world"]
        path1 = _staging_dir_path("user/repo", "model-a", prompts)
        path2 = _staging_dir_path("user/repo", "model-a", prompts)
        assert path1 == path2

    def test_staging_dir_different_for_different_args(self, cache_dir):
        """Different arguments produce different staging paths."""
        prompts = ["hello", "world"]
        path_a = _staging_dir_path("user/repo", "model-a", prompts)
        path_b = _staging_dir_path("user/repo", "model-b", prompts)
        path_c = _staging_dir_path("user/other-repo", "model-a", prompts)
        path_d = _staging_dir_path("user/repo", "model-a", ["different"])
        assert len({path_a, path_b, path_c, path_d}) == 4

    def test_staging_dir_order_dependent(self, cache_dir):
        """Different prompt order produces different staging path.

        Labels and metadata are positionally aligned with prompts, so
        reordering prompts with the same labels would produce wrong data.
        """
        path1 = _staging_dir_path("user/repo", "m", ["b", "a"])
        path2 = _staging_dir_path("user/repo", "m", ["a", "b"])
        assert path1 != path2

    def test_staging_dir_includes_shard_bytes(self, cache_dir):
        """Different shard_max_bytes produces different staging path."""
        path1 = _staging_dir_path("user/repo", "m", ["x"], shard_max_bytes=1024)
        path2 = _staging_dir_path("user/repo", "m", ["x"], shard_max_bytes=2048)
        assert path1 != path2

    def test_staging_dir_includes_labels(self, cache_dir):
        """Different labels produce different staging path."""
        path1 = _staging_dir_path("user/repo", "m", ["x"], labels=[0])
        path2 = _staging_dir_path("user/repo", "m", ["x"], labels=[1])
        path3 = _staging_dir_path("user/repo", "m", ["x"], labels=None)
        assert len({path1, path2, path3}) == 3

    def test_staging_dir_includes_metadata(self, cache_dir):
        """Different metadata produce different staging path."""
        path1 = _staging_dir_path("user/repo", "m", ["x"], metadata=[{"k": "a"}])
        path2 = _staging_dir_path("user/repo", "m", ["x"], metadata=[{"k": "b"}])
        path3 = _staging_dir_path("user/repo", "m", ["x"], metadata=None)
        assert len({path1, path2, path3}) == 3

    def test_staging_dir_includes_tensors_filter(self, cache_dir):
        """Different tensors filter produces different staging path."""
        path1 = _staging_dir_path("user/repo", "m", ["x"], tensors=["hidden_layers"])
        path2 = _staging_dir_path("user/repo", "m", ["x"], tensors=["logits_topk"])
        path3 = _staging_dir_path("user/repo", "m", ["x"], tensors=None)
        assert len({path1, path2, path3}) == 3

    def test_staging_dir_under_cache(self, cache_dir):
        """Staging dir lives under the cache directory."""
        path = _staging_dir_path("user/repo", "m", ["x"])
        assert str(path).startswith(str(cache_dir))
        assert "staging" in path.parts

    @requires_pyarrow
    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    @patch("huggingface_hub.HfApi")
    def test_push_dataset_resumes_from_sentinel(
        self, MockHfApi, mock_pyarrow, mock_deps, populated_cache, cache_dir
    ):
        """When sentinel exists, consolidation is skipped entirely."""
        prompts = populated_cache
        repo_id = "user/resume-test"
        model_name = TEST_MODEL

        mock_api = MagicMock()
        MockHfApi.return_value = mock_api

        staging_dir = _staging_dir_path(repo_id, model_name, prompts)
        staging_dir.mkdir(parents=True, exist_ok=True)
        (staging_dir / "tensors").mkdir(exist_ok=True)
        (staging_dir / "index").mkdir(exist_ok=True)

        # Write a minimal lmprobe_info.json
        info = _build_lmprobe_info(model_name, len(prompts), {})
        with open(staging_dir / INFO_FILENAME, "w") as f:
            json.dump(info, f)

        # Write a dummy parquet so upload_large_folder has something
        import pyarrow as pa
        import pyarrow.parquet as pq

        table = pa.table({
            "text": pa.array(prompts, type=pa.string()),
            "label": pa.array([None] * len(prompts), type=pa.int32()),
            "num_tokens": pa.array([5] * len(prompts), type=pa.int32()),
        })
        pq.write_table(table, str(staging_dir / PARQUET_PATH))

        # Write README
        with open(staging_dir / "README.md", "w") as f:
            f.write("test")

        # Touch the sentinel
        (staging_dir / _STAGE_SENTINEL).touch()

        with patch(
            "lmprobe.sharing._consolidate_and_shard"
        ) as mock_consolidate:
            # Note: exist_ok=False (the default) — resume should force True
            push_dataset(
                repo_id, model_name, prompts,
            )
            # Consolidation should NOT have been called
            mock_consolidate.assert_not_called()
            # Upload should still have been called
            mock_api.upload_large_folder.assert_called_once()
            # create_repo should have been called with exist_ok=True
            # even though we didn't pass exist_ok=True to push_dataset
            mock_api.create_repo.assert_called_once_with(
                repo_id,
                exist_ok=True,
                private=False,
                repo_type="dataset",
            )

    @requires_pyarrow
    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    @patch("huggingface_hub.HfApi")
    def test_push_dataset_creates_sentinel_after_staging(
        self, MockHfApi, mock_pyarrow, mock_deps, populated_cache, cache_dir
    ):
        """Sentinel is created after consolidation completes."""
        prompts = populated_cache
        repo_id = "user/sentinel-test"
        model_name = TEST_MODEL

        mock_api = MagicMock()
        MockHfApi.return_value = mock_api

        staging_dir = _staging_dir_path(repo_id, model_name, prompts)
        sentinel = staging_dir / _STAGE_SENTINEL

        # Sentinel should not exist before push
        assert not sentinel.exists()

        # Make upload fail so we can inspect state before cleanup
        mock_api.upload_large_folder.side_effect = Exception("upload fail")

        with pytest.raises(Exception, match="upload fail"):
            push_dataset(
                repo_id, model_name, prompts, exist_ok=True,
            )

        # Sentinel should exist after consolidation even though upload failed
        assert sentinel.exists()

    @requires_pyarrow
    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    @patch("huggingface_hub.HfApi")
    def test_staging_dir_cleaned_on_success(
        self, MockHfApi, mock_pyarrow, mock_deps, populated_cache, cache_dir
    ):
        """Staging directory is removed after successful upload."""
        prompts = populated_cache
        repo_id = "user/cleanup-test"
        model_name = TEST_MODEL

        mock_api = MagicMock()
        MockHfApi.return_value = mock_api

        staging_dir = _staging_dir_path(repo_id, model_name, prompts)

        push_dataset(
            repo_id, model_name, prompts, exist_ok=True,
        )

        # Staging dir should be cleaned up
        assert not staging_dir.exists()


class TestLastTokenShardSplitting:
    """Tests for last-token shard splitting (issue #165)."""

    def test_last_token_shards_created(self, populated_raw_cache):
        """Consolidation with raw layers creates separate lt and rest shards."""
        from safetensors import safe_open

        prompts, seq_lens, layers, _ = populated_raw_cache
        tensor_types = {
            "raw_layers": layers,
            "pooled": {},
            "has_logits": False,
            "logits_top_k": None,
            "has_perplexity": False,
        }
        tmpdir, tensor_descriptors, prompt_metadata = _consolidate_and_shard(
            model_name=TEST_MODEL,
            prompts=prompts,
            kept_indices=[0, 1, 2],
            tensor_types=tensor_types,
            labels=None,
            shard_max_bytes=1_000_000_000,
            repo_id="user/lt-test",
        )

        desc = tensor_descriptors["hidden_layers"]
        assert desc["storage"] == "full_sequence"
        assert "last_token_shards" in desc
        lt_count = desc["last_token_shards"]
        assert lt_count >= 1

        # Last-token shard should have exactly num_prompts tokens
        # (one per prompt)
        lt_shards = desc["shards"][:lt_count]
        lt_total_tokens = sum(s["num_tokens"] for s in lt_shards)
        assert lt_total_tokens == len(prompts)

        # Rest shards should have sum(seq_lens) - num_prompts tokens
        rest_shards = desc["shards"][lt_count:]
        rest_total_tokens = sum(s["num_tokens"] for s in rest_shards)
        assert rest_total_tokens == sum(seq_lens) - len(prompts)

        # Verify shard files exist and have correct shapes
        for layer in layers:
            # Last-token shard
            lt_file = tmpdir / f"tensors/hidden_layer{layer:03d}_shard000.safetensors"
            assert lt_file.exists()
            with safe_open(str(lt_file), framework="pt") as f:
                t = f.get_tensor(f"hidden.layer_{layer}")
                assert t.shape == (len(prompts), HIDDEN_DIM)

            # Rest shard (should exist if any prompt has > 1 token)
            rest_file = tmpdir / f"tensors/hidden_layer{layer:03d}_shard001.safetensors"
            assert rest_file.exists()
            with safe_open(str(rest_file), framework="pt") as f:
                t = f.get_tensor(f"hidden.layer_{layer}")
                expected_rest_tokens = sum(sl - 1 for sl in seq_lens)
                assert t.shape == (expected_rest_tokens, HIDDEN_DIM)

        shutil.rmtree(tmpdir, ignore_errors=True)

    def test_per_token_arrays_in_metadata(self, populated_raw_cache):
        """prompt_metadata has token_shard_ids and token_shard_offsets."""
        prompts, seq_lens, layers, _ = populated_raw_cache
        tensor_types = {
            "raw_layers": layers,
            "pooled": {},
            "has_logits": False,
            "logits_top_k": None,
            "has_perplexity": False,
        }
        _, tensor_descriptors, prompt_metadata = _consolidate_and_shard(
            model_name=TEST_MODEL,
            prompts=prompts,
            kept_indices=[0, 1, 2],
            tensor_types=tensor_types,
            labels=None,
            shard_max_bytes=1_000_000_000,
            repo_id="user/lt-test",
        )

        lt_count = tensor_descriptors["hidden_layers"]["last_token_shards"]

        for pm in prompt_metadata:
            assert "token_shard_ids" in pm
            assert "token_shard_offsets" in pm
            num_tok = pm["num_tokens"]
            shard_ids = pm["token_shard_ids"]
            shard_offsets = pm["token_shard_offsets"]

            # Arrays have correct length
            assert len(shard_ids) == num_tok
            assert len(shard_offsets) == num_tok

            # Last token is in a last-token shard (index < lt_count)
            assert shard_ids[-1] < lt_count

            # Non-last tokens are in rest shards (index >= lt_count)
            if num_tok > 1:
                for sid in shard_ids[:-1]:
                    assert sid >= lt_count

            # Universal shard_index points to last-token shard
            assert pm["shard_index"] == shard_ids[-1]

    def test_per_token_offsets_consistency(self, populated_raw_cache):
        """token_shard_offsets are valid indices into shard tensors."""
        from safetensors import safe_open

        prompts, seq_lens, layers, _ = populated_raw_cache
        tensor_types = {
            "raw_layers": layers,
            "pooled": {},
            "has_logits": False,
            "logits_top_k": None,
            "has_perplexity": False,
        }
        tmpdir, tensor_descriptors, prompt_metadata = _consolidate_and_shard(
            model_name=TEST_MODEL,
            prompts=prompts,
            kept_indices=[0, 1, 2],
            tensor_types=tensor_types,
            labels=None,
            shard_max_bytes=1_000_000_000,
            repo_id="user/lt-test",
        )

        # Load all shard sizes for layer 0
        shard_sizes = {}
        for shard_desc_idx, shard_desc in enumerate(
            tensor_descriptors["hidden_layers"]["shards"]
        ):
            fname = f"tensors/hidden_layer000_shard{shard_desc_idx:03d}.safetensors"
            fpath = tmpdir / fname
            if fpath.exists():
                with safe_open(str(fpath), framework="pt") as f:
                    t = f.get_tensor("hidden.layer_0")
                    shard_sizes[shard_desc_idx] = t.shape[0]

        # Verify all offsets are within shard bounds
        for pm in prompt_metadata:
            for sid, soff in zip(
                pm["token_shard_ids"], pm["token_shard_offsets"]
            ):
                assert sid in shard_sizes, f"shard {sid} not found"
                assert 0 <= soff < shard_sizes[sid], (
                    f"offset {soff} out of bounds for shard {sid} "
                    f"(size {shard_sizes[sid]})"
                )

        shutil.rmtree(tmpdir, ignore_errors=True)

    @requires_pyarrow
    def test_parquet_has_list_columns(self, populated_raw_cache):
        """Parquet index contains token_shard_ids/offsets list columns."""
        import pyarrow.parquet as pq

        prompts, seq_lens, layers, _ = populated_raw_cache
        tensor_types = {
            "raw_layers": layers,
            "pooled": {},
            "has_logits": False,
            "logits_top_k": None,
            "has_perplexity": False,
        }
        tmpdir, _, prompt_metadata = _consolidate_and_shard(
            model_name=TEST_MODEL,
            prompts=prompts,
            kept_indices=[0, 1, 2],
            tensor_types=tensor_types,
            labels=None,
            shard_max_bytes=1_000_000_000,
            repo_id="user/lt-test",
        )

        _write_parquet_index(tmpdir, prompt_metadata)
        table = pq.read_table(str(tmpdir / PARQUET_PATH))
        columns = table.column_names

        assert "token_shard_ids" in columns
        assert "token_shard_offsets" in columns

        # Verify round-trip: values match metadata
        idx = table.to_pydict()
        for i, pm in enumerate(prompt_metadata):
            assert idx["token_shard_ids"][i] == pm["token_shard_ids"]
            assert idx["token_shard_offsets"][i] == pm["token_shard_offsets"]

        shutil.rmtree(tmpdir, ignore_errors=True)

    def test_enumerate_shard_files_includes_last_token(
        self, populated_raw_cache
    ):
        """_enumerate_shard_files lists both lt and rest shard files."""
        prompts, seq_lens, layers, _ = populated_raw_cache
        tensor_types = {
            "raw_layers": layers,
            "pooled": {},
            "has_logits": False,
            "logits_top_k": None,
            "has_perplexity": False,
        }
        plan = _compute_shard_plan(
            model_name=TEST_MODEL,
            prompts=prompts,
            kept_indices=[0, 1, 2],
            tensor_types=tensor_types,
            labels=None,
            shard_max_bytes=1_000_000_000,
            repo_id="user/lt-test",
        )

        files = _enumerate_shard_files(plan)
        # Should have at least 2 shards per layer (lt + rest)
        assert len(files) >= len(layers) * 2

        # Verify all listed files get created during consolidation
        tmpdir, _, _ = _consolidate_and_shard(
            model_name=TEST_MODEL,
            prompts=prompts,
            kept_indices=[0, 1, 2],
            tensor_types=tensor_types,
            labels=None,
            shard_max_bytes=1_000_000_000,
            repo_id="user/lt-test",
        )
        for f in files:
            assert (tmpdir / f).exists(), f"Expected shard file {f} not found"

        shutil.rmtree(tmpdir, ignore_errors=True)


# =============================================================================
# Migration tests
# =============================================================================


@requires_pyarrow
class TestMigrateDataset:
    """Tests for migrate_dataset: converting old full_sequence datasets
    to use last-token shard splitting."""

    @staticmethod
    def _build_old_format_dataset(dest: Path) -> dict:
        """Create an old-format full_sequence dataset without lt splitting."""
        import pyarrow as pa
        import pyarrow.parquet as pq
        from safetensors.torch import save_file

        prompts = ["alpha", "beta", "gamma"]
        seq_lens = [3, 5, 4]
        layers_list = [0, 1]
        n_prompts = len(prompts)

        (dest / "tensors").mkdir(parents=True, exist_ok=True)
        (dest / "index").mkdir(parents=True, exist_ok=True)

        gen = torch.Generator().manual_seed(42)
        original_tensors: dict[str, dict[int, torch.Tensor]] = {}
        for prompt, sl in zip(prompts, seq_lens):
            original_tensors[prompt] = {}
            for layer in layers_list:
                original_tensors[prompt][layer] = torch.randn(
                    sl, HIDDEN_DIM, generator=gen,
                )

        for layer in layers_list:
            rows = [original_tensors[p][layer] for p in prompts]
            tensor = torch.cat(rows, dim=0)
            key = f"hidden.layer_{layer}"
            fname = (
                f"tensors/hidden_layer{layer:03d}"
                f"_shard000.safetensors"
            )
            save_file({key: tensor}, str(dest / fname))

        tok_offset = 0
        token_offsets = []
        for sl in seq_lens:
            token_offsets.append(tok_offset)
            tok_offset += sl

        from lmprobe.sharing import _embed_metadata_in_schema

        total_tokens = sum(seq_lens)
        lmprobe_info = {
            "format_version": FORMAT_VERSION,
            "model": {"name": TEST_MODEL, "revision": None},
            "num_prompts": n_prompts,
            "prompt_ordering": "random",
            "tensors": {
                "hidden_layers": {
                    "type": "hidden",
                    "layers": layers_list,
                    "dim": HIDDEN_DIM,
                    "dtype": "float32",
                    "layout": "per_layer",
                    "storage": "full_sequence",
                    "shards": [{
                        "num_prompts": n_prompts,
                        "num_tokens": total_tokens,
                    }],
                },
            },
            "provenance": {
                "lmprobe_version": "0.8.0",
                "created_at": "2024-01-01T00:00:00+00:00",
            },
        }

        columns = {
            "text": pa.array(prompts, type=pa.string()),
            "label": pa.array([0, 1, 0], type=pa.int32()),
            "num_tokens": pa.array(seq_lens, type=pa.int32()),
            "shard_index": pa.array(
                [0] * n_prompts, type=pa.int32(),
            ),
            "row_offset": pa.array(
                list(range(n_prompts)), type=pa.int32(),
            ),
            "shard_index_hidden": pa.array(
                [0] * n_prompts, type=pa.int32(),
            ),
            "row_offset_hidden": pa.array(
                list(range(n_prompts)), type=pa.int32(),
            ),
            "token_offset_hidden": pa.array(
                token_offsets, type=pa.int64(),
            ),
            "token_offset": pa.array(
                token_offsets, type=pa.int64(),
            ),
        }
        table = pa.table(columns)
        table = _embed_metadata_in_schema(table, lmprobe_info)
        pq.write_table(table, str(dest / PARQUET_PATH))

        return {
            "prompts": prompts,
            "seq_lens": seq_lens,
            "layers": layers_list,
            "original_tensors": original_tensors,
        }

    def test_migrate_rejects_non_full_sequence(self, tmp_path):
        """migrate_dataset raises on pooled datasets."""
        import pyarrow as pa
        import pyarrow.parquet as pq

        from lmprobe.sharing import _embed_metadata_in_schema, migrate_dataset

        old_dir = tmp_path / "old"
        old_dir.mkdir()
        (old_dir / "index").mkdir(parents=True)

        info = {
            "format_version": FORMAT_VERSION,
            "model": {"name": TEST_MODEL},
            "num_prompts": 1,
            "tensors": {
                "hidden_layers": {
                    "storage": "pooled",
                    "layers": [0],
                    "dim": HIDDEN_DIM,
                    "shards": [],
                },
            },
        }

        table = pa.table({
            "text": ["a"], "label": [0], "num_tokens": [1],
            "shard_index": [0], "row_offset": [0],
        })
        table = _embed_metadata_in_schema(table, info)
        pq.write_table(table, str(old_dir / PARQUET_PATH))

        def mock_dl(repo_id, filename, **kwargs):
            return str(old_dir / filename)

        with (
            patch("lmprobe.sharing._check_hub_deps"),
            patch("lmprobe.sharing._check_pyarrow"),
            patch(
                "huggingface_hub.hf_hub_download",
                side_effect=mock_dl,
            ),
        ):
            with pytest.raises(ValueError, match="not full_sequence"):
                migrate_dataset("user/pooled-ds")

    def test_migrate_rejects_already_migrated(self, tmp_path):
        """migrate_dataset raises on already-migrated datasets."""
        import pyarrow as pa
        import pyarrow.parquet as pq

        from lmprobe.sharing import _embed_metadata_in_schema, migrate_dataset

        old_dir = tmp_path / "old"
        old_dir.mkdir()
        (old_dir / "index").mkdir(parents=True)

        info = {
            "format_version": FORMAT_VERSION,
            "model": {"name": TEST_MODEL},
            "num_prompts": 1,
            "tensors": {
                "hidden_layers": {
                    "storage": "full_sequence",
                    "last_token_shards": 1,
                    "layers": [0],
                    "dim": HIDDEN_DIM,
                    "shards": [],
                },
            },
        }

        table = pa.table({
            "text": ["a"], "label": [0], "num_tokens": [1],
            "shard_index": [0], "row_offset": [0],
        })
        table = _embed_metadata_in_schema(table, info)
        pq.write_table(table, str(old_dir / PARQUET_PATH))

        def mock_dl(repo_id, filename, **kwargs):
            return str(old_dir / filename)

        with (
            patch("lmprobe.sharing._check_hub_deps"),
            patch("lmprobe.sharing._check_pyarrow"),
            patch(
                "huggingface_hub.hf_hub_download",
                side_effect=mock_dl,
            ),
        ):
            with pytest.raises(ValueError, match="already has"):
                migrate_dataset("user/already-migrated")

    def test_migrate_dry_run(self, tmp_path):
        """dry_run computes plan without uploading."""
        from lmprobe.sharing import migrate_dataset

        old_dir = tmp_path / "old"
        old_dir.mkdir()
        self._build_old_format_dataset(old_dir)

        def mock_dl(repo_id, filename, **kwargs):
            return str(old_dir / filename)

        with (
            patch("lmprobe.sharing._check_hub_deps"),
            patch("lmprobe.sharing._check_pyarrow"),
            patch(
                "huggingface_hub.hf_hub_download",
                side_effect=mock_dl,
            ),
        ):
            result = migrate_dataset("user/test-ds", dry_run=True)

        assert "[DRY RUN]" in result
        assert "3 prompts" in result
        assert "2 layers" in result

    def test_migrate_roundtrip(self, tmp_path):
        """Full migration: old format -> migrate -> verify."""
        import pyarrow.parquet as pq
        from safetensors import safe_open

        from lmprobe.sharing import migrate_dataset

        old_dir = tmp_path / "old"
        old_dir.mkdir()
        data = self._build_old_format_dataset(old_dir)
        prompts = data["prompts"]
        seq_lens = data["seq_lens"]
        layers_list = data["layers"]
        original_tensors = data["original_tensors"]

        migrated_dir = tmp_path / "migrated"
        migrated_dir.mkdir()
        mock_api = MagicMock()

        def capture_upload(repo_id, folder_path, **kwargs):
            src = Path(folder_path)
            for f in src.rglob("*"):
                if f.is_file():
                    rel = f.relative_to(src)
                    d = migrated_dir / rel
                    d.parent.mkdir(parents=True, exist_ok=True)
                    shutil.copy2(str(f), str(d))

        mock_api.upload_large_folder.side_effect = capture_upload

        def mock_dl(repo_id, filename, **kwargs):
            return str(old_dir / filename)

        with (
            patch("lmprobe.sharing._check_hub_deps"),
            patch("lmprobe.sharing._check_pyarrow"),
            patch(
                "huggingface_hub.hf_hub_download",
                side_effect=mock_dl,
            ),
            patch(
                "huggingface_hub.HfApi",
                return_value=mock_api,
            ),
        ):
            result = migrate_dataset("user/test-ds")

        assert "huggingface.co" in result

        # Verify lmprobe_info.json
        with open(migrated_dir / INFO_FILENAME) as f:
            info = json.load(f)

        hidden_desc = info["tensors"]["hidden_layers"]
        assert hidden_desc["storage"] == "full_sequence"
        assert "last_token_shards" in hidden_desc
        lt_count = hidden_desc["last_token_shards"]
        assert lt_count >= 1

        lt_shards = hidden_desc["shards"][:lt_count]
        assert sum(
            s["num_tokens"] for s in lt_shards
        ) == len(prompts)

        rest_shards = hidden_desc["shards"][lt_count:]
        assert sum(
            s["num_tokens"] for s in rest_shards
        ) == sum(seq_lens) - len(prompts)

        # Verify parquet index
        idx = pq.read_table(
            str(migrated_dir / PARQUET_PATH),
        ).to_pydict()
        assert "token_shard_ids" in idx
        assert "token_shard_offsets" in idx

        for i in range(len(prompts)):
            sids = idx["token_shard_ids"][i]
            soffs = idx["token_shard_offsets"][i]
            num_tok = seq_lens[i]

            assert len(sids) == num_tok
            assert len(soffs) == num_tok
            assert sids[-1] < lt_count
            if num_tok > 1:
                for sid in sids[:-1]:
                    assert sid >= lt_count

        # Verify shard data integrity
        for layer in layers_list:
            lt_fname = (
                f"tensors/hidden_layer{layer:03d}"
                f"_shard000.safetensors"
            )
            with safe_open(
                str(migrated_dir / lt_fname), framework="pt",
            ) as sf:
                lt_t = sf.get_tensor(f"hidden.layer_{layer}")
            assert lt_t.shape == (len(prompts), HIDDEN_DIM)

            for i, prompt in enumerate(prompts):
                orig_last = original_tensors[prompt][layer][-1]
                assert torch.allclose(lt_t[i], orig_last, atol=1e-6)

            rest_fname = (
                f"tensors/hidden_layer{layer:03d}"
                f"_shard{lt_count:03d}.safetensors"
            )
            with safe_open(
                str(migrated_dir / rest_fname), framework="pt",
            ) as sf:
                rest_t = sf.get_tensor(f"hidden.layer_{layer}")

            expected_rest = sum(sl - 1 for sl in seq_lens)
            assert rest_t.shape == (expected_rest, HIDDEN_DIM)

            offset = 0
            for i, prompt in enumerate(prompts):
                n_rest = seq_lens[i] - 1
                orig_rest = original_tensors[prompt][layer][:-1]
                assert torch.allclose(
                    rest_t[offset:offset + n_rest],
                    orig_rest, atol=1e-6,
                )
                offset += n_rest


@requires_pyarrow
class TestUpgradeDatasetFormat:
    """Tests for upgrade_dataset_format: v1→v2 metadata migration."""

    def _make_v1_dataset(self, tmp_path):
        """Create a realistic v1 dataset with JSON sidecar and parquet index."""
        import pyarrow as pa
        import pyarrow.parquet as pq

        (tmp_path / "index").mkdir(parents=True)

        # Realistic v1 metadata (based on real-world latent-lab dataset)
        v1_info = {
            "format_version": "1.1",
            "model": {"name": "meta-llama/Llama-3.1-405B", "revision": None},
            "num_prompts": 3,
            "prompt_ordering": "random",
            "tensors": {
                "hidden_layers": {
                    "type": "hidden",
                    "layers": [0, 1],
                    "dim": HIDDEN_DIM,
                    "dtype": "float16",
                    "layout": "per_layer",
                    "storage": "full_sequence",
                    "shards": [
                        {"num_prompts": 3, "num_tokens": 3},
                        {"num_prompts": 2, "num_tokens": 8},
                    ],
                    "last_token_shards": 1,
                }
            },
            "provenance": {
                "lmprobe_version": "0.8.0",
                "extraction_backend": "ndif",
                "created_at": "2025-03-19T00:00:00+00:00",
            },
        }

        # Write v1 JSON sidecar
        info_path = tmp_path / INFO_FILENAME
        info_path.write_text(json.dumps(v1_info))

        # Write v1 parquet with per-type hidden columns
        table = pa.table({
            "text": pa.array(["p0", "p1", "p2"], type=pa.string()),
            "label": pa.array([None, None, None], type=pa.int32()),
            "num_tokens": pa.array([5, 7, 4], type=pa.int32()),
            "shard_index": pa.array([0, 0, 0], type=pa.int32()),
            "row_offset": pa.array([0, 1, 2], type=pa.int32()),
            # v1 per-type hidden columns (should be dropped by migration)
            "shard_index_hidden": pa.array([0, 0, 0], type=pa.int32()),
            "row_offset_hidden": pa.array([0, 1, 2], type=pa.int32()),
            "token_offset_hidden": pa.array([0, 5, 12], type=pa.int64()),
        })
        pq.write_table(table, str(tmp_path / PARQUET_PATH))

        return tmp_path, v1_info

    def _run_upgrade(self, tmp_path, dest_path):
        """Run upgrade_dataset_format with mocked HF API, capturing uploaded parquet."""
        uploaded = {}

        def mock_download(repo_id, filename, **kwargs):
            return str(tmp_path / filename)

        def mock_create_commit(repo_id, operations, **kwargs):
            for op in operations:
                p = getattr(op, "path_or_fileobj", None)
                if isinstance(p, str) and Path(p).exists():
                    shutil.copy2(p, dest_path)
                    uploaded[op.path_in_repo] = str(dest_path)

        with patch(
            "huggingface_hub.hf_hub_download", side_effect=mock_download,
        ), patch(
            "huggingface_hub.HfApi", return_value=MagicMock(
                create_commit=mock_create_commit,
            ),
        ), patch(
            "huggingface_hub.CommitOperationAdd",
            side_effect=lambda **kw: MagicMock(**kw),
        ), patch(
            "huggingface_hub.CommitOperationDelete",
            side_effect=lambda **kw: MagicMock(**kw),
        ):
            url = upgrade_dataset_format("user/test-v1")

        return url, uploaded

    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    def test_injects_file_pattern_and_key_pattern(
        self, mock_pa, mock_deps, tmp_path
    ):
        """Migration should add file_pattern/key_pattern to tensor descriptors."""
        remote_dir, _ = self._make_v1_dataset(tmp_path)
        dest = tmp_path / "uploaded.parquet"
        url, uploaded = self._run_upgrade(remote_dir, dest)

        assert "user/test-v1" in url

        import pyarrow.parquet as pq
        schema = pq.read_schema(str(dest))
        tensors_raw = json.loads(schema.metadata[b"lmprobe:tensors"])
        hidden = tensors_raw["hidden_layers"]

        assert hidden["file_pattern"] == (
            "tensors/hidden_layer{layer:03d}_shard{shard:03d}.safetensors"
        )
        assert hidden["key_pattern"] == "hidden.layer_{layer}"

    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    def test_drops_per_type_hidden_columns(
        self, mock_pa, mock_deps, tmp_path
    ):
        """Migration should drop shard_index_hidden etc. from parquet."""
        remote_dir, _ = self._make_v1_dataset(tmp_path)
        dest = tmp_path / "uploaded.parquet"
        self._run_upgrade(remote_dir, dest)

        import pyarrow.parquet as pq
        cols = set(pq.read_table(str(dest)).column_names)

        assert "shard_index" in cols
        assert "row_offset" in cols
        assert "shard_index_hidden" not in cols
        assert "row_offset_hidden" not in cols
        assert "token_offset_hidden" not in cols

    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    def test_stamps_format_version_2(
        self, mock_pa, mock_deps, tmp_path
    ):
        """Migration should set format_version to 2.0."""
        remote_dir, _ = self._make_v1_dataset(tmp_path)
        dest = tmp_path / "uploaded.parquet"
        self._run_upgrade(remote_dir, dest)

        import pyarrow.parquet as pq
        schema = pq.read_schema(str(dest))
        version = json.loads(schema.metadata[b"lmprobe:format_version"])
        assert version == FORMAT_VERSION

    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    def test_renames_tensor_types_key(
        self, mock_pa, mock_deps, tmp_path
    ):
        """Very old datasets with tensor_types key should be renamed to tensors."""
        import pyarrow as pa
        import pyarrow.parquet as pq

        (tmp_path / "index").mkdir(parents=True)

        v1_info = {
            "format_version": "1.0",
            "model": {"name": TEST_MODEL, "revision": None},
            "num_prompts": 2,
            "tensor_types": {
                "hidden_layers": {
                    "type": "hidden",
                    "layers": [0],
                    "dim": HIDDEN_DIM,
                    "dtype": "float32",
                    "layout": "per_layer",
                    "storage": "pooled",
                    "shards": [{"num_prompts": 2}],
                }
            },
            "provenance": {},
        }
        (tmp_path / INFO_FILENAME).write_text(json.dumps(v1_info))
        table = pa.table({
            "text": pa.array(["a", "b"], type=pa.string()),
            "shard_index": pa.array([0, 0], type=pa.int32()),
            "row_offset": pa.array([0, 1], type=pa.int32()),
        })
        pq.write_table(table, str(tmp_path / PARQUET_PATH))

        dest = tmp_path / "uploaded.parquet"
        self._run_upgrade(tmp_path, dest)

        schema = pq.read_schema(str(dest))
        meta = schema.metadata
        assert b"lmprobe:tensors" in meta
        assert b"lmprobe:tensor_types" not in meta
        tensors = json.loads(meta[b"lmprobe:tensors"])
        assert "file_pattern" in tensors["hidden_layers"]

    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    def test_logits_get_file_pattern(
        self, mock_pa, mock_deps, tmp_path
    ):
        """Migration should inject file_pattern into logits_topk descriptor."""
        import pyarrow as pa
        import pyarrow.parquet as pq

        (tmp_path / "index").mkdir(parents=True)

        v1_info = {
            "format_version": "1.1",
            "model": {"name": TEST_MODEL, "revision": None},
            "num_prompts": 2,
            "tensors": {
                "hidden_layers": {
                    "type": "hidden",
                    "layers": [0],
                    "dim": HIDDEN_DIM,
                    "dtype": "float32",
                    "layout": "per_layer",
                    "storage": "pooled",
                    "shards": [{"num_prompts": 2}],
                },
                "logits_topk": {
                    "type": "logits_topk",
                    "k": 50,
                    "dtype": "float32",
                    "pooling": "last_token",
                    "shards": [
                        {"file": "tensors/logits_topk_000.safetensors",
                         "num_prompts": 2},
                    ],
                },
            },
            "provenance": {},
        }
        (tmp_path / INFO_FILENAME).write_text(json.dumps(v1_info))
        table = pa.table({
            "text": pa.array(["a", "b"], type=pa.string()),
            "shard_index": pa.array([0, 0], type=pa.int32()),
            "row_offset": pa.array([0, 1], type=pa.int32()),
            "shard_index_logits": pa.array([0, 0], type=pa.int32()),
            "row_offset_logits": pa.array([0, 1], type=pa.int32()),
        })
        pq.write_table(table, str(tmp_path / PARQUET_PATH))

        dest = tmp_path / "uploaded.parquet"
        self._run_upgrade(tmp_path, dest)

        schema = pq.read_schema(str(dest))
        tensors = json.loads(schema.metadata[b"lmprobe:tensors"])
        assert tensors["logits_topk"]["file_pattern"] == (
            "tensors/logits_topk_{shard:03d}.safetensors"
        )
        # Logits columns should NOT be dropped
        cols = set(pq.read_table(str(dest)).column_names)
        assert "shard_index_logits" in cols
        assert "row_offset_logits" in cols


# =============================================================================
# NEW: Additional coverage tests
# =============================================================================


class TestNewCheckFormatVersion:
    """Cover _check_format_version minor-version warning and edge cases."""

    def test_minor_version_warning(self):
        """Minor version mismatch produces a warning when check_minor=True."""
        import warnings as _w

        from lmprobe.sharing import _check_format_version

        info = {"format_version": "2.99"}
        with _w.catch_warnings(record=True) as caught:
            _w.simplefilter("always")
            _check_format_version(info, check_minor=True)
        assert len(caught) == 1
        assert "newer than supported" in str(caught[0].message)

    def test_minor_version_no_warn_when_disabled(self):
        """No warning when check_minor=False."""
        import warnings as _w

        from lmprobe.sharing import _check_format_version

        info = {"format_version": "2.99"}
        with _w.catch_warnings(record=True) as caught:
            _w.simplefilter("always")
            _check_format_version(info, check_minor=False)
        assert len(caught) == 0

    def test_major_version_mismatch_raises(self):
        from lmprobe.sharing import _check_format_version

        info = {"format_version": "99.0"}
        with pytest.raises(ValueError, match="Incompatible format version"):
            _check_format_version(info)

    def test_same_version_no_warning(self):
        import warnings as _w

        from lmprobe.sharing import _check_format_version

        info = {"format_version": FORMAT_VERSION}
        with _w.catch_warnings(record=True) as caught:
            _w.simplefilter("always")
            _check_format_version(info, check_minor=True)
        assert len(caught) == 0

    def test_missing_version_defaults_to_1_0(self):
        """Missing format_version defaults to 1.0 (no error for major <= local)."""
        from lmprobe.sharing import _check_format_version

        _check_format_version({})  # Should not raise


class TestNewExtractModelName:
    """Cover _extract_model_name error path."""

    def test_missing_model_raises(self):
        from lmprobe.sharing import _extract_model_name

        with pytest.raises(KeyError, match="model.name"):
            _extract_model_name({})

    def test_non_dict_model_raises(self):
        from lmprobe.sharing import _extract_model_name

        with pytest.raises(KeyError, match="model.name"):
            _extract_model_name({"model": "just-a-string"})

    def test_valid_model_returns_name(self):
        from lmprobe.sharing import _extract_model_name

        name = _extract_model_name({"model": {"name": "foo/bar"}})
        assert name == "foo/bar"


class TestNewHiddenShardFilename:
    """Cover _hidden_shard_filename."""

    def test_basic_formatting(self):
        from lmprobe.sharing import _hidden_shard_filename

        assert _hidden_shard_filename(0, 0) == (
            "tensors/hidden_layer000_shard000.safetensors"
        )
        assert _hidden_shard_filename(12, 5) == (
            "tensors/hidden_layer012_shard005.safetensors"
        )


class TestNewSelectTensorTypes:
    """Cover _select_tensor_types."""

    def test_no_filter_returns_all(self):
        from lmprobe.sharing import _select_tensor_types

        descriptors = {"hidden_layers": {}, "logits_topk": {}}
        result = _select_tensor_types(descriptors, None)
        assert result == ["hidden_layers", "logits_topk"]

    def test_filter_returns_subset(self):
        from lmprobe.sharing import _select_tensor_types

        descriptors = {"hidden_layers": {}, "logits_topk": {}}
        result = _select_tensor_types(descriptors, ["logits_topk"])
        assert result == ["logits_topk"]

    def test_filter_with_nonexistent_key(self):
        from lmprobe.sharing import _select_tensor_types

        descriptors = {"hidden_layers": {}}
        result = _select_tensor_types(descriptors, ["nonexistent"])
        assert result == []


class TestNewComputeShardBoundaries:
    """Cover _compute_shard_boundaries edge cases."""

    def test_zero_row_bytes(self):
        from lmprobe.sharing import _compute_shard_boundaries

        result = _compute_shard_boundaries(10, 0, 1000)
        assert result == [10]

    def test_negative_row_bytes(self):
        from lmprobe.sharing import _compute_shard_boundaries

        result = _compute_shard_boundaries(5, -1, 1000)
        assert result == [5]

    def test_single_shard(self):
        from lmprobe.sharing import _compute_shard_boundaries

        result = _compute_shard_boundaries(3, 100, 1000)
        assert result == [3]

    def test_multiple_shards(self):
        from lmprobe.sharing import _compute_shard_boundaries

        result = _compute_shard_boundaries(10, 100, 300)
        assert sum(result) == 10
        assert all(s <= 3 for s in result)


class TestNewDeterministicShuffle:
    """Cover _deterministic_seed and _shuffle_indices."""

    def test_seed_is_deterministic(self):
        from lmprobe.sharing import _deterministic_seed

        s1 = _deterministic_seed("user/test")
        s2 = _deterministic_seed("user/test")
        assert s1 == s2

    def test_different_repos_different_seeds(self):
        from lmprobe.sharing import _deterministic_seed

        s1 = _deterministic_seed("user/a")
        s2 = _deterministic_seed("user/b")
        assert s1 != s2

    def test_shuffle_is_permutation(self):
        from lmprobe.sharing import _shuffle_indices

        result = _shuffle_indices(5, seed=42)
        assert sorted(result) == [0, 1, 2, 3, 4]

    def test_shuffle_is_deterministic(self):
        from lmprobe.sharing import _shuffle_indices

        r1 = _shuffle_indices(10, seed=42)
        r2 = _shuffle_indices(10, seed=42)
        assert r1 == r2


class TestNewParallelPreload:
    """Cover _parallel_preload basic execution."""

    def test_basic_execution(self):
        from lmprobe.sharing import _parallel_preload

        prompts = ["a", "b", "c"]

        def load_fn(idx, prompt):
            return (idx, prompt.upper())

        result = _parallel_preload(
            prompts, load_fn, lambda: None, workers=2,
        )
        assert result == ["A", "B", "C"]

    def test_with_executor(self):
        from concurrent.futures import ThreadPoolExecutor

        from lmprobe.sharing import _parallel_preload

        prompts = ["x", "y"]

        def load_fn(idx, prompt):
            return (idx, len(prompt))

        with ThreadPoolExecutor(max_workers=1) as pool:
            result = _parallel_preload(
                prompts, load_fn, lambda: 0, workers=1,
                executor=pool,
            )
        assert result == [1, 1]


class TestNewConsolidatePreloadValidation:
    """Cover preload parameter validation in _consolidate_and_shard."""

    def test_invalid_preload_raises(self, populated_cache):
        tensor_types = {
            "raw_layers": [],
            "pooled": {"last_token": [0]},
            "has_logits": False,
            "logits_top_k": None,
            "has_perplexity": False,
        }
        with pytest.raises(ValueError, match="preload must be one of"):
            _consolidate_and_shard(
                model_name=TEST_MODEL,
                prompts=populated_cache,
                kept_indices=[0, 1, 2],
                tensor_types=tensor_types,
                labels=None,
                shard_max_bytes=1_000_000_000,
                repo_id="user/test",
                preload="invalid",
            )


@requires_pyarrow
class TestNewExtractMetadataFromParquet:
    """Cover _extract_metadata_from_parquet roundtrip and edge cases."""

    def test_roundtrip(self, tmp_path):
        import pyarrow as pa
        import pyarrow.parquet as pq

        from lmprobe.sharing import (
            _embed_metadata_in_schema,
            _extract_metadata_from_parquet,
        )

        info = {
            "format_version": "2.0",
            "model": {"name": "test/model"},
            "num_prompts": 5,
            "tensors": {"hidden_layers": {"layers": [0, 1]}},
        }

        table = pa.table({"x": pa.array([1, 2])})
        table = _embed_metadata_in_schema(table, info)
        path = tmp_path / "test.parquet"
        pq.write_table(table, str(path))

        extracted = _extract_metadata_from_parquet(path)
        assert extracted is not None
        assert extracted["format_version"] == "2.0"
        assert extracted["model"]["name"] == "test/model"
        assert extracted["num_prompts"] == 5
        assert extracted["tensors"]["hidden_layers"]["layers"] == [0, 1]

    def test_no_metadata_returns_none(self, tmp_path):
        import pyarrow as pa
        import pyarrow.parquet as pq

        from lmprobe.sharing import _extract_metadata_from_parquet

        table = pa.table({"x": pa.array([1])})
        path = tmp_path / "plain.parquet"
        pq.write_table(table, str(path))

        assert _extract_metadata_from_parquet(path) is None

    def test_numeric_format_version_coerced_to_string(self, tmp_path):
        """Numeric format_version is coerced to string on extraction."""
        import pyarrow as pa
        import pyarrow.parquet as pq

        from lmprobe.sharing import _extract_metadata_from_parquet

        # Manually embed a numeric format_version
        meta = {
            b"lmprobe:format_version": b"2",
        }
        table = pa.table({"x": pa.array([1])})
        table = table.replace_schema_metadata(meta)
        path = tmp_path / "numver.parquet"
        pq.write_table(table, str(path))

        extracted = _extract_metadata_from_parquet(path)
        assert extracted is not None
        assert extracted["format_version"] == "2"
        assert isinstance(extracted["format_version"], str)


@requires_pyarrow
class TestNewLoadDatasetMetadata:
    """Cover _load_dataset_metadata v1 format error."""

    @patch("lmprobe.sharing._check_hub_deps")
    def test_v1_format_raises(self, mock_deps, tmp_path):
        import pyarrow as pa
        import pyarrow.parquet as pq

        from lmprobe.sharing import _load_dataset_metadata

        # Write a parquet file with NO lmprobe metadata (v1-style)
        (tmp_path / "index").mkdir(parents=True)
        table = pa.table({"text": pa.array(["a"])})
        pq.write_table(table, str(tmp_path / PARQUET_PATH))

        def mock_download(repo_id, filename, **kwargs):
            return str(tmp_path / filename)

        with patch(
            "huggingface_hub.hf_hub_download", side_effect=mock_download,
        ):
            with pytest.raises(ValueError, match="v1 format"):
                _load_dataset_metadata("user/v1-dataset")


@requires_pyarrow
class TestNewFetchDatasetMetadata:
    """Cover fetch_dataset_metadata."""

    def _setup_v2_parquet(self, tmp_path):
        import pyarrow as pa
        import pyarrow.parquet as pq

        from lmprobe.sharing import _embed_metadata_in_schema

        (tmp_path / "index").mkdir(parents=True)

        lmprobe_info = {
            "format_version": "2.0",
            "model": {"name": TEST_MODEL, "revision": None},
            "num_prompts": 2,
            "tensors": {
                "hidden_layers": {
                    "type": "hidden",
                    "layers": [0, 1],
                    "dim": HIDDEN_DIM,
                    "layout": "per_layer",
                    "pooling": "last_token",
                    "shards": [{"num_prompts": 2}],
                }
            },
            "provenance": {},
        }

        table = pa.table({
            "text": pa.array(["hello", "world"], type=pa.string()),
            "label": pa.array([1, 0], type=pa.int32()),
            "num_tokens": pa.array([3, 4], type=pa.int32()),
            "shard_index": pa.array([0, 0], type=pa.int32()),
            "row_offset": pa.array([0, 1], type=pa.int32()),
        })
        table = _embed_metadata_in_schema(table, lmprobe_info)
        pq.write_table(table, str(tmp_path / PARQUET_PATH))
        return tmp_path

    @patch("lmprobe.sharing._check_hub_deps")
    def test_returns_dataset_metadata(self, mock_deps, tmp_path):
        from lmprobe.sharing import DatasetMetadata, fetch_dataset_metadata

        remote_dir = self._setup_v2_parquet(tmp_path / "remote")

        def mock_download(repo_id, filename, **kwargs):
            return str(remote_dir / filename)

        with patch(
            "huggingface_hub.hf_hub_download", side_effect=mock_download,
        ):
            meta = fetch_dataset_metadata("user/test")

        assert isinstance(meta, DatasetMetadata)
        assert meta.model_name == TEST_MODEL
        assert meta.available_layers == [0, 1]
        assert meta.num_prompts == 2
        assert meta.format_version == "2.0"
        assert meta.prompts == ["hello", "world"]
        assert "hidden_layers" in meta.tensor_descriptors


@requires_pyarrow
class TestNewLoadLabelsForPrompts:
    """Cover _load_labels_for_prompts."""

    def _setup_parquet_with_labels(self, tmp_path, labels):
        import pyarrow as pa
        import pyarrow.parquet as pq

        from lmprobe.sharing import _embed_metadata_in_schema

        (tmp_path / "index").mkdir(parents=True)

        lmprobe_info = {
            "format_version": "2.0",
            "model": {"name": TEST_MODEL, "revision": None},
            "tensors": {},
            "provenance": {},
        }

        table = pa.table({
            "text": pa.array(["a", "b", "c"], type=pa.string()),
            "label": pa.array(labels, type=pa.int32()),
            "num_tokens": pa.array([1, 1, 1], type=pa.int32()),
            "shard_index": pa.array([0, 0, 0], type=pa.int32()),
            "row_offset": pa.array([0, 1, 2], type=pa.int32()),
        })
        table = _embed_metadata_in_schema(table, lmprobe_info)
        pq.write_table(table, str(tmp_path / PARQUET_PATH))
        return tmp_path

    @patch("lmprobe.sharing._check_hub_deps")
    def test_returns_labels_array(self, mock_deps, tmp_path):
        import numpy as np

        from lmprobe.sharing import _load_labels_for_prompts

        remote_dir = self._setup_parquet_with_labels(
            tmp_path / "remote", [1, 0, 1]
        )

        def mock_download(repo_id, filename, **kwargs):
            return str(remote_dir / filename)

        with patch(
            "huggingface_hub.hf_hub_download", side_effect=mock_download,
        ):
            labels = _load_labels_for_prompts(
                "user/test", ["a", "b", "c"]
            )

        assert labels is not None
        np.testing.assert_array_equal(labels, [1, 0, 1])

    @patch("lmprobe.sharing._check_hub_deps")
    def test_returns_none_for_missing_prompt(self, mock_deps, tmp_path):
        from lmprobe.sharing import _load_labels_for_prompts

        remote_dir = self._setup_parquet_with_labels(
            tmp_path / "remote", [1, 0, 1]
        )

        def mock_download(repo_id, filename, **kwargs):
            return str(remote_dir / filename)

        with patch(
            "huggingface_hub.hf_hub_download", side_effect=mock_download,
        ):
            labels = _load_labels_for_prompts(
                "user/test", ["a", "missing"]
            )

        assert labels is None


@requires_pyarrow
class TestNewLoadActivations:
    """Cover load_activations function."""

    def _setup_remote(self, tmp_path):
        """Create a minimal v2 dataset for testing load_activations."""
        import pyarrow as pa
        import pyarrow.parquet as pq
        from safetensors.torch import save_file

        from lmprobe.sharing import _embed_metadata_in_schema

        (tmp_path / "tensors").mkdir(parents=True)
        (tmp_path / "index").mkdir(parents=True)

        prompts = ["p0", "p1", "p2"]
        tensor = torch.randn(3, HIDDEN_DIM)
        save_file(
            {"hidden.layer_0": tensor},
            str(tmp_path / "tensors" / "hidden_layer000_shard000.safetensors"),
        )

        lmprobe_info = {
            "format_version": "2.0",
            "model": {"name": TEST_MODEL, "revision": None},
            "num_prompts": 3,
            "prompt_ordering": "random",
            "tensors": {
                "hidden_layers": {
                    "type": "hidden",
                    "layers": [0],
                    "dim": HIDDEN_DIM,
                    "dtype": "float32",
                    "layout": "per_layer",
                    "pooling": "last_token",
                    "row_bytes": HIDDEN_DIM * 4,
                    "shards": [{"num_prompts": 3}],
                }
            },
            "provenance": {},
        }

        table = pa.table({
            "text": pa.array(prompts, type=pa.string()),
            "label": pa.array([1, 0, 1], type=pa.int32()),
            "num_tokens": pa.array([5, 5, 5], type=pa.int32()),
            "shard_index": pa.array([0, 0, 0], type=pa.int32()),
            "row_offset": pa.array([0, 1, 2], type=pa.int32()),
        })
        table = _embed_metadata_in_schema(table, lmprobe_info)
        pq.write_table(table, str(tmp_path / PARQUET_PATH))

        return tmp_path, prompts, tensor

    @patch("lmprobe.sharing._check_hub_deps")
    def test_load_activations_as_dict(self, mock_deps, tmp_path, cache_dir):
        from lmprobe.sharing import load_activations

        remote_dir, prompts, tensor = self._setup_remote(tmp_path / "remote")

        def mock_download(repo_id, filename, **kwargs):
            return str(remote_dir / filename)

        with patch(
            "huggingface_hub.hf_hub_download", side_effect=mock_download,
        ):
            result = load_activations(
                "user/test", show_progress=False
            )

        assert isinstance(result, dict)
        assert 0 in result
        assert result[0].shape == (3, HIDDEN_DIM)

    @patch("lmprobe.sharing._check_hub_deps")
    def test_load_activations_as_array(self, mock_deps, tmp_path, cache_dir):
        from lmprobe.sharing import load_activations

        remote_dir, prompts, tensor = self._setup_remote(tmp_path / "remote")

        def mock_download(repo_id, filename, **kwargs):
            return str(remote_dir / filename)

        with patch(
            "huggingface_hub.hf_hub_download", side_effect=mock_download,
        ):
            result = load_activations(
                "user/test", as_dict=False, show_progress=False,
            )

        import numpy as np

        assert isinstance(result, np.ndarray)
        assert result.shape == (3, 1, HIDDEN_DIM)

    @patch("lmprobe.sharing._check_hub_deps")
    def test_load_activations_with_labels(
        self, mock_deps, tmp_path, cache_dir,
    ):
        from lmprobe.sharing import load_activations

        remote_dir, prompts, tensor = self._setup_remote(tmp_path / "remote")

        def mock_download(repo_id, filename, **kwargs):
            return str(remote_dir / filename)

        with patch(
            "huggingface_hub.hf_hub_download", side_effect=mock_download,
        ):
            result, labels = load_activations(
                "user/test", return_labels=True, show_progress=False,
            )

        import numpy as np

        assert isinstance(result, dict)
        assert labels is not None
        assert len(labels) == 3
        np.testing.assert_array_equal(labels, [1, 0, 1])

    @patch("lmprobe.sharing._check_hub_deps")
    def test_load_activations_invalid_layer_raises(
        self, mock_deps, tmp_path, cache_dir,
    ):
        from lmprobe.sharing import load_activations

        remote_dir, prompts, tensor = self._setup_remote(tmp_path / "remote")

        def mock_download(repo_id, filename, **kwargs):
            return str(remote_dir / filename)

        with patch(
            "huggingface_hub.hf_hub_download", side_effect=mock_download,
        ):
            with pytest.raises(ValueError, match="Layers.*not in dataset"):
                load_activations(
                    "user/test", layers=[99], show_progress=False,
                )

    @patch("lmprobe.sharing._check_hub_deps")
    def test_load_activations_specific_prompts(
        self, mock_deps, tmp_path, cache_dir,
    ):
        from lmprobe.sharing import load_activations

        remote_dir, prompts, tensor = self._setup_remote(tmp_path / "remote")

        def mock_download(repo_id, filename, **kwargs):
            return str(remote_dir / filename)

        with patch(
            "huggingface_hub.hf_hub_download", side_effect=mock_download,
        ):
            result = load_activations(
                "user/test", prompts=["p0", "p1"],
                show_progress=False,
            )

        assert 0 in result
        assert result[0].shape == (2, HIDDEN_DIM)


@requires_pyarrow
class TestNewWriteParquetExtraTypes:
    """Cover extra column type inference in _write_parquet_index."""

    def test_list_of_int_columns(self, tmp_path):
        import pyarrow.parquet as pq

        (tmp_path / "index").mkdir()
        prompt_metadata = [
            {"text": "a", "label": 1, "num_tokens": 3,
             "shard_index": 0, "row_offset": 0,
             "token_shard_ids": [0, 0, 1]},
            {"text": "b", "label": 0, "num_tokens": 2,
             "shard_index": 0, "row_offset": 1,
             "token_shard_ids": [0, 1]},
        ]
        _write_parquet_index(tmp_path, prompt_metadata)

        table = pq.read_table(str(tmp_path / PARQUET_PATH))
        assert "token_shard_ids" in table.column_names
        assert table.column("token_shard_ids").to_pylist() == [
            [0, 0, 1], [0, 1],
        ]

    def test_list_of_float_columns(self, tmp_path):
        import pyarrow.parquet as pq

        (tmp_path / "index").mkdir()
        prompt_metadata = [
            {"text": "a", "label": 1, "num_tokens": 3,
             "shard_index": 0, "row_offset": 0,
             "token_perplexity": [1.5, 2.3]},
            {"text": "b", "label": 0, "num_tokens": 2,
             "shard_index": 0, "row_offset": 1,
             "token_perplexity": [0.9]},
        ]
        _write_parquet_index(tmp_path, prompt_metadata)

        table = pq.read_table(str(tmp_path / PARQUET_PATH))
        assert "token_perplexity" in table.column_names

    def test_list_of_string_columns(self, tmp_path):
        import pyarrow.parquet as pq

        (tmp_path / "index").mkdir()
        prompt_metadata = [
            {"text": "a", "label": 1, "num_tokens": 3,
             "shard_index": 0, "row_offset": 0,
             "token_strings": ["hello", "world"]},
            {"text": "b", "label": 0, "num_tokens": 2,
             "shard_index": 0, "row_offset": 1,
             "token_strings": ["foo"]},
        ]
        _write_parquet_index(tmp_path, prompt_metadata)

        table = pq.read_table(str(tmp_path / PARQUET_PATH))
        assert "token_strings" in table.column_names
        assert table.column("token_strings").to_pylist() == [
            ["hello", "world"], ["foo"],
        ]

    def test_null_labels(self, tmp_path):
        """All-None labels still produce a valid Parquet column."""
        import pyarrow.parquet as pq

        (tmp_path / "index").mkdir()
        prompt_metadata = [
            {"text": "a", "label": None, "num_tokens": 3,
             "shard_index": 0, "row_offset": 0},
        ]
        _write_parquet_index(tmp_path, prompt_metadata)

        table = pq.read_table(str(tmp_path / PARQUET_PATH))
        assert table.num_rows == 1

    def test_embedded_metadata_in_parquet(self, tmp_path):
        """lmprobe_info is embedded in Parquet schema when provided."""

        from lmprobe.sharing import _extract_metadata_from_parquet

        (tmp_path / "index").mkdir()
        lmprobe_info = {
            "format_version": "2.0",
            "model": {"name": "test"},
            "tensors": {},
        }
        prompt_metadata = [
            {"text": "x", "label": 1, "num_tokens": 3,
             "shard_index": 0, "row_offset": 0},
        ]
        _write_parquet_index(tmp_path, prompt_metadata, lmprobe_info)

        extracted = _extract_metadata_from_parquet(
            tmp_path / PARQUET_PATH,
        )
        assert extracted is not None
        assert extracted["format_version"] == "2.0"

    def test_shard_index_logits_column(self, tmp_path):
        """shard_index_logits and row_offset_logits as per-type columns."""
        import pyarrow.parquet as pq

        (tmp_path / "index").mkdir()
        prompt_metadata = [
            {"text": "a", "label": 1, "num_tokens": 3,
             "shard_index": 0, "row_offset": 0,
             "shard_index_logits": 1, "row_offset_logits": 0},
            {"text": "b", "label": 0, "num_tokens": 4,
             "shard_index": 0, "row_offset": 1,
             "shard_index_logits": 1, "row_offset_logits": 1},
        ]
        _write_parquet_index(tmp_path, prompt_metadata)

        table = pq.read_table(str(tmp_path / PARQUET_PATH))
        assert "shard_index_logits" in table.column_names
        assert "row_offset_logits" in table.column_names
        assert table.column("shard_index_logits").to_pylist() == [1, 1]

    def test_all_none_extra_column(self, tmp_path):
        """Extra column with all None values defaults to string type."""
        import pyarrow as pa
        import pyarrow.parquet as pq

        (tmp_path / "index").mkdir()
        prompt_metadata = [
            {"text": "a", "label": 1, "num_tokens": 3,
             "shard_index": 0, "row_offset": 0,
             "optional_field": None},
            {"text": "b", "label": 0, "num_tokens": 4,
             "shard_index": 0, "row_offset": 1,
             "optional_field": None},
        ]
        _write_parquet_index(tmp_path, prompt_metadata)

        table = pq.read_table(str(tmp_path / PARQUET_PATH))
        assert "optional_field" in table.column_names
        assert table.column("optional_field").type == pa.string()


class TestNewFilterTensorTypesEdgeCases:
    """Cover _filter_tensor_types unknown key warning and perplexity."""

    def test_unknown_key_logged(self, caplog):
        import logging

        available = {
            "raw_layers": [0],
            "pooled": {},
            "has_logits": False,
            "logits_top_k": None,
            "has_perplexity": False,
            "has_token_perplexity": False,
        }
        with caplog.at_level(logging.WARNING, logger="lmprobe.sharing"):
            result = _filter_tensor_types(available, ["bogus_key"])
        assert "Unknown tensor filter key" in caplog.text
        # Result should have nothing enabled
        assert result["raw_layers"] == []
        assert result["pooled"] == {}
        assert result["has_logits"] is False

    def test_perplexity_filter(self):
        available = {
            "raw_layers": [0],
            "pooled": {"last_token": [0]},
            "has_logits": False,
            "logits_top_k": None,
            "has_perplexity": True,
            "has_token_perplexity": True,
        }
        result = _filter_tensor_types(available, ["perplexity"])
        assert result["has_perplexity"] is True
        assert result["has_token_perplexity"] is True
        # Other types should be excluded
        assert result["raw_layers"] == []
        assert result["pooled"] == {}


class TestNewPushNoTensorsRaises:
    """Cover push_dataset error when no tensor types remain after filter."""

    @requires_pyarrow
    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    def test_no_tensor_types_raises(
        self, mock_pyarrow, mock_deps, populated_cache,
    ):
        """Filtering to nonexistent tensor types raises ValueError."""
        with pytest.raises(ValueError, match="No tensor types available"):
            push_dataset(
                repo_id="user/test",
                model_name=TEST_MODEL,
                prompts=populated_cache,
                tensors=["logits_topk"],  # cache has no logits
            )


@requires_pyarrow
class TestNewPushDatasetShuffleParam:
    """Cover push_dataset shuffle=False param forwarding."""

    @patch("lmprobe.sharing._check_hub_deps")
    @patch("lmprobe.sharing._check_pyarrow")
    @patch("huggingface_hub.HfApi")
    def test_push_shuffle_false(
        self, MockHfApi, mock_pyarrow, mock_deps, populated_cache,
    ):
        """shuffle=False is forwarded and preserves order."""
        mock_api = MagicMock()
        MockHfApi.return_value = mock_api

        uploaded_files = {}

        def capture_upload(repo_id, folder_path, **kwargs):
            import pyarrow.parquet as pq

            folder = Path(folder_path)
            parquet_file = folder / PARQUET_PATH
            if parquet_file.exists():
                table = pq.read_table(str(parquet_file))
                uploaded_files["texts"] = table.column("text").to_pylist()

        mock_api.upload_large_folder.side_effect = capture_upload

        push_dataset(
            repo_id="user/noshuffle",
            model_name=TEST_MODEL,
            prompts=populated_cache,
            exist_ok=True,
            shuffle=False,
        )

        assert uploaded_files["texts"] == populated_cache


class TestNewManifestHelpers:
    """Cover _new_manifest, _save_manifest, _load_manifest."""

    def test_new_manifest_structure(self):
        m = _new_manifest("user/test")
        assert m["format"] == "stream_manifest_v1"
        assert m["repo_id"] == "user/test"
        assert m["completed_shards"] == []
        assert m["metadata_uploaded"] is False

    def test_save_and_load_manifest(self, tmp_path):
        m = _new_manifest("user/test")
        m["completed_shards"] = ["shard_0"]
        _save_manifest(tmp_path, m)

        loaded = _load_manifest(tmp_path)
        assert loaded is not None
        assert loaded["completed_shards"] == ["shard_0"]

    def test_load_missing_manifest(self, tmp_path):
        assert _load_manifest(tmp_path) is None


class TestNewStagingDirPath:
    """Cover _staging_dir_path deterministic hashing."""

    def test_deterministic(self, cache_dir):
        p1 = _staging_dir_path("user/test", "model", ["a", "b"])
        p2 = _staging_dir_path("user/test", "model", ["a", "b"])
        assert p1 == p2

    def test_different_params_different_path(self, cache_dir):
        p1 = _staging_dir_path("user/test", "model", ["a"])
        p2 = _staging_dir_path("user/test", "model", ["a", "b"])
        assert p1 != p2

    def test_stream_changes_path(self, cache_dir):
        p1 = _staging_dir_path("user/test", "model", ["a"])
        p2 = _staging_dir_path("user/test", "model", ["a"], stream=True)
        assert p1 != p2

    def test_labels_change_path(self, cache_dir):
        p1 = _staging_dir_path("user/test", "model", ["a"])
        p2 = _staging_dir_path("user/test", "model", ["a"], labels=[1])
        assert p1 != p2


class TestNewMakeHiddenLoader:
    """Cover _make_hidden_loader and preload paths."""

    def test_pooled_loader(self, populated_cache):
        from lmprobe.sharing import _make_hidden_loader

        loader = _make_hidden_loader(
            TEST_MODEL, [0], use_raw=False,
            hidden_strategy="last_token",
        )
        idx, result = loader(0, populated_cache[0])
        assert idx == 0
        assert isinstance(result, dict)
        assert "hidden.layer_0" in result

    def test_pooled_loader_with_extract_key(self, populated_cache):
        from lmprobe.sharing import _make_hidden_loader

        loader = _make_hidden_loader(
            TEST_MODEL, [0], use_raw=False,
            hidden_strategy="last_token",
            extract_key="hidden.layer_0",
        )
        idx, result = loader(0, populated_cache[0])
        assert idx == 0
        assert result is not None
        assert result.shape[-1] == HIDDEN_DIM

    def test_no_strategy_returns_empty(self, populated_cache):
        from lmprobe.sharing import _make_hidden_loader

        loader = _make_hidden_loader(
            TEST_MODEL, [0], use_raw=False,
            hidden_strategy=None,
        )
        idx, result = loader(0, populated_cache[0])
        assert idx == 0
        assert result == {}

    def test_raw_loader(self, populated_raw_cache):
        from lmprobe.sharing import _make_hidden_loader

        prompts, seq_lens, layers, _ = populated_raw_cache
        loader = _make_hidden_loader(
            TEST_MODEL, layers, use_raw=True,
            hidden_strategy=None,
        )
        idx, result = loader(0, prompts[0])
        assert idx == 0
        assert "hidden.layer_0" in result


class TestNewPreloadLayer:
    """Cover _preload_layer."""

    def test_preload_one_layer(self, populated_cache):
        from lmprobe.sharing import _preload_layer

        result = _preload_layer(
            TEST_MODEL, populated_cache, layer=0,
            use_raw=False, hidden_strategy="last_token",
            workers=2,
        )
        assert len(result) == 3
        for tensor in result:
            assert tensor is not None
            assert tensor.shape[-1] == HIDDEN_DIM


class TestNewPreloadFull:
    """Cover _preload_full."""

    def test_preload_all_layers(self, populated_cache):
        from lmprobe.sharing import _preload_full

        result = _preload_full(
            TEST_MODEL, populated_cache, layers=[0, 1],
            use_raw=False, hidden_strategy="last_token",
            workers=2,
        )
        assert len(result) == 3
        for entry in result:
            assert "hidden.layer_0" in entry
            assert "hidden.layer_1" in entry


class TestNewConsolidatePreloadModes:
    """Cover preload='per_layer' and preload='full' in _consolidate_and_shard."""

    def test_preload_per_layer(self, populated_cache):
        tensor_types = {
            "raw_layers": [],
            "pooled": {"last_token": [0, 1]},
            "has_logits": False,
            "logits_top_k": None,
            "has_perplexity": False,
        }
        tmpdir, tensor_descriptors, prompt_metadata = _consolidate_and_shard(
            model_name=TEST_MODEL,
            prompts=populated_cache,
            kept_indices=[0, 1, 2],
            tensor_types=tensor_types,
            labels=None,
            shard_max_bytes=1_000_000_000,
            repo_id="user/test-per-layer",
            preload="per_layer",
            preload_workers=2,
        )
        assert len(prompt_metadata) == 3
        assert "hidden_layers" in tensor_descriptors
        for layer in [0, 1]:
            fname = f"tensors/hidden_layer{layer:03d}_shard000.safetensors"
            assert (tmpdir / fname).exists()
        shutil.rmtree(tmpdir, ignore_errors=True)

    def test_preload_full(self, populated_cache):
        tensor_types = {
            "raw_layers": [],
            "pooled": {"last_token": [0, 1]},
            "has_logits": False,
            "logits_top_k": None,
            "has_perplexity": False,
        }
        tmpdir, tensor_descriptors, prompt_metadata = _consolidate_and_shard(
            model_name=TEST_MODEL,
            prompts=populated_cache,
            kept_indices=[0, 1, 2],
            tensor_types=tensor_types,
            labels=None,
            shard_max_bytes=1_000_000_000,
            repo_id="user/test-full",
            preload="full",
            preload_workers=2,
        )
        assert len(prompt_metadata) == 3
        assert "hidden_layers" in tensor_descriptors
        for layer in [0, 1]:
            fname = f"tensors/hidden_layer{layer:03d}_shard000.safetensors"
            assert (tmpdir / fname).exists()
        shutil.rmtree(tmpdir, ignore_errors=True)


@requires_pyarrow
class TestNewLoadActivationDatasetColocated:
    """Cover v1.0 co-located layout in load_activation_dataset."""

    @patch("lmprobe.sharing._check_hub_deps")
    def test_colocated_layout(self, mock_deps, tmp_path):
        """load_activation_dataset handles co-located (v1.0) layout."""
        import pyarrow as pa
        import pyarrow.parquet as pq
        from safetensors.torch import save_file

        from lmprobe.sharing import _embed_metadata_in_schema

        (tmp_path / "tensors").mkdir(parents=True)
        (tmp_path / "index").mkdir(parents=True)

        tensor = torch.randn(2, HIDDEN_DIM)
        save_file(
            {"hidden.layer_0": tensor},
            str(tmp_path / "tensors" / "shard_000.safetensors"),
        )

        lmprobe_info = {
            "format_version": "2.0",
            "model": {"name": TEST_MODEL, "revision": None},
            "num_prompts": 2,
            "tensors": {
                "hidden_layers": {
                    "type": "hidden",
                    "layers": [0],
                    "dim": HIDDEN_DIM,
                    "dtype": "float32",
                    "layout": "co_located",
                    "pooling": "last_token",
                    "row_bytes": HIDDEN_DIM * 4,
                    "shards": [
                        {
                            "file": "tensors/shard_000.safetensors",
                            "num_prompts": 2,
                        }
                    ],
                }
            },
            "provenance": {},
        }

        table = pa.table({
            "text": pa.array(["p0", "p1"], type=pa.string()),
            "label": pa.array([1, 0], type=pa.int32()),
            "num_tokens": pa.array([5, 5], type=pa.int32()),
            "shard_index": pa.array([0, 0], type=pa.int32()),
            "row_offset": pa.array([0, 1], type=pa.int32()),
        })
        table = _embed_metadata_in_schema(table, lmprobe_info)
        pq.write_table(table, str(tmp_path / PARQUET_PATH))

        def mock_download(repo_id, filename, **kwargs):
            return str(tmp_path / filename)

        import warnings as _w

        with (
            patch(
                "huggingface_hub.hf_hub_download",
                side_effect=mock_download,
            ),
            _w.catch_warnings(record=True),
        ):
            result, info = load_activation_dataset(
                "user/test-colocated",
                layers=[0],
            )

        assert "hidden.layer_0" in result
        assert result["hidden.layer_0"].shape == (2, HIDDEN_DIM)
        assert torch.allclose(result["hidden.layer_0"], tensor)

    @patch("lmprobe.sharing._check_hub_deps")
    def test_colocated_layer_filter_warns(self, mock_deps, tmp_path):
        """layers param on co-located dataset produces a warning."""
        import pyarrow as pa
        import pyarrow.parquet as pq
        from safetensors.torch import save_file

        from lmprobe.sharing import _embed_metadata_in_schema

        (tmp_path / "tensors").mkdir(parents=True)
        (tmp_path / "index").mkdir(parents=True)

        save_file(
            {"hidden.layer_0": torch.randn(1, HIDDEN_DIM)},
            str(tmp_path / "tensors" / "shard_000.safetensors"),
        )

        lmprobe_info = {
            "format_version": "2.0",
            "model": {"name": TEST_MODEL, "revision": None},
            "num_prompts": 1,
            "tensors": {
                "hidden_layers": {
                    "type": "hidden",
                    "layers": [0],
                    "dim": HIDDEN_DIM,
                    "dtype": "float32",
                    "layout": "co_located",
                    "pooling": "last_token",
                    "row_bytes": HIDDEN_DIM * 4,
                    "shards": [
                        {
                            "file": "tensors/shard_000.safetensors",
                            "num_prompts": 1,
                        }
                    ],
                }
            },
            "provenance": {},
        }

        table = pa.table({
            "text": pa.array(["p0"], type=pa.string()),
            "label": pa.array([1], type=pa.int32()),
            "num_tokens": pa.array([5], type=pa.int32()),
            "shard_index": pa.array([0], type=pa.int32()),
            "row_offset": pa.array([0], type=pa.int32()),
        })
        table = _embed_metadata_in_schema(table, lmprobe_info)
        pq.write_table(table, str(tmp_path / PARQUET_PATH))

        def mock_download(repo_id, filename, **kwargs):
            return str(tmp_path / filename)

        import warnings as _w

        with (
            patch(
                "huggingface_hub.hf_hub_download",
                side_effect=mock_download,
            ),
            _w.catch_warnings(record=True) as caught,
        ):
            _w.simplefilter("always")
            load_activation_dataset(
                "user/test-colocated-warn",
                layers=[0],
            )

        warn_texts = [str(w.message) for w in caught]
        assert any(
            "co-located" in t for t in warn_texts
        ), f"Expected co-located warning, got: {warn_texts}"


class TestNewComputeShardPlanEdgeCases:
    """Cover _compute_shard_plan edge cases like skipping prompts."""

    def test_skipped_prompts_warning(self, cache_dir, caplog):
        """Prompts not in cache during metadata scan are skipped."""
        import logging

        # Create cache with only 1 of 3 prompts
        prompts = ["cached", "missing1", "missing2"]
        save_prompt_pooled_activations(
            TEST_MODEL, "cached", [0], torch.randn(1, HIDDEN_DIM), "last_token"
        )

        tensor_types = {
            "pooled": {"last_token": [0]},
            "has_logits": False,
            "logits_top_k": None,
            "raw_layers": [],
        }

        with caplog.at_level(logging.WARNING, logger="lmprobe.sharing"):
            plan = _compute_shard_plan(
                model_name=TEST_MODEL,
                prompts=prompts,
                kept_indices=[0, 1, 2],
                tensor_types=tensor_types,
                labels=None,
                shard_max_bytes=1_000_000_000,
                repo_id="user/test-skip",
            )

        assert len(plan["prompt_metadata"]) == 1
        assert "Skipped" in caplog.text

    def test_all_prompts_missing_raises(self, cache_dir):
        """All prompts missing during metadata scan raises ValueError."""
        tensor_types = {
            "pooled": {"last_token": [0]},
            "has_logits": False,
            "logits_top_k": None,
            "raw_layers": [],
        }

        with pytest.raises(ValueError, match="No prompts could be loaded"):
            _compute_shard_plan(
                model_name=TEST_MODEL,
                prompts=["missing"],
                kept_indices=[0],
                tensor_types=tensor_types,
                labels=None,
                shard_max_bytes=1_000_000_000,
                repo_id="user/test-none",
            )

    def test_no_shuffle_preserves_order(self, populated_cache):
        tensor_types = {
            "pooled": {"last_token": [0]},
            "has_logits": False,
            "logits_top_k": None,
            "raw_layers": [],
        }

        plan = _compute_shard_plan(
            model_name=TEST_MODEL,
            prompts=populated_cache,
            kept_indices=list(range(len(populated_cache))),
            tensor_types=tensor_types,
            labels=None,
            shard_max_bytes=1_000_000_000,
            repo_id="user/test-noshuffle",
            shuffle=False,
        )

        output_texts = [p["text"] for p in plan["prompt_metadata"]]
        assert output_texts == populated_cache

    def test_metadata_forwarded(self, populated_cache):
        """Extra metadata dicts appear in prompt_metadata."""
        tensor_types = {
            "pooled": {"last_token": [0]},
            "has_logits": False,
            "logits_top_k": None,
            "raw_layers": [],
        }
        meta = [
            {"source": "train"},
            {"source": "train"},
            {"source": "test"},
        ]

        plan = _compute_shard_plan(
            model_name=TEST_MODEL,
            prompts=populated_cache,
            kept_indices=list(range(len(populated_cache))),
            tensor_types=tensor_types,
            labels=None,
            shard_max_bytes=1_000_000_000,
            repo_id="user/test-meta",
            metadata=meta,
            shuffle=False,
        )

        for pm in plan["prompt_metadata"]:
            assert "source" in pm
