"""orbit explore — browse and discover LeRobot datasets on HuggingFace."""

from __future__ import annotations

from dataclasses import dataclass

from rich.table import Table

from orbit.utils.display import console


@dataclass
class DatasetResult:
    """A dataset search result."""

    repo_id: str
    robot_type: str | None
    num_episodes: int | None
    task: str | None
    downloads: int
    likes: int
    grade: str | None = None  # cached ORBIT grade if available


def search_datasets(
    robot: str | None = None,
    task: str | None = None,
    min_grade: str | None = None,
    policy: str | None = None,
    min_episodes: int | None = None,
    limit: int = 20,
) -> list[DatasetResult]:
    """Search HuggingFace for LeRobot-compatible datasets."""
    from huggingface_hub import HfApi

    api = HfApi()

    # Search for datasets tagged with lerobot
    search_query = "lerobot"
    if robot:
        search_query += f" {robot}"
    if task:
        search_query += f" {task}"

    try:
        datasets = list(api.list_datasets(
            search=search_query,
            sort="downloads",
            limit=limit * 3,  # fetch extra to filter
        ))
    except Exception as e:
        console.print(f"[bold red]Error:[/bold red] Could not search HuggingFace: {e}")
        return []

    results: list[DatasetResult] = []
    for ds in datasets:
        tags = ds.tags or []

        # Filter: must have lerobot tag or be in lerobot namespace
        is_lerobot = any("lerobot" in t.lower() for t in tags) or (
            ds.id and ds.id.startswith("lerobot/")
        )
        if not is_lerobot:
            continue

        # Extract robot type from tags
        robot_type = None
        for tag in tags:
            if tag.startswith("robot_"):
                robot_type = tag.replace("robot_", "")
                break

        # Filter by robot type
        if robot and robot_type and robot.lower() not in robot_type.lower():
            if robot.lower() not in ds.id.lower():
                continue

        result = DatasetResult(
            repo_id=ds.id,
            robot_type=robot_type,
            num_episodes=None,  # would need card parsing
            task=None,
            downloads=ds.downloads or 0,
            likes=ds.likes or 0,
        )
        results.append(result)

        if len(results) >= limit:
            break

    return results


def print_explore_results(results: list[DatasetResult]) -> None:
    """Print search results as a Rich table."""
    if not results:
        console.print("[yellow]No matching datasets found.[/yellow]")
        return

    table = Table(title="LeRobot Datasets", border_style="cyan")
    table.add_column("#", style="dim", width=3)
    table.add_column("Dataset", style="bold")
    table.add_column("Robot", style="cyan")
    table.add_column("Downloads", justify="right")
    table.add_column("Likes", justify="right", style="yellow")
    table.add_column("Grade", justify="center")

    for i, r in enumerate(results, 1):
        grade_str = ""
        if r.grade:
            color = {"A": "green", "B": "green", "C": "yellow", "D": "yellow", "F": "red"}.get(r.grade, "dim")
            grade_str = f"[{color}]{r.grade}[/{color}]"

        table.add_row(
            str(i),
            r.repo_id,
            r.robot_type or "-",
            str(r.downloads),
            str(r.likes),
            grade_str or "-",
        )

    console.print()
    console.print(table)
    console.print(f"\n[dim]Showing {len(results)} results. Use orbit analyze <dataset> for details.[/dim]")
