"""Public /v1 released datasets + artifacts."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from ...models.v1.datasets import DatasetDetail, DatasetSummary, WebDatasetShardArtifact
from ...models.v1.pagination import Page, Pagination
from ._helpers import parse_detail

if TYPE_CHECKING:
    from ..._http import HttpTransport


class DatasetsResource:
    """client.datasets — released datasets and training artifacts."""

    def __init__(self, http: HttpTransport) -> None:
        self._http = http

    def list(self, *, page_size: int = 50, cursor: str | None = None) -> Page[DatasetSummary]:
        params: dict[str, Any] = {"page_size": page_size}
        if cursor is not None:
            params["cursor"] = cursor
        data = self._http.get("/datasets", params=params)
        items = [DatasetSummary.model_validate(i) for i in data.get("data", [])]
        pagination = Pagination.model_validate(data.get("pagination", {}) or {})
        return Page[DatasetSummary](
            data=items, pagination=pagination, request_id=data.get("request_id")
        )

    def iter(self, *, page_size: int = 50):
        from ...models.v1.pagination import auto_paginate

        return auto_paginate(self.list, page_size=page_size)

    def get(self, dataset_id_or_name: str) -> DatasetDetail:
        data = self._http.get(f"/datasets/{dataset_id_or_name}")
        return parse_detail(data, DatasetDetail)

    def list_webdataset_artifacts(
        self,
        dataset_id_or_name: str,
        *,
        variant: str | None = None,
        revision: str = "main",
        page_size: int = 200,
        cursor: str | None = None,
    ) -> Page[WebDatasetShardArtifact]:
        params: dict[str, Any] = {"page_size": page_size, "revision": revision}
        if cursor is not None:
            params["cursor"] = cursor
        if variant is not None:
            params["variant"] = variant

        data = self._http.get(
            f"/datasets/{dataset_id_or_name}/artifacts/webdataset",
            params=params,
        )
        items = [WebDatasetShardArtifact.model_validate(i) for i in data.get("data", [])]
        pagination = Pagination.model_validate(data.get("pagination", {}) or {})
        return Page[WebDatasetShardArtifact](
            data=items, pagination=pagination, request_id=data.get("request_id")
        )

    def webdataset_urls(
        self,
        dataset_id_or_name: str,
        *,
        variant: str | None = None,
        revision: str = "main",
        page_size: int = 500,
    ) -> list[str]:
        """Return all shard download URLs for a dataset (auto-paginates)."""
        from ...models.v1.pagination import auto_paginate

        return [
            shard.download_url
            for shard in auto_paginate(
                self.list_webdataset_artifacts,
                dataset_id_or_name=dataset_id_or_name,
                variant=variant,
                revision=revision,
                page_size=page_size,
            )
            if shard.download_url
        ]
