"""Public /v1 embeddings (space-agnostic)."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from ...models.v1.embeddings import (
    EmbeddingHit,
    EmbeddingsBatchItem,
    EmbeddingVectorItem,
)
from ...models.v1.pagination import Page, Pagination

if TYPE_CHECKING:
    from ..._http import HttpTransport


class EmbeddingsResource:
    """client.embeddings — query/similar + fetch vectors."""

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

    def query(
        self,
        *,
        query_text: str | None = None,
        query_image_url: str | None = None,
        query_image_base64: str | None = None,
        representation_profile: str = "general",
        page_size: int = 50,
        cursor: str | None = None,
        corpus: dict[str, Any] | None = None,
    ) -> Page[EmbeddingHit]:
        payload: dict[str, Any] = {
            "representation_profile": representation_profile,
            "page_size": page_size,
        }
        if cursor is not None:
            payload["cursor"] = cursor
        if query_text is not None:
            payload["query_text"] = query_text
        if query_image_url is not None:
            payload["query_image_url"] = query_image_url
        if query_image_base64 is not None:
            payload["query_image_base64"] = query_image_base64
        # Default to internal to avoid requiring dataset IDs for external.
        payload["corpus"] = corpus or {"source": "internal"}

        data = self._http.post("/embeddings/query", json=payload)
        items = [EmbeddingHit.model_validate(i) for i in data.get("data", [])]
        pagination = Pagination.model_validate(data.get("pagination", {}) or {})
        return Page[EmbeddingHit](
            data=items, pagination=pagination, request_id=data.get("request_id")
        )

    def similar(
        self,
        *,
        entity_id: str,
        entity_type: str = "clip",
        representation_profile: str = "general",
        page_size: int = 50,
        cursor: str | None = None,
        exclude_self: bool = True,
        corpus: dict[str, Any] | None = None,
    ) -> Page[EmbeddingHit]:
        payload: dict[str, Any] = {
            "entity_id": entity_id,
            "entity_type": entity_type,
            "representation_profile": representation_profile,
            "page_size": page_size,
            "exclude_self": exclude_self,
        }
        if cursor is not None:
            payload["cursor"] = cursor
        payload["corpus"] = corpus or {"source": "internal"}

        data = self._http.post("/embeddings/similar", json=payload)
        items = [EmbeddingHit.model_validate(i) for i in data.get("data", [])]
        pagination = Pagination.model_validate(data.get("pagination", {}) or {})
        return Page[EmbeddingHit](
            data=items, pagination=pagination, request_id=data.get("request_id")
        )

    def get_clip_embedding(self, clip_id: str) -> EmbeddingVectorItem:
        data = self._http.get(f"/clips/{clip_id}/embedding")
        return EmbeddingVectorItem.model_validate(data.get("data") or {})

    def batch(
        self,
        *,
        entity_ids: list[str],
        entity_type: str = "clip",
        representation_profile: str = "general",
        include_vector: bool = False,
    ) -> list[EmbeddingsBatchItem]:
        payload: dict[str, Any] = {
            "entity_ids": entity_ids,
            "entity_type": entity_type,
            "representation_profile": representation_profile,
            "include_vector": include_vector,
        }
        data = self._http.post("/embeddings/batch", json=payload)
        return [EmbeddingsBatchItem.model_validate(i) for i in data.get("data", [])]
