import json
from pathlib import Path

from typer.testing import CliRunner

from schemas_hf import cli, inspect
from schemas_hf.naming import module_filename


class StubRegistry:
    def __init__(self, names):
        self.names = names
        self.discovered = False

    def discover(self):
        self.discovered = True

    def list_names(self):
        return self.names


def test_cli_sync_dry_run(tmp_path):
    cache_dir = tmp_path / "hub"
    repo_dir = cache_dir / "datasets--owner--name" / "snapshots" / "123"
    repo_dir.mkdir(parents=True)
    (repo_dir / "dataset_info.json").write_text(json.dumps({"features": {}}))

    output_dir = tmp_path / "out"
    runner = CliRunner()
    result = runner.invoke(
        cli.app,
        ["sync", "--cache-dir", str(cache_dir), "--output", str(output_dir), "--dry-run"],
    )

    assert result.exit_code == 0
    assert "Generated 0 schema" in result.stdout
    assert not output_dir.exists()


def test_cli_instantiate_writes_file(tmp_path):
    destination = tmp_path / "schemas"
    runner = CliRunner()
    result = runner.invoke(
        cli.app,
        ["instantiate", "ag_news", "--dest", str(destination)],
    )

    assert result.exit_code == 0
    target_file = destination / "ag_news.py"
    assert target_file.exists()
    contents = target_file.read_text()
    assert "AgNewsRecord" in contents


def test_cli_instantiate_with_config(tmp_path):
    destination = tmp_path / "schemas"
    runner = CliRunner()
    result = runner.invoke(
        cli.app,
        ["instantiate", "imdb", "--config", "plain_text", "--dest", str(destination)],
    )

    assert result.exit_code == 0
    target_file = destination / "imdb__plain_text.py"
    assert target_file.exists()


def test_cli_list(monkeypatch):
    stub = StubRegistry(["ag_news", "glue:mrpc"])
    monkeypatch.setattr(cli, "registry", stub)

    runner = CliRunner()
    result = runner.invoke(cli.app, ["list"])

    assert result.exit_code == 0
    assert "ag_news" in result.stdout
    assert stub.discovered


def test_cli_info(monkeypatch):
    details = inspect.DatasetDetails(
        dataset_id="ag_news",
        config_name=None,
        description="desc",
        homepage="https://example.com",
        license="mit",
        citation="cite",
        features={"text": {"_type": "Value", "dtype": "string"}},
        splits={"train": {"num_examples": 100}},
        builder_configs=["default"],
        size_in_bytes=123,
        download_size=456,
    )

    stub = StubRegistry(["ag_news", "ai2_arc:ARC-Easy"])
    monkeypatch.setattr(cli, "registry", stub)
    monkeypatch.setattr(cli, "describe_dataset", lambda dataset, config: details)
    monkeypatch.setattr(cli, "summarize_feature", lambda spec: "string")

    runner = CliRunner()
    result = runner.invoke(cli.app, ["info", "ag_news"])

    assert result.exit_code == 0
    assert "Dataset : ag_news" in result.stdout
    assert "Columns:" in result.stdout


def test_cli_info_handles_error(monkeypatch):
    stub = StubRegistry(["ag_news"])
    monkeypatch.setattr(cli, "registry", stub)

    def boom(*args, **kwargs):
        raise inspect.DatasetInfoError("unavailable")

    monkeypatch.setattr(cli, "describe_dataset", boom)

    runner = CliRunner()
    result = runner.invoke(cli.app, ["info", "ag_news"])

    assert result.exit_code == 1
    assert "Unable to load metadata" in result.output


def test_match_schema_name_fuzzy():
    names = ["ai2_arc:ARC-Easy", "lukaemon/bbh:boolean_expressions"]
    match = cli._match_schema_name("ai-arc", names)
    assert match == "ai2_arc:ARC-Easy"
    match = cli._match_schema_name("boolean", names)
    assert match == "lukaemon/bbh:boolean_expressions"


def test_cli_init_loader(tmp_path):
    runner = CliRunner()
    result = runner.invoke(
        cli.app,
        [
            "init-loader",
            "demo_loader",
            "-d",
            "ag_news",
            "-d",
            "glue:mrpc@validation",
            "--output",
            str(tmp_path),
            "--force",
        ],
    )
    assert result.exit_code == 0, result.output
    created = tmp_path / "demo_loader.py"
    assert created.exists()
    text = created.read_text()
    assert '"id": "ag_news"' in text
    assert '"id": "glue"' in text


def test_cli_make_local_schema_jsonl(tmp_path):
    # Prepare a small JSONL file with a couple of rows.
    jsonl_path = tmp_path / "data.jsonl"
    rows = [
        {"id": 1, "text": "hello", "tags": ["a", "b"]},
        {"id": 2, "text": "world", "tags": ["c"]},
    ]
    jsonl_path.write_text("\n".join(json.dumps(r) for r in rows), encoding="utf-8")

    runner = CliRunner()
    result = runner.invoke(
        cli.app,
        [
            "make-local-schema",
            "local/demo_cli",
            "--format",
            "jsonl",
            "--path",
            str(jsonl_path),
            "--output",
            str(tmp_path),
            "--sample-rows",
            "10",
        ],
    )

    assert result.exit_code == 0, result.output

    # The schema module should have been written under the output directory.
    module_name = module_filename("local/demo_cli", None)
    module_path = tmp_path / module_name
    assert module_path.exists()

    # Import the generated module and instantiate the dataclass.
    import importlib.util
    from schemas_hf.base import DatasetSchema

    spec = importlib.util.spec_from_file_location("local_demo_schema_cli", module_path)
    assert spec is not None and spec.loader is not None
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)  # type: ignore[call-arg]

    record_cls = None
    for attr in dir(module):
        obj = getattr(module, attr)
        if isinstance(obj, type) and issubclass(obj, DatasetSchema) and obj is not DatasetSchema:
            record_cls = obj
            break

    assert record_cls is not None
    instance = record_cls(id=1, text="hello", tags=["a", "b"])
    assert instance.id == 1
    assert instance.text == "hello"
    assert instance.tags == ["a", "b"]


def test_cli_make_local_schema_no_matches(tmp_path):
    runner = CliRunner()
    result = runner.invoke(
        cli.app,
        [
            "make-local-schema",
            "local/empty",
            "--format",
            "jsonl",
            "--path",
            str(tmp_path / "does_not_exist_*.jsonl"),
            "--output",
            str(tmp_path),
        ],
    )

    assert result.exit_code != 0
    assert "No files matched the provided --path patterns." in result.output


def test_cli_publish_hf_invokes_publisher(monkeypatch, tmp_path):
    source = tmp_path / "data.jsonl"
    source.write_text('{"id": 1}\n', encoding="utf-8")

    called: dict[str, object] = {}

    def fake_publish(dataset_id, split, private=True, token=None, max_shard_size=None, image_column=None):
        called["dataset_id"] = dataset_id
        called["split"] = split
        called["private"] = private
        called["token"] = token
        called["max_shard_size"] = max_shard_size

    monkeypatch.setattr(cli, "publish_hf_dataset", fake_publish)

    runner = CliRunner()
    result = runner.invoke(
        cli.app,
        [
            "publish-hf",
            "owner/name",
            "--format",
            "jsonl",
            "--path",
            str(source),
            "--split",
            "train",
            "--private",
        ],
    )

    assert result.exit_code == 0, result.output
    assert called["dataset_id"] == "owner/name"
    split = called["split"]
    # LocalSplitSpec is defined in schemas_hf.publish, but here we only
    # assert on the attributes we care about to keep the test decoupled.
    assert getattr(split, "name") == "train"
    assert getattr(split, "format") == "jsonl"
    assert getattr(split, "paths") == [source]
    assert called["private"] is True
    assert called["max_shard_size"] is None


def test_cli_publish_hf_warns_on_asset_columns(monkeypatch, tmp_path):
    # Create a small Parquet file with an asset-like column.
    import pyarrow as pa  # type: ignore[import]
    import pyarrow.parquet as pq  # type: ignore[import]

    table = pa.table(
        {
            "image_path": ["img_align_celeba/000001.jpg"],
            "Smiling": [1],
        }
    )
    parquet_path = tmp_path / "celeba_attrs.parquet"
    pq.write_table(table, parquet_path)

    called: dict[str, object] = {}

    def fake_publish(dataset_id, split, private=True, token=None, max_shard_size=None, image_column=None):
        called["dataset_id"] = dataset_id
        called["split"] = split

    # Avoid touching HF Hub or schema generation.
    monkeypatch.setattr(cli, "publish_hf_dataset", fake_publish)
    monkeypatch.setattr(cli, "load_dataset_builder", lambda *a, **k: None)
    monkeypatch.setattr(
        cli,
        "generate_from_cache",
        lambda *a, **k: [],
    )

    runner = CliRunner()
    result = runner.invoke(
        cli.app,
        [
        "publish-hf",
        "uwunion/celeba-attrs",
        "--format",
        "parquet",
        "--no-images",
        "--path",
        str(parquet_path),
        "--split",
        "train",
        "--private",
        ],
    )

    assert result.exit_code == 0, result.output
    # Ensure the warning about asset-like columns was emitted.
    assert "asset-like columns" in result.output
    assert "image_path" in result.output


def test_cli_publish_hf_prompts_and_sets_image_column(monkeypatch, tmp_path):
    # Create a small Parquet file with an asset-like column.
    import pyarrow as pa  # type: ignore[import]
    import pyarrow.parquet as pq  # type: ignore[import]

    table = pa.table(
        {
            "image_path": ["img_align_celeba/000001.jpg"],
            "Smiling": [1],
        }
    )
    parquet_path = tmp_path / "celeba_attrs.parquet"
    pq.write_table(table, parquet_path)

    called: dict[str, object] = {}

    def fake_publish(dataset_id, split, private=True, token=None, max_shard_size=None, image_column=None):
        called["dataset_id"] = dataset_id
        called["split"] = split
        called["image_column"] = image_column

    # Avoid touching HF Hub or schema generation.
    monkeypatch.setattr(cli, "publish_hf_dataset", fake_publish)
    monkeypatch.setattr(cli, "load_dataset_builder", lambda *a, **k: None)
    monkeypatch.setattr(
        cli,
        "generate_from_cache",
        lambda *a, **k: [],
    )

    runner = CliRunner()
    # Simulate user answering "y" when asked about uploading images.
    result = runner.invoke(
        cli.app,
        [
            "publish-hf",
            "uwunion/celeba-attrs",
            "--format",
            "parquet",
            "--path",
            str(parquet_path),
            "--split",
            "train",
            "--private",
        ],
        input="y\n",
    )

    assert result.exit_code == 0, result.output
    # Ensure we asked about asset-like columns and set image_column.
    assert "asset-like columns" in result.output
    assert "image_path" in result.output
    assert called.get("image_column") == "image_path"


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

    def fake_load_dataset_builder(dataset_id, **kwargs):
        calls.setdefault("builders", []).append((dataset_id, kwargs))

    def fake_generate_from_cache(cache_dir, output, dry_run=False, only=None):
        calls["cache_dir"] = cache_dir
        calls["output"] = output
        calls["dry_run"] = dry_run
        calls["only"] = list(only or [])
        # Return a dummy path list to simulate generated schemas.
        return [output / "sapientinc__maze_30x30_hard_1k.py"]

    monkeypatch.setattr(cli, "generate_from_cache", fake_generate_from_cache)
    monkeypatch.setattr(cli, "load_dataset_builder", fake_load_dataset_builder)

    runner = CliRunner()
    result = runner.invoke(
        cli.app,
        [
            "download",
            "sapientinc/maze-30x30-hard-1k",
            "sapientinc/sudoku-extreme",
        ],
    )

    assert result.exit_code == 0, result.output
    # Ensure builders were called for each dataset id.
    assert calls.get("builders") == [
        ("sapientinc/maze-30x30-hard-1k", {}),
        ("sapientinc/sudoku-extreme", {}),
    ]
    # Ensure generate_from_cache was called with the expected filters.
    only = calls.get("only")
    assert "sapientinc/maze-30x30-hard-1k" in only
    assert "sapientinc/sudoku-extreme" in only


def test_cli_download_warns_on_failure(monkeypatch):
    def fake_load_dataset_builder(dataset_id, **kwargs):
        raise RuntimeError("boom")

    monkeypatch.setattr(cli, "load_dataset_builder", fake_load_dataset_builder)

    runner = CliRunner()
    result = runner.invoke(
        cli.app,
        [
            "download",
            "sapientinc/maze-30x30-hard-1k",
        ],
    )

    assert result.exit_code != 0
    assert "Unable to download metadata for 'sapientinc/maze-30x30-hard-1k'" in result.output


def test_cli_sync_space_uses_hfapi_and_builder(monkeypatch, tmp_path):
    class DummyEntry:
        def __init__(self, id: str) -> None:
            self.id = id

    calls: dict[str, object] = {}

    class DummyApi:
        def list_datasets(self, author: str):
            calls["author"] = author
            return [DummyEntry("uwunion/arcagi"), DummyEntry("uwunion/arcagi2")]

    def fake_whoami():
        return {"name": "uwunion"}

    def fake_builder(dataset_id, **kwargs):
        calls.setdefault("builders", []).append((dataset_id, kwargs))

    def fake_generate_from_cache(cache_dir, output, dry_run=False, only=None, skip_existing=False):
        calls["cache_dir"] = cache_dir
        calls["output"] = output
        calls["dry_run"] = dry_run
        calls["only"] = list(only or [])
        calls["skip_existing"] = skip_existing
        return [output / "uwunion__arcagi.py", output / "uwunion__arcagi2.py"]

    monkeypatch.setattr(cli, "HfApi", DummyApi)
    monkeypatch.setattr(cli, "whoami", fake_whoami)
    monkeypatch.setattr(cli, "load_dataset_builder", fake_builder)
    monkeypatch.setattr(cli, "generate_from_cache", fake_generate_from_cache)

    runner = CliRunner()
    result = runner.invoke(
        cli.app,
        [
            "sync-space",
        ],
    )

    assert result.exit_code == 0, result.output
    assert calls["author"] == "uwunion"
    assert calls["builders"] == [
        ("uwunion/arcagi", {}),
        ("uwunion/arcagi2", {}),
    ]
    only = calls["only"]
    assert "uwunion/arcagi" in only and "uwunion/arcagi2" in only
    assert calls["skip_existing"] is True


def test_cli_embeddings_add_and_list_spaces(tmp_path):
    store_path = tmp_path / "embeddings.json"
    runner = CliRunner()

    # Add a new embedding space.
    result = runner.invoke(
        cli.app,
        [
            "embeddings",
            "add",
            "--space",
            "text_gemma",
            "--model",
            "google/gemma-2-2b-it",
            "--store",
            str(store_path),
        ],
    )
    assert result.exit_code == 0, result.output
    assert store_path.exists()

    data = json.loads(store_path.read_text())
    assert "spaces" in data
    assert data["spaces"]["text_gemma"]["model"] == "google/gemma-2-2b-it"

    # List spaces and ensure our new one appears.
    result = runner.invoke(
        cli.app,
        [
            "embeddings",
            "list",
            "--store",
            str(store_path),
        ],
    )
    assert result.exit_code == 0, result.output
    assert "text_gemma" in result.stdout
    assert "google/gemma-2-2b-it" in result.stdout


def test_cli_embeddings_link_field(tmp_path):
    store_path = tmp_path / "embeddings.json"
    runner = CliRunner()

    # First define an embedding space to link to.
    result = runner.invoke(
        cli.app,
        [
            "embeddings",
            "add",
            "--space",
            "text_gemma",
            "--model",
            "google/gemma-2-2b-it",
            "--store",
            str(store_path),
        ],
    )
    assert result.exit_code == 0, result.output

    # Link a schema field to that space.
    result = runner.invoke(
        cli.app,
        [
            "embeddings",
            "link-field",
            "ag_news",
            "--field",
            "text",
            "--space",
            "text_gemma",
            "--store",
            str(store_path),
        ],
    )
    assert result.exit_code == 0, result.output

    data = json.loads(store_path.read_text())
    assert "links" in data
    assert data["links"]["ag_news"]["text"] == "text_gemma"

    # Config-specific link should use dataset:config as key.
    result = runner.invoke(
        cli.app,
        [
            "embeddings",
            "link-field",
            "imdb",
            "--config",
            "plain_text",
            "--field",
            "text",
            "--space",
            "text_gemma",
            "--store",
            str(store_path),
        ],
    )
    assert result.exit_code == 0, result.output

    data = json.loads(store_path.read_text())
    assert data["links"]["imdb:plain_text"]["text"] == "text_gemma"
