import json
from pathlib import Path
from typing import Dict, List

import pandas as pd

from any4robot import merge_lerobot_datasets_from_folder


def _write_json(path: Path, payload: Dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload), encoding="utf-8")


def _write_jsonl(path: Path, rows: List[Dict]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(row) + "\n")


def _create_minimal_lerobot_dataset(root: Path, dataset_name: str, task_name: str) -> Path:
    dataset_root = root / dataset_name
    meta_root = dataset_root / "meta"
    data_chunk_root = dataset_root / "data" / "chunk-000"
    data_chunk_root.mkdir(parents=True, exist_ok=True)

    info = {
        "codebase_version": "0.1.0",
        "data_path": "data",
        "video_path": "videos",
        "features": {},
        "total_episodes": 1,
        "total_tasks": 1,
        "robot_type": "test_bot",
        "fps": 30,
    }
    _write_json(meta_root / "info.json", info)
    _write_jsonl(
        meta_root / "episodes.jsonl",
        [{"episode_index": 0, "task": task_name, "task_index": 0, "length": 2}],
    )
    _write_jsonl(meta_root / "tasks.jsonl", [{"task_index": 0, "task": task_name}])

    df = pd.DataFrame(
        {
            "frame_index": [0, 1],
            "episode_index": [0, 0],
            "task_index": [0, 0],
            "timestamp": [0.0, 0.033],
            "action": [[0.1], [0.2]],
            "observation.state": [[1.0], [1.1]],
        }
    )
    df.to_parquet(data_chunk_root / "episode_000000.parquet", index=False)
    return dataset_root


def test_merge_lerobot_datasets_from_folder_sdk(tmp_path: Path) -> None:
    source_root = tmp_path / "multi_dataset_folder"
    _create_minimal_lerobot_dataset(source_root, "dataset_a", "pick")
    _create_minimal_lerobot_dataset(source_root, "dataset_b", "place")

    output_path = tmp_path / "merged_dataset"
    ok = merge_lerobot_datasets_from_folder(
        source_root=source_root,
        output_path=output_path,
    )
    assert ok is True

    merged_episodes = (output_path / "meta" / "episodes.jsonl").read_text(encoding="utf-8").strip().splitlines()
    merged_tasks = (output_path / "meta" / "tasks.jsonl").read_text(encoding="utf-8").strip().splitlines()
    merged_info = json.loads((output_path / "meta" / "info.json").read_text(encoding="utf-8"))

    assert len(merged_episodes) == 2
    assert len(merged_tasks) == 2
    assert merged_info["total_episodes"] == 2
    assert (output_path / "data" / "chunk-000" / "episode_000000.parquet").exists()
    assert (output_path / "data" / "chunk-000" / "episode_000001.parquet").exists()
