from __future__ import annotations

import subprocess
import sys
from pathlib import Path

from dovla_cil.data.datasets import CILDataset
from dovla_cil.data.lerobot_export import LeRobotExportConfig, export_lerobot_style_dataset
from dovla_cil.generation.pipeline import generate_cil_dataset
from dovla_cil.tasks.library import built_in_toy_tasks
from dovla_cil.utils.io import iter_jsonl, read_json


def _make_toy_cil(tmp_path: Path) -> Path:
    dataset_dir = tmp_path / "cil"
    generate_cil_dataset(
        backend="toy",
        tasks=built_in_toy_tasks()[:2],
        out_dir=dataset_dir,
        num_states_per_task=2,
        k=4,
        seed=19,
        shard_size=8,
        inline_observations=True,
    )
    return dataset_dir


def test_lerobot_style_export_selects_best_record_per_group(tmp_path: Path) -> None:
    dataset_dir = _make_toy_cil(tmp_path)
    out_dir = tmp_path / "lerobot"

    metadata = export_lerobot_style_dataset(
        LeRobotExportConfig(
            dataset_dir=dataset_dir,
            out_dir=out_dir,
            max_groups=3,
            copy_images=False,
        )
    )

    rows = list(iter_jsonl(out_dir / "train.jsonl"))
    dataset = CILDataset(dataset_dir)

    assert metadata["schema_version"] == "dovla-cil-lerobot-export/v0"
    assert metadata["num_episodes"] == 3
    assert len(rows) == 3
    assert (out_dir / "tasks.jsonl").exists()
    assert read_json(out_dir / "metadata.json") == metadata
    for row in rows:
        group = dataset.get_group(row["cil"]["group_id"])
        assert row["reward"] == max(record.reward.score for record in group)
        assert row["cil"]["record_id"] in {record.record_id for record in group}
        assert row["task"]
        assert row["observation"]["image"] is None
        assert "action_chunk" in row


def test_lerobot_style_export_cli_runs_without_network(tmp_path: Path) -> None:
    dataset_dir = _make_toy_cil(tmp_path)
    out_dir = tmp_path / "cli-export"

    result = subprocess.run(
        [
            sys.executable,
            "scripts/export_lerobot_dataset.py",
            "--dataset",
            str(dataset_dir),
            "--out",
            str(out_dir),
            "--max-groups",
            "2",
            "--no-images",
        ],
        check=True,
        capture_output=True,
        text=True,
    )

    assert "dovla-cil-lerobot-export/v0" in result.stdout
    assert len(list(iter_jsonl(out_dir / "train.jsonl"))) == 2


def test_task_balanced_export_covers_tasks_deterministically(tmp_path: Path) -> None:
    dataset_dir = _make_toy_cil(tmp_path)
    first_out = tmp_path / "balanced-first"
    second_out = tmp_path / "balanced-second"
    config_kwargs = {
        "dataset_dir": dataset_dir,
        "max_groups": 4,
        "group_sampling": "task_balanced",
        "seed": 7,
        "copy_images": False,
    }

    export_lerobot_style_dataset(LeRobotExportConfig(out_dir=first_out, **config_kwargs))
    export_lerobot_style_dataset(LeRobotExportConfig(out_dir=second_out, **config_kwargs))
    first = list(iter_jsonl(first_out / "train.jsonl"))
    second = list(iter_jsonl(second_out / "train.jsonl"))

    assert {row["cil"]["task_id"] for row in first} == {
        "toy_pick_red_mug",
        "toy_put_red_mug_in_blue_bowl",
    }
    assert [row["cil"]["group_id"] for row in first] == [
        row["cil"]["group_id"] for row in second
    ]