"""Dataset format knowledge base.

Specifications, compatibility matrix, and conversion paths for every
supported robotics dataset format.
"""

from __future__ import annotations

FORMAT_KNOWLEDGE: dict[str, dict] = {
    "lerobot-v3": {
        "name": "LeRobot v3 (current)",
        "description": "Latest LeRobot format using Parquet + video files",
        "structure": "meta/info.json, meta/episodes.jsonl, data/*.parquet, videos/",
        "key_field": "features (in info.json)",
        "compatible_with": ["lerobot"],
        "incompatible_with": ["groot"],
        "convert_to": {
            "lerobot-v2": "orbit convert {path} --to lerobot-v2 -o {output}",
            "hdf5": "orbit convert {path} --to hdf5 -o {output}",
        },
    },
    "lerobot-v2": {
        "name": "LeRobot v2 (legacy)",
        "description": "Older LeRobot format, required by GR00T N1.6",
        "structure": "meta/info.json, data/, videos/",
        "key_field": "video_keys (in info.json)",
        "compatible_with": ["lerobot", "groot"],
        "incompatible_with": [],
        "convert_to": {
            "lerobot-v3": "orbit convert {path} --to lerobot-v3 -o {output}",
            "hdf5": "orbit convert {path} --to hdf5 -o {output}",
        },
    },
    "rlds": {
        "name": "RLDS (TFRecord)",
        "description": "TensorFlow Datasets format used by Google Brain / RT-X",
        "structure": "*.tfrecord files with tf.train.Example schema",
        "compatible_with": ["openvla"],
        "incompatible_with": ["lerobot", "groot"],
        "convert_to": {
            "lerobot-v3": "orbit convert {path} --to lerobot-v3 -o {output}",
            "hdf5": "orbit convert {path} --to hdf5 -o {output}",
        },
    },
    "hdf5": {
        "name": "HDF5",
        "description": "Hierarchical Data Format — common in ACT/ALOHA pipelines",
        "structure": "*.hdf5 with /action, /observations groups per episode",
        "compatible_with": [],
        "incompatible_with": [],
        "convert_to": {
            "lerobot-v3": "orbit convert {path} --to lerobot-v3 -o {output}",
        },
    },
    "rosbag": {
        "name": "ROS Bag (MCAP/bag)",
        "description": "ROS 1 (.bag) or ROS 2 (.mcap) recorded topics",
        "structure": ".bag or .mcap files with timestamped topic messages",
        "compatible_with": [],
        "incompatible_with": [],
        "convert_to": {
            "lerobot-v3": (
                "orbit convert {path} --to lerobot-v3 -o {output} "
                "--action-topic /joint_commands --state-topic /joint_states"
            ),
        },
        "notes": "Requires specifying --action-topic and --state-topic for conversion",
    },
    "numpy": {
        "name": "Raw numpy/CSV",
        "description": "Unstructured numpy arrays or CSV files in a directory",
        "structure": "Directory with .npy/.npz/.csv files",
        "compatible_with": [],
        "incompatible_with": [],
        "convert_to": {
            "lerobot-v3": "orbit convert {path} --to lerobot-v3 -o {output}",
        },
    },
}

# Conversion compatibility matrix: source -> target -> command
CONVERSION_MATRIX: dict[str, dict[str, str]] = {}
for fmt_key, fmt_info in FORMAT_KNOWLEDGE.items():
    CONVERSION_MATRIX[fmt_key] = fmt_info.get("convert_to", {})


def get_format_summary() -> str:
    """Generate a concise text summary of all formats for the system prompt."""
    parts: list[str] = []
    for key, fmt in FORMAT_KNOWLEDGE.items():
        compat = ", ".join(fmt.get("compatible_with", [])) or "none directly"
        parts.append(f"- **{fmt['name']}**: {fmt['description']}. Compatible with: {compat}")
    return "\n".join(parts)


def get_conversion_advice(source_format: str, target_framework: str) -> str | None:
    """Get specific conversion advice for a source format → target framework."""
    from orbit.assist.knowledge.frameworks import FRAMEWORK_KNOWLEDGE

    fw = FRAMEWORK_KNOWLEDGE.get(target_framework)
    if not fw:
        return None

    needed = fw["format_required"].lower()

    # Check if already compatible
    fmt_info = FORMAT_KNOWLEDGE.get(source_format, {})
    if target_framework in fmt_info.get("compatible_with", []):
        return None  # No conversion needed

    # Find conversion path
    conversions = fmt_info.get("convert_to", {})
    for target_fmt, cmd in conversions.items():
        if target_fmt.lower() in needed or needed in target_fmt.lower():
            return f"Convert {source_format} → {target_fmt}: {cmd}"

    return f"No direct conversion path from {source_format} to {needed}. Try converting to lerobot-v3 first."
