from pathlib import Path

from schemas_hf.publish import LocalSplitSpec, publish_hf_dataset


class DummyDatasetDict:
    def __init__(self) -> None:
        self.push_args: tuple | None = None
        self.push_kwargs: dict | None = None

    def push_to_hub(self, *args, **kwargs) -> None:
        self.push_args = args
        self.push_kwargs = kwargs


def test_publish_hf_dataset_uses_load_and_push(monkeypatch, tmp_path):
    calls: dict[str, object] = {}

    def fake_load_dataset(builder: str, data_files: dict[str, list[str]]):
        calls["builder"] = builder
        calls["data_files"] = data_files
        ds = DummyDatasetDict()
        calls["dataset"] = ds
        return ds

    # Patch datasets.load_dataset inside the publish module.
    monkeypatch.setattr(
        "schemas_hf.publish.datasets.load_dataset",
        fake_load_dataset,
    )

    # Prepare a couple of dummy files.
    f1 = tmp_path / "a.jsonl"
    f2 = tmp_path / "b.jsonl"
    f1.write_text('{"id": 1}\n')
    f2.write_text('{"id": 2}\n')

    split = LocalSplitSpec(name="train", format="jsonl", paths=[f1, f2])
    publish_hf_dataset("owner/name", split, private=True, token="dummy-token")

    # Check that the correct builder and data_files were used.
    assert calls["builder"] == "json"
    data_files = calls["data_files"]
    assert isinstance(data_files, dict)
    assert "train" in data_files
    assert data_files["train"] == [str(f1), str(f2)]

    # Check that push_to_hub was called with the right arguments.
    ds = calls["dataset"]
    assert isinstance(ds, DummyDatasetDict)
    assert ds.push_args == ("owner/name",)
    assert ds.push_kwargs == {"private": True, "token": "dummy-token"}


def test_publish_hf_dataset_with_image_column(monkeypatch, tmp_path):
    calls: dict[str, object] = {}

    class DummySplit:
        def __init__(self) -> None:
            self.calls: list[tuple[str, tuple, dict]] = []
            self.column_names = ["image_path"]

        def map(self, fn, *args, **kwargs):
            self.calls.append(("map", args, kwargs))
            return self

        def cast_column(self, name, feature, *args, **kwargs):
            self.calls.append(("cast_column", (name, feature), kwargs))
            return self

        def remove_columns(self, cols):
            self.calls.append(("remove_columns", (cols,), {}))
            return self

    class DummyDatasetDict(dict):
        def __init__(self) -> None:
            super().__init__({"train": DummySplit()})
            self.push_args: tuple | None = None
            self.push_kwargs: dict | None = None

        def push_to_hub(self, *args, **kwargs):
            calls["push"] = (args, kwargs)

    def fake_load_dataset(builder: str, data_files: dict[str, list[str]]):
        calls["builder"] = builder
        calls["data_files"] = data_files
        return DummyDatasetDict()

    monkeypatch.setattr(
        "schemas_hf.publish.datasets.load_dataset",
        fake_load_dataset,
    )

    # Patch Image so we don't depend on datasets[vision] extras at test time.
    class DummyImage:
        pass

    monkeypatch.setattr(
        "schemas_hf.publish.Image",
        DummyImage,
        raising=False,
    )

    f = tmp_path / "images.parquet"
    f.write_text("dummy", encoding="utf-8")

    split = LocalSplitSpec(name="train", format="parquet", paths=[f])
    publish_hf_dataset(
        "owner/images",
        split,
        private=False,
        token=None,
        max_shard_size=None,
        image_column="image_path",
    )

    assert calls["builder"] == "parquet"
    assert isinstance(calls["data_files"]["train"], list)
    # Ensure push_to_hub was called.
    assert "push" in calls
