from pathlib import Path

from schemas_hf.custom import infer_features, write_schema_module


def test_infer_features_scalars_and_lists():
    examples = [
        {
            "id": 1,
            "score": 3.5,
            "text": "hello",
            "flag": True,
            "ints": [1, 2, 3],
            "grid": [[1, 2], [3, 4]],
        },
        {
            "id": 2,
            "score": 4.0,
            "text": "world",
            "flag": False,
            "ints": [4, 5],
            "grid": [[5, 6]],
        },
    ]

    features = infer_features(examples)

    assert features["id"]["_type"] == "Value"
    assert features["id"]["dtype"].startswith("int")

    assert features["score"]["_type"] == "Value"
    assert features["score"]["dtype"].startswith("float")

    assert features["text"]["_type"] == "Value"
    assert features["text"]["dtype"] == "string"

    assert features["flag"]["_type"] == "Value"
    assert features["flag"]["dtype"] == "bool"

    ints = features["ints"]
    assert ints["_type"] == "Sequence"
    assert ints["feature"]["_type"] == "Value"
    assert ints["feature"]["dtype"].startswith("int")

    grid = features["grid"]
    assert grid["_type"] == "Sequence"
    inner = grid["feature"]
    assert inner["_type"] == "Sequence"
    assert inner["feature"]["_type"] == "Value"
    assert inner["feature"]["dtype"].startswith("int")


def test_infer_features_mixed_types_fallback_to_any():
    examples = [
        {"value": 1},
        {"value": "two"},
        {"value": 3},
    ]

    features = infer_features(examples)
    assert features["value"]["_type"] == "Value"
    # Unknown / mixed types should fall back to an unknown dtype
    assert features["value"]["dtype"] == "unknown"


def test_write_schema_module_creates_importable_dataclass(tmp_path: Path):
    features = {
        "id": {"_type": "Value", "dtype": "int64"},
        "text": {"_type": "Value", "dtype": "string"},
        "tags": {
            "_type": "Sequence",
            "feature": {"_type": "Value", "dtype": "string"},
        },
    }

    module_path = write_schema_module(
        dataset_id="local/demo_infer",
        config_name=None,
        features=features,
        output_dir=tmp_path,
    )

    assert module_path.exists()
    # Import the generated module and verify the dataclass shape
    import importlib.util

    spec = importlib.util.spec_from_file_location("local_demo_schema", 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]

    # Expect a Record class named after the dataset id
    from schemas_hf.base import DatasetSchema

    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="x", tags=["a", "b"])
    assert instance.id == 1
    assert instance.text == "x"
    assert instance.tags == ["a", "b"]
