"""Format conversion engine: any format → LeRobot v3."""

from __future__ import annotations

import json
import logging
from pathlib import Path

import numpy as np
import pandas as pd

from orbit.analyzer.dataset_loader import DatasetInfo

logger = logging.getLogger(__name__)

# Supported output formats
SUPPORTED_OUTPUT_FORMATS = ["lerobot-v3"]


def convert(
    source: str,
    output_dir: str,
    output_format: str = "lerobot-v3",
    source_format: str | None = None,
    fps: float | None = None,
    **loader_kwargs,
) -> Path:
    """Convert a dataset from any supported format to the target format.

    Args:
        source: Input path, file, or HuggingFace repo ID.
        output_dir: Output directory path.
        output_format: Target format (currently only "lerobot-v3").
        source_format: Explicit source format override.
        fps: Override FPS for output dataset.
        **loader_kwargs: Extra kwargs for source loader (e.g., action_topic).

    Returns:
        Path to the output directory.
    """
    if output_format not in SUPPORTED_OUTPUT_FORMATS:
        raise ValueError(
            f"Unsupported output format '{output_format}'. "
            f"Supported: {', '.join(SUPPORTED_OUTPUT_FORMATS)}"
        )

    from orbit.loaders import detect_and_load_episodes, detect_and_load_info

    # Load source data
    logger.info("Loading metadata from %s...", source)
    info = detect_and_load_info(source, format=source_format, **loader_kwargs)

    logger.info("Loading episode data from %s...", source)
    episodes_df = detect_and_load_episodes(
        source, max_episodes=None, format=source_format, **loader_kwargs
    )

    if episodes_df.empty:
        raise ValueError(f"No episode data loaded from {source}")

    # Override FPS if specified
    if fps is not None:
        info = DatasetInfo(
            repo_id=info.repo_id,
            num_episodes=info.num_episodes,
            num_frames=info.num_frames,
            fps=fps,
            robot_type=info.robot_type,
            camera_names=info.camera_names,
            video_keys=info.video_keys,
            shapes=info.shapes,
            duration_seconds=info.num_frames / fps if fps > 0 else 0.0,
        )

    out = Path(output_dir).expanduser().resolve()
    write_lerobot_v3(info, episodes_df, out)

    return out


def write_lerobot_v3(
    info: DatasetInfo,
    episodes_df: pd.DataFrame,
    output_dir: Path,
) -> None:
    """Write a LeRobot v3 dataset to local disk.

    Creates the standard LeRobot v3 directory structure:
        output_dir/
        ├── meta/
        │   ├── info.json
        │   └── episodes/
        │       └── chunk-000/
        │           └── file-000.parquet
        └── data/
            └── chunk-000/
                ├── episode_000000.parquet
                ├── episode_000001.parquet
                └── ...
    """
    output_dir.mkdir(parents=True, exist_ok=True)

    # Determine shapes from data
    shapes = dict(info.shapes) if info.shapes else {}
    if not shapes and not episodes_df.empty:
        # Infer from first non-null action/state
        for _, row in episodes_df.iterrows():
            if "action" in episodes_df.columns and row.get("action") is not None:
                action = row["action"]
                if isinstance(action, (list, np.ndarray)):
                    shapes["action"] = [len(action)]
                break
        for _, row in episodes_df.iterrows():
            if "state" in episodes_df.columns and row.get("state") is not None:
                state = row["state"]
                if isinstance(state, (list, np.ndarray)):
                    shapes["state"] = [len(state)]
                break

    # Build features dict for info.json
    features = {}
    if "action" in shapes:
        features["action"] = {
            "dtype": "float32",
            "shape": shapes["action"],
            "names": None,
        }
    if "state" in shapes:
        features["observation.state"] = {
            "dtype": "float32",
            "shape": shapes["state"],
            "names": None,
        }

    # Compute episode stats
    episode_groups = episodes_df.groupby("episode_index")
    unique_episodes = sorted(episodes_df["episode_index"].unique())
    num_episodes = len(unique_episodes)
    num_frames = len(episodes_df)

    # 1. Write meta/info.json
    meta_dir = output_dir / "meta"
    meta_dir.mkdir(parents=True, exist_ok=True)

    info_dict = {
        "codebase_version": "v3.0",
        "robot_type": info.robot_type or "unknown",
        "total_episodes": num_episodes,
        "total_frames": num_frames,
        "fps": info.fps or 30.0,
        "features": features,
        "splits": {"train": f"0:{num_episodes}"},
        "data_path": "data/chunk-{chunk:03d}/episode_{episode:06d}.parquet",
    }

    with open(meta_dir / "info.json", "w") as f:
        json.dump(info_dict, f, indent=2)

    # 2. Write meta/episodes parquet
    episodes_meta_dir = meta_dir / "episodes" / "chunk-000"
    episodes_meta_dir.mkdir(parents=True, exist_ok=True)

    episode_rows = []
    for new_idx, ep_idx in enumerate(unique_episodes):
        ep_data = episode_groups.get_group(ep_idx)
        episode_rows.append(
            {
                "episode_index": new_idx,
                "length": len(ep_data),
            }
        )

    ep_meta_df = pd.DataFrame(episode_rows)
    ep_meta_df.to_parquet(episodes_meta_dir / "file-000.parquet", index=False)

    # 3. Write per-episode data parquet files
    data_dir = output_dir / "data" / "chunk-000"
    data_dir.mkdir(parents=True, exist_ok=True)

    for new_idx, ep_idx in enumerate(unique_episodes):
        ep_data = episode_groups.get_group(ep_idx).copy()

        # Remap episode index to sequential
        ep_data["episode_index"] = new_idx

        # Ensure frame_index is sequential within episode
        ep_data["frame_index"] = range(len(ep_data))

        # Ensure timestamp exists
        if "timestamp" not in ep_data.columns or ep_data["timestamp"].isna().all():
            fps_val = info.fps or 30.0
            ep_data["timestamp"] = [i / fps_val for i in range(len(ep_data))]

        # Rename state -> observation.state for LeRobot compatibility
        if "state" in ep_data.columns:
            ep_data = ep_data.rename(columns={"state": "observation.state"})

        # Select output columns
        out_cols = ["episode_index", "frame_index", "timestamp"]
        if "action" in ep_data.columns:
            out_cols.append("action")
        if "observation.state" in ep_data.columns:
            out_cols.append("observation.state")

        # Add any extra columns
        for col in ep_data.columns:
            if col not in out_cols:
                out_cols.append(col)

        ep_data[out_cols].to_parquet(
            data_dir / f"episode_{new_idx:06d}.parquet", index=False
        )

    logger.info(
        "Wrote LeRobot v3 dataset: %d episodes, %d frames → %s",
        num_episodes,
        num_frames,
        output_dir,
    )
