"""Dataset management SDK.

Usage:
    from bandito.datasets import get_dataset, list_datasets, create_dataset

    # List datasets for a project
    datasets = list_datasets("my-project")

    # Create a dataset from traces matching criteria
    dataset = create_dataset(
        name="golden-set",
        project="my-project",
        description="Known good traces",
    )

    # Get dataset with items
    ds = get_dataset(dataset_id=123)
    for item in ds.items:
        trace = ds.get_trace(item.trace_id)
"""

from __future__ import annotations

from typing import Any

from .cli.api import BanditoClient
from .cli.retrieval import RetrievalLayer


def _get_client() -> BanditoClient:
    """Get BanditoClient instance."""
    return BanditoClient()


def create_dataset(
    name: str,
    project: str,
    description: str | None = None,
) -> dict[str, Any]:
    """Create a new dataset.
    
    Args:
        name: Dataset name.
        project: Project name.
        description: Optional description.
    
    Returns:
        Dataset dict with id, name, project, description, created_at.
    """
    client = _get_client()
    return client.create_dataset(name=name, project=project, description=description)


def list_datasets(project: str) -> list[dict[str, Any]]:
    """List all datasets for a project.
    
    Args:
        project: Project name.
    
    Returns:
        List of dataset dicts.
    """
    client = _get_client()
    return client.list_datasets(project=project)


def get_dataset(dataset_id: int) -> Dataset:
    """Get a dataset with its items.
    
    Args:
        dataset_id: Dataset ID.
    
    Returns:
        Dataset object with items and trace data.
    """
    client = _get_client()
    data = client.get_dataset(dataset_id)
    return Dataset(client=client, **data)


def delete_dataset(dataset_id: int) -> None:
    """Delete a dataset.
    
    Args:
        dataset_id: Dataset ID.
    """
    client = _get_client()
    client.delete_dataset(dataset_id)


class Dataset:
    """A dataset with items and trace data."""

    def __init__(
        self,
        client: BanditoClient,
        id: int,
        name: str,
        project: str,
        description: str | None = None,
        created_by: str | None = None,
        created_at: str | None = None,
        updated_at: str | None = None,
        items: list[dict[str, Any]] | None = None,
    ):
        self.client = client
        self.id = id
        self.name = name
        self.project = project
        self.description = description
        self.created_by = created_by
        self.created_at = created_at
        self.updated_at = updated_at
        self._items = items or []
        self._retrieval = RetrievalLayer(project)

    @property
    def item_count(self) -> int:
        """Number of items in the dataset."""
        return len(self._items)

    @property
    def items(self) -> list[dict[str, Any]]:
        """Get dataset items (without full trace data)."""
        return self._items

    def cases(self) -> list[dict[str, Any]]:
        """Get dataset items with full trace data.
        
        Returns:
            List of dicts with trace_id, input, output, metadata.
        """
        results = []
        for item in self._items:
            trace = self._retrieval.get_trace(item["trace_id"])
            if trace:
                results.append({
                    "trace_id": item["trace_id"],
                    "input": trace.get("input"),
                    "output": trace.get("output"),
                    "metadata": item.get("metadata"),
                    "source": item.get("source"),
                })
        return results

    def add_item(
        self,
        trace_id: str,
        source: str = "manual",
        metadata: dict | None = None,
    ) -> dict[str, Any]:
        """Add a trace to the dataset.
        
        Args:
            trace_id: Trace ID to add.
            source: Source of the item (human, judge, assertion, manual).
            metadata: Optional metadata dict.
        
        Returns:
            DatasetItem dict.
        """
        item = self.client.add_dataset_item(
            dataset_id=self.id,
            trace_id=trace_id,
            source=source,
            metadata=metadata,
        )
        self._items.append(item)
        return item

    def remove_item(self, item_id: int) -> None:
        """Remove an item from the dataset.
        
        Args:
            item_id: Item ID to remove.
        """
        self.client.delete_dataset_item(self.id, item_id)
        self._items = [i for i in self._items if i["id"] != item_id]
