import json

import numpy as np
import pytest

from dagzoo.filtering.deferred_filter import _iter_packed_split_datasets, run_deferred_filter
from dagzoo.io.parquet_writer import write_packed_parquet_shards_stream
from dagzoo.types import DatasetBundle


def _bundle_with_embedded_config(
    seed: int,
    *,
    dataset_seed: int | None = None,
    dataset_index: int | None = None,
) -> DatasetBundle:
    metadata = {
        "seed": seed,
        "filter": {"mode": "deferred", "status": "not_run"},
        "config": {
            "dataset": {"task": "classification"},
            "filter": {"enabled": True},
        },
    }
    if dataset_seed is not None:
        metadata["dataset_seed"] = int(dataset_seed)
    if dataset_index is not None:
        metadata["dataset_index"] = int(dataset_index)

    return DatasetBundle(
        X_train=np.array(
            [
                [0.0, 1.0],
                [1.0, 0.0],
                [2.0, 1.0],
            ],
            dtype=np.float32,
        ),
        y_train=np.array([0, 1, 0], dtype=np.int64),
        X_test=np.array([[1.5, 0.5], [0.5, 1.5]], dtype=np.float32),
        y_test=np.array([1, 0], dtype=np.int64),
        feature_types=["num", "num"],
        metadata=metadata,
    )


def _bundle_without_config(seed: int) -> DatasetBundle:
    bundle = _bundle_with_embedded_config(seed)
    bundle.metadata["config"] = {"dataset": {"task": "classification"}}
    return bundle


def _load_ndjson(path) -> list[dict[str, object]]:
    payload: list[dict[str, object]] = []
    for line in path.read_text(encoding="utf-8").splitlines():
        if line.strip():
            payload.append(json.loads(line))
    return payload


def _write_ndjson_records(path, records: list[dict[str, object]]) -> None:
    path.write_text(
        "".join(json.dumps(record, sort_keys=True) + "\n" for record in records),
        encoding="utf-8",
    )


def _write_split_table(
    path,
    *,
    dataset_indices: list[int],
    row_indices: list[int],
    x_rows: list[list[float]],
    y_rows: list[int],
) -> None:
    pyarrow = pytest.importorskip("pyarrow")
    pyarrow_parquet = pytest.importorskip("pyarrow.parquet")

    table = pyarrow.table(
        {
            "dataset_index": pyarrow.array(dataset_indices, type=pyarrow.int64()),
            "row_index": pyarrow.array(row_indices, type=pyarrow.int64()),
            "x": pyarrow.array(x_rows, type=pyarrow.list_(pyarrow.float32())),
            "y": pyarrow.array(y_rows, type=pyarrow.int64()),
        }
    )
    pyarrow_parquet.write_table(table, path, compression="zstd")


def test_iter_packed_split_datasets_handles_dataset_split_across_record_batches(
    tmp_path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    pyarrow = pytest.importorskip("pyarrow")

    table = pyarrow.table(
        {
            "dataset_index": pyarrow.array([0, 0, 0, 1, 1], type=pyarrow.int64()),
            "row_index": pyarrow.array([0, 1, 2, 0, 1], type=pyarrow.int64()),
            "x": pyarrow.array(
                [[0.0, 0.5], [1.0, 1.5], [2.0, 2.5], [3.0, 3.5], [4.0, 4.5]],
                type=pyarrow.list_(pyarrow.float32()),
            ),
            "y": pyarrow.array([0, 1, 0, 1, 0], type=pyarrow.int64()),
        }
    )
    batches = [
        table.slice(0, 2).to_batches()[0],
        table.slice(2, 1).to_batches()[0],
        table.slice(3, 2).to_batches()[0],
    ]

    class _FakeParquetFile:
        def __init__(self, _path) -> None:
            self.schema_arrow = table.schema

        def iter_batches(self, *, columns):
            assert columns == ["dataset_index", "row_index", "x", "y"]
            return iter(batches)

    monkeypatch.setattr("dagzoo.filtering.deferred_filter.pq.ParquetFile", _FakeParquetFile)

    split_path = tmp_path / "train.parquet"
    datasets = list(_iter_packed_split_datasets(split_path))

    assert [dataset.dataset_index for dataset in datasets] == [0, 1]
    assert datasets[0].x.shape == (3, 2)
    assert datasets[0].y.tolist() == [0, 1, 0]
    assert np.allclose(datasets[0].x[:, 0], np.array([0.0, 1.0, 2.0], dtype=np.float32))
    assert datasets[1].x.shape == (2, 2)
    assert datasets[1].y.tolist() == [1, 0]


def test_iter_packed_split_datasets_handles_mixed_feature_widths_within_record_batch(
    tmp_path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    pyarrow = pytest.importorskip("pyarrow")

    table = pyarrow.table(
        {
            "dataset_index": pyarrow.array([0, 0, 1, 1], type=pyarrow.int64()),
            "row_index": pyarrow.array([0, 1, 0, 1], type=pyarrow.int64()),
            "x": pyarrow.array(
                [[0.0, 0.5], [1.0, 1.5], [2.0, 2.5, 2.75], [3.0, 3.5, 3.75]],
                type=pyarrow.list_(pyarrow.float32()),
            ),
            "y": pyarrow.array([0, 1, 1, 0], type=pyarrow.int64()),
        }
    )
    batches = table.to_batches(max_chunksize=4)

    class _FakeParquetFile:
        def __init__(self, _path) -> None:
            self.schema_arrow = table.schema

        def iter_batches(self, *, columns):
            assert columns == ["dataset_index", "row_index", "x", "y"]
            return iter(batches)

    monkeypatch.setattr("dagzoo.filtering.deferred_filter.pq.ParquetFile", _FakeParquetFile)

    split_path = tmp_path / "train.parquet"
    datasets = list(_iter_packed_split_datasets(split_path))

    assert [dataset.dataset_index for dataset in datasets] == [0, 1]
    assert datasets[0].x.shape == (2, 2)
    assert datasets[0].y.tolist() == [0, 1]
    assert datasets[1].x.shape == (2, 3)
    assert datasets[1].y.tolist() == [1, 0]


def test_run_deferred_filter_writes_manifest_and_summary(
    tmp_path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    pytest.importorskip("pyarrow.parquet")

    in_dir = tmp_path / "input"
    out_dir = tmp_path / "filter_out"
    bundles = [_bundle_with_embedded_config(101), _bundle_with_embedded_config(102)]
    _ = write_packed_parquet_shards_stream(bundles, in_dir, shard_size=2, compression="zstd")

    def _stub_filter(*_args, **_kwargs):
        seed = int(_kwargs["seed"])
        accepted = bool(seed % 2)
        details = {"wins_ratio": 1.0 if accepted else 0.0, "n_valid_oob": 128}
        if not accepted:
            details["reason"] = "below_threshold"
        return accepted, details

    monkeypatch.setattr(
        "dagzoo.filtering.deferred_filter._apply_extra_trees_filter_numpy", _stub_filter
    )

    result = run_deferred_filter(in_dir=in_dir, out_dir=out_dir)
    assert result.total_datasets == 2
    assert result.accepted_datasets == 1
    assert result.rejected_datasets == 1

    summary = json.loads(result.summary_path.read_text(encoding="utf-8"))
    assert summary["accepted_datasets"] == 1
    assert summary["rejected_datasets"] == 1
    assert summary["rejected_reason_counts"]["below_threshold"] == 1

    manifest_records = _load_ndjson(result.manifest_path)
    assert len(manifest_records) == 2
    statuses = {str(record["status"]) for record in manifest_records}
    assert statuses == {"accepted", "rejected"}
    for record in manifest_records:
        filter_payload = record["filter"]
        assert isinstance(filter_payload, dict)
        assert filter_payload["mode"] == "deferred"
        assert filter_payload["status"] in {"accepted", "rejected"}


def test_run_deferred_filter_writes_curated_output_for_accepted_only(
    tmp_path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    pyarrow_parquet = pytest.importorskip("pyarrow.parquet")

    in_dir = tmp_path / "input"
    out_dir = tmp_path / "filter_out"
    curated_out = tmp_path / "curated"
    bundles = [
        _bundle_with_embedded_config(201),
        _bundle_with_embedded_config(202),
        _bundle_with_embedded_config(203),
    ]
    _ = write_packed_parquet_shards_stream(bundles, in_dir, shard_size=3, compression="zstd")

    def _stub_filter(*_args, **_kwargs):
        seed = int(_kwargs["seed"])
        accepted = bool(seed % 2)
        return accepted, {"wins_ratio": 1.0 if accepted else 0.0, "n_valid_oob": 128}

    monkeypatch.setattr(
        "dagzoo.filtering.deferred_filter._apply_extra_trees_filter_numpy", _stub_filter
    )

    result = run_deferred_filter(in_dir=in_dir, out_dir=out_dir, curated_out_dir=curated_out)
    assert result.curated_accepted_datasets == 2

    shard_dir = curated_out / "shard_00000"
    assert shard_dir.exists()
    metadata_records = _load_ndjson(shard_dir / "metadata.ndjson")
    assert [int(record["dataset_index"]) for record in metadata_records] == [0, 2]

    train_table = pyarrow_parquet.read_table(shard_dir / "train.parquet")
    dataset_indices = {int(value) for value in train_table.column("dataset_index").to_pylist()}
    assert dataset_indices == {0, 2}


def test_run_deferred_filter_requires_embedded_filter_config(
    tmp_path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    pytest.importorskip("pyarrow.parquet")

    in_dir = tmp_path / "input"
    out_dir = tmp_path / "filter_out"
    bundles = [_bundle_without_config(7)]
    _ = write_packed_parquet_shards_stream(bundles, in_dir, shard_size=1, compression="zstd")

    monkeypatch.setattr(
        "dagzoo.filtering.deferred_filter._apply_extra_trees_filter_numpy",
        lambda *_args, **_kwargs: (True, {"wins_ratio": 1.0, "n_valid_oob": 128}),
    )

    with pytest.raises(ValueError, match="requires embedded metadata\\.config\\.filter"):
        _ = run_deferred_filter(in_dir=in_dir, out_dir=out_dir)


def test_run_deferred_filter_prefers_dataset_seed_when_present(
    tmp_path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    pytest.importorskip("pyarrow.parquet")

    in_dir = tmp_path / "input"
    out_dir = tmp_path / "filter_out"
    bundles = [
        _bundle_with_embedded_config(101, dataset_seed=501, dataset_index=0),
        _bundle_with_embedded_config(102, dataset_seed=502, dataset_index=1),
    ]
    _ = write_packed_parquet_shards_stream(bundles, in_dir, shard_size=2, compression="zstd")

    replay_seeds: list[int] = []

    def _stub_filter(*_args, **_kwargs):
        replay_seeds.append(int(_kwargs["seed"]))
        return True, {"wins_ratio": 1.0, "n_valid_oob": 128}

    monkeypatch.setattr(
        "dagzoo.filtering.deferred_filter._apply_extra_trees_filter_numpy", _stub_filter
    )

    result = run_deferred_filter(in_dir=in_dir, out_dir=out_dir)

    assert result.total_datasets == 2
    assert result.accepted_datasets == 2
    assert replay_seeds == [501, 502]


def test_run_deferred_filter_rejects_stale_filter_output_dir(
    tmp_path,
) -> None:
    pytest.importorskip("pyarrow.parquet")

    in_dir = tmp_path / "input"
    out_dir = tmp_path / "filter_out"
    _ = write_packed_parquet_shards_stream(
        [_bundle_with_embedded_config(301)],
        in_dir,
        shard_size=1,
        compression="zstd",
    )

    out_dir.mkdir()
    (out_dir / "filter_manifest.ndjson").write_text("", encoding="utf-8")

    with pytest.raises(RuntimeError, match="already contains prior artifacts"):
        _ = run_deferred_filter(in_dir=in_dir, out_dir=out_dir)


def test_run_deferred_filter_rejects_extra_split_rows_beyond_metadata(
    tmp_path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    pytest.importorskip("pyarrow.parquet")

    in_dir = tmp_path / "input"
    out_dir = tmp_path / "filter_out"
    _ = write_packed_parquet_shards_stream(
        [_bundle_with_embedded_config(401), _bundle_with_embedded_config(402)],
        in_dir,
        shard_size=2,
        compression="zstd",
    )

    metadata_path = in_dir / "shard_00000" / "metadata.ndjson"
    records = _load_ndjson(metadata_path)
    _write_ndjson_records(metadata_path, [records[0]])

    monkeypatch.setattr(
        "dagzoo.filtering.deferred_filter._apply_extra_trees_filter_numpy",
        lambda *_args, **_kwargs: (True, {"wins_ratio": 1.0, "n_valid_oob": 128}),
    )

    with pytest.raises(ValueError, match="extra dataset rows beyond metadata coverage"):
        _ = run_deferred_filter(in_dir=in_dir, out_dir=out_dir)
    assert list(out_dir.iterdir()) == []

    with pytest.raises(ValueError, match="extra dataset rows beyond metadata coverage"):
        _ = run_deferred_filter(in_dir=in_dir, out_dir=out_dir)
    assert list(out_dir.iterdir()) == []


def test_run_deferred_filter_rejects_non_monotonic_split_rows(
    tmp_path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    pytest.importorskip("pyarrow.parquet")

    shard_dir = tmp_path / "input" / "shard_00000"
    shard_dir.mkdir(parents=True)
    metadata_path = shard_dir / "metadata.ndjson"
    train_path = shard_dir / "train.parquet"
    test_path = shard_dir / "test.parquet"

    metadata_records = [
        {
            "dataset_index": 0,
            "n_train": 2,
            "n_test": 1,
            "n_features": 2,
            "feature_types": ["num", "num"],
            "metadata": {
                "seed": 11,
                "filter": {"mode": "deferred", "status": "not_run"},
                "config": {
                    "dataset": {"task": "classification"},
                    "filter": {"enabled": True},
                },
            },
        },
        {
            "dataset_index": 1,
            "n_train": 1,
            "n_test": 1,
            "n_features": 2,
            "feature_types": ["num", "num"],
            "metadata": {
                "seed": 12,
                "filter": {"mode": "deferred", "status": "not_run"},
                "config": {
                    "dataset": {"task": "classification"},
                    "filter": {"enabled": True},
                },
            },
        },
    ]
    _write_ndjson_records(metadata_path, metadata_records)
    _write_split_table(
        train_path,
        dataset_indices=[0, 1, 0],
        row_indices=[0, 0, 1],
        x_rows=[[0.0, 0.0], [1.0, 1.0], [0.5, 0.5]],
        y_rows=[0, 1, 0],
    )
    _write_split_table(
        test_path,
        dataset_indices=[0, 1],
        row_indices=[0, 0],
        x_rows=[[0.25, 0.25], [1.25, 1.25]],
        y_rows=[0, 1],
    )

    monkeypatch.setattr(
        "dagzoo.filtering.deferred_filter._apply_extra_trees_filter_numpy",
        lambda *_args, **_kwargs: (True, {"wins_ratio": 1.0, "n_valid_oob": 128}),
    )

    with pytest.raises(ValueError, match="monotonically increasing dataset_index"):
        _ = run_deferred_filter(in_dir=shard_dir, out_dir=tmp_path / "filter_out")


def test_run_deferred_filter_rejects_lineage_symlinks_during_curated_copy(
    tmp_path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    pytest.importorskip("pyarrow.parquet")

    in_dir = tmp_path / "input"
    out_dir = tmp_path / "filter_out"
    curated_out = tmp_path / "curated_out"
    _ = write_packed_parquet_shards_stream(
        [_bundle_with_embedded_config(501)],
        in_dir,
        shard_size=1,
        compression="zstd",
    )

    lineage_dir = in_dir / "shard_00000" / "lineage"
    lineage_dir.mkdir()
    outside_path = tmp_path / "outside.txt"
    outside_path.write_text("sentinel\n", encoding="utf-8")
    try:
        (lineage_dir / "escape.txt").symlink_to(outside_path)
    except OSError:
        pytest.skip("symlinks unavailable in this environment")

    monkeypatch.setattr(
        "dagzoo.filtering.deferred_filter._apply_extra_trees_filter_numpy",
        lambda *_args, **_kwargs: (True, {"wins_ratio": 1.0, "n_valid_oob": 128}),
    )

    with pytest.raises(RuntimeError, match="must not be a symlink"):
        _ = run_deferred_filter(in_dir=in_dir, out_dir=out_dir, curated_out_dir=curated_out)
    assert list(out_dir.iterdir()) == []
    assert list(curated_out.iterdir()) == []

    with pytest.raises(RuntimeError, match="must not be a symlink"):
        _ = run_deferred_filter(in_dir=in_dir, out_dir=out_dir, curated_out_dir=curated_out)
    assert list(out_dir.iterdir()) == []
    assert list(curated_out.iterdir()) == []


def test_run_deferred_filter_cleans_up_curated_output_after_split_exhaustion_failure(
    tmp_path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    pytest.importorskip("pyarrow.parquet")

    in_dir = tmp_path / "input"
    out_dir = tmp_path / "filter_out"
    curated_out = tmp_path / "curated_out"
    _ = write_packed_parquet_shards_stream(
        [_bundle_with_embedded_config(601), _bundle_with_embedded_config(602)],
        in_dir,
        shard_size=2,
        compression="zstd",
    )

    metadata_path = in_dir / "shard_00000" / "metadata.ndjson"
    records = _load_ndjson(metadata_path)
    _write_ndjson_records(metadata_path, [records[0]])

    monkeypatch.setattr(
        "dagzoo.filtering.deferred_filter._apply_extra_trees_filter_numpy",
        lambda *_args, **_kwargs: (True, {"wins_ratio": 1.0, "n_valid_oob": 128}),
    )

    with pytest.raises(ValueError, match="extra dataset rows beyond metadata coverage"):
        _ = run_deferred_filter(in_dir=in_dir, out_dir=out_dir, curated_out_dir=curated_out)
    assert list(out_dir.iterdir()) == []
    assert list(curated_out.iterdir()) == []

    with pytest.raises(ValueError, match="extra dataset rows beyond metadata coverage"):
        _ = run_deferred_filter(in_dir=in_dir, out_dir=out_dir, curated_out_dir=curated_out)
    assert list(out_dir.iterdir()) == []
    assert list(curated_out.iterdir()) == []
