from typing import Any, Literal, cast

from ..models import Dataset, DatasetCreating, PageDatasetExample
from ..ty import Page
from .base import BaseClient


class Data(BaseClient):
    def all(
        self,
        session: bool | None = None,
        *,
        page: int | None = None,
        size: int | None = None,
    ) -> Page[str]:
        res = self.request(
            "GET",
            endpoint="all",
            data={"session": session},
            page=page,
            size=size,
            return_model=Page[str],
        )
        return cast(Page[str], res)

    async def all_async(
        self,
        session: bool | None = None,
        *,
        page: int | None = None,
        size: int | None = None,
    ) -> Page[str]:
        res = await self.request_async(
            "GET",
            endpoint="all",
            data={"session": session},
            page=page,
            size=size,
            return_model=Page[str],
        )
        return cast(Page[str], res)

    def create(self, body: DatasetCreating) -> Dataset:
        res = self.request("POST", endpoint="create", data=body, return_model=Dataset)
        return cast(Dataset, res)

    async def create_async(self, body: DatasetCreating) -> Dataset:
        res = await self.request_async(
            "POST", endpoint="create", data=body, return_model=Dataset
        )
        return cast(Dataset, res)

    def read_examples(
        self,
        datasets: list[str],
        *,
        page: int | None = None,
        size: int | None = None,
        order: Literal["asc", "desc"] | None = "asc",
    ) -> PageDatasetExample:
        params: dict[str, Any] = {}
        if page is not None:
            params["page"] = page
        if size is not None:
            params["size"] = size
        if order is not None:
            params["order"] = order

        data = {"datasets": datasets}

        res = self.request(
            "POST",
            endpoint="read-examples",
            params=params,
            data=data,
            return_model=PageDatasetExample,
        )
        return cast(PageDatasetExample, res)

    async def read_examples_async(
        self,
        datasets: list[str],
        *,
        page: int | None = None,
        size: int | None = None,
        order: Literal["asc", "desc"] | None = "asc",
    ) -> PageDatasetExample:
        params: dict[str, Any] = {}
        if page is not None:
            params["page"] = page
        if size is not None:
            params["size"] = size
        if order is not None:
            params["order"] = order

        data = {"datasets": datasets}

        res = await self.request_async(
            "POST",
            endpoint="read-examples",
            params=params,
            data=data,
            return_model=PageDatasetExample,
        )
        return cast(PageDatasetExample, res)
