"""Type stubs for the luci search engine."""

from typing import Any, TypedDict

# --- Mapping types ---

class FieldMapping(TypedDict, total=False):
    """A single field's mapping definition."""

    type: str
    """Field type: ``"text"``, ``"keyword"``, ``"integer"``, ``"long"``,
    ``"float"``, ``"double"``, ``"boolean"``, ``"date"``,
    ``"dense_vector"``, ``"geo_point"``, ``"nested"``."""
    analyzer: str
    """Analyzer name for text fields (default: ``"standard"``)."""
    dims: int
    """Vector dimensions for ``dense_vector`` fields."""
    store: bool
    """Whether to store the field value in ``_source``."""
    index: bool
    """Whether to index the field for search."""
    doc_values: bool
    """Whether to create columnar storage for sorting/aggregations."""
    norms: bool
    """Whether to store field length norms for BM25 scoring."""

class MappingProperties(TypedDict):
    """The ``properties`` section of a mapping."""

    properties: dict[str, FieldMapping]

class MappingDef(TypedDict, total=False):
    """ES-compatible field mapping definition."""

    properties: dict[str, FieldMapping]
    mappings: MappingProperties
    dynamic: str
    """Dynamic mapping mode: ``"true"`` (default) or ``"false"``."""

# --- Query types ---

class TermQuery(TypedDict):
    term: dict[str, str | dict[str, str]]

class TermsQuery(TypedDict):
    terms: dict[str, list[str]]

class MatchQuery(TypedDict):
    match: dict[str, str | dict[str, str]]

class MatchPhraseQuery(TypedDict):
    match_phrase: dict[str, str | dict[str, str]]

class ExistsQuery(TypedDict):
    exists: dict[str, str]

class PrefixQuery(TypedDict):
    prefix: dict[str, str | dict[str, str]]

class MatchAllQuery(TypedDict):
    match_all: dict[str, Any]

class MatchNoneQuery(TypedDict):
    match_none: dict[str, Any]

class BoolClauses(TypedDict, total=False):
    """Clauses for a ``bool`` query."""

    must: list[Query]
    should: list[Query]
    must_not: list[Query]
    filter: list[Query]

class BoolQuery(TypedDict):
    bool: BoolClauses

class ConstantScoreInner(TypedDict, total=False):
    filter: Query
    boost: float

class ConstantScoreQuery(TypedDict):
    constant_score: ConstantScoreInner

class GeoDistanceInner(TypedDict, total=False):
    distance: str
    """Distance string, e.g. ``"10km"``, ``"5mi"``, ``"1000m"``."""

class GeoDistanceQuery(TypedDict):
    geo_distance: dict[str, Any]

class GeoBoundingBoxField(TypedDict):
    top_left: dict[str, float] | list[float] | str
    bottom_right: dict[str, float] | list[float] | str

class GeoBoundingBoxQuery(TypedDict):
    geo_bounding_box: dict[str, GeoBoundingBoxField]

class NestedInner(TypedDict):
    path: str
    query: Query

class NestedQuery(TypedDict):
    nested: NestedInner

type Query = (
    TermQuery
    | TermsQuery
    | MatchQuery
    | MatchPhraseQuery
    | BoolQuery
    | ExistsQuery
    | PrefixQuery
    | ConstantScoreQuery
    | GeoDistanceQuery
    | GeoBoundingBoxQuery
    | NestedQuery
    | MatchAllQuery
    | MatchNoneQuery
    | dict[str, Any]
)
"""Any ES-compatible query dict."""

# --- KNN types ---

class KnnRequest(TypedDict, total=False):
    """kNN search configuration."""

    field: str
    """The ``dense_vector`` field to search."""
    query_vector: list[float]
    """The query vector."""
    k: int
    """Number of nearest neighbors to return (default: 10)."""
    num_candidates: int
    """Number of candidates to consider during HNSW traversal (default: 2*k)."""

# --- Aggregation types ---

class FieldAgg(TypedDict):
    """Aggregation on a single field."""

    field: str

class TermsAgg(TypedDict, total=False):
    """Terms bucket aggregation."""

    field: str
    size: int
    """Maximum number of buckets to return (default: 10)."""

class RangeEntry(TypedDict, total=False):
    key: str
    from_: float
    to: float

class RangeAgg(TypedDict):
    field: str
    ranges: list[dict[str, float | str | None]]

class HistogramAgg(TypedDict):
    field: str
    interval: float

type AggregationDef = dict[str, dict[str, Any]]
"""A named aggregation definition, e.g. ``{"avg": {"field": "price"}}``."""

type Aggregations = dict[str, AggregationDef | dict[str, Any]]
"""Named aggregations with optional ``aggs`` sub-aggregations."""

# --- Search request ---

class SearchRequest(TypedDict, total=False):
    """Full search request with query, kNN, aggregations, and size."""

    query: Query
    """The text/structured query."""
    knn: KnnRequest
    """Optional kNN vector search. Combined with ``query`` via RRF."""
    aggs: Aggregations
    """Optional aggregations to compute alongside the query."""
    aggregations: Aggregations
    """Alias for ``aggs``."""
    size: int
    """Maximum number of hits to return (default: 10)."""

# --- Search response ---

class SearchHit(TypedDict):
    """A single search result."""

    _score: float
    """Relevance score (BM25 for text, RRF for hybrid, 1.0 for filters)."""
    _doc_id: int
    """Internal document ID within the segment."""
    _source: dict[str, Any]
    """The original document."""

class MetricAggResult(TypedDict, total=False):
    """Result from a metric aggregation (avg, sum, min, max, etc.)."""

    value: float | None
    count: float
    min: float
    max: float
    sum: float

class BucketEntry(TypedDict, total=False):
    """A single bucket in a bucket aggregation result."""

    key: str | float
    """Bucket key (term value, histogram boundary, etc.)."""
    doc_count: int
    """Number of documents in this bucket."""

class BucketAggResult(TypedDict):
    """Result from a bucket aggregation (terms, range, histogram, etc.)."""

    buckets: list[BucketEntry]

type AggResult = dict[str, object]

class SearchResults(TypedDict, total=False):
    """Search response."""

    hits: list[SearchHit]
    """Matching documents, sorted by score descending."""
    total_hits: dict[str, Any]
    """Total hit count: ``{"value": N, "relation": "eq"|"gte"}``."""
    aggregations: dict[str, AggResult]
    """Aggregation results (present if ``aggs`` was specified)."""

# --- GeoPoint types ---

class GeoPointObject(TypedDict):
    """Geographic point as an object."""

    lat: float
    lon: float

type GeoPoint = GeoPointObject | list[float] | str
"""A geographic point: ``{"lat": 40.7, "lon": -73.9}`` or
``[-73.9, 40.7]`` (GeoJSON) or ``"40.7,-73.9"``."""

# --- Index class ---

class Index:
    """A Luci search index.

    The primary entry point for indexing and searching documents.
    Wraps the Rust ``luci_index::Index`` via PyO3.

    Example::

        import luci

        # Create with explicit mappings
        index = luci.Index.create("data.luci", {
            "properties": {
                "title": {"type": "text"},
                "tag": {"type": "keyword"},
                "price": {"type": "float"},
            }
        })

        # Or create with fully dynamic mappings
        index = luci.Index.create("data.luci")

        index.add({"title": "Hello World", "tag": "greeting", "price": 9.99})
        index.commit()

        results = index.search({"match": {"title": "hello"}})

        # Open existing index (mappings loaded from disk)
        index = luci.Index.open("data.luci")
    """

    @staticmethod
    def create(
        path: str,
        mapping: MappingDef | MappingProperties | dict[str, Any] | None = None,
        memory_budget: int | None = None,
    ) -> Index:
        """Create a new index.

        Args:
            path: File path for the index (e.g. ``"data.luci"``).
            mapping: Optional ES-compatible mapping definition. If omitted,
                all field types are inferred dynamically from documents.
            memory_budget: Memory budget in bytes before auto-flush to a new
                segment. Smaller values produce more segments (better query
                parallelism, more merge work). Default: 64 MB (67108864).
                Not persisted — applies only to this session's writes.

        Returns:
            A new ``Index`` instance.

        Raises:
            ValueError: If the mapping is invalid.

        """

    @staticmethod
    def open(path: str, memory_budget: int | None = None) -> Index:
        """Open an existing index.

        Mappings are loaded from the index file automatically.

        Args:
            path: File path of the index.
            memory_budget: Memory budget in bytes before auto-flush to a new
                segment. Default: 64 MB (67108864).
                Not persisted — applies only to this session's writes.

        Returns:
            An ``Index`` instance connected to the existing index.

        Raises:
            ValueError: If the path doesn't exist or the file is corrupted.

        """

    def add(self, doc: dict[str, Any]) -> None:
        """Add a document to the index.

        The document is buffered in memory until :meth:`commit` is called.
        Unknown fields are handled according to the mapping's dynamic mode.

        Args:
            doc: A dict representing the document.

        Raises:
            ValueError: If the document is not a valid JSON object.

        """

    def commit(self) -> None:
        """Commit all buffered documents, making them searchable.

        Persists field mappings alongside segment data. After commit returns,
        all previously added documents are visible to search.

        Raises:
            ValueError: If the commit fails (e.g., I/O error).

        """

    def force_merge(self, max_segments: int = 1) -> None:
        """Force-merge all segments down to at most ``max_segments``.

        Call after bulk indexing to optimize for query performance.
        This is expensive and should not be called routinely.

        Args:
            max_segments: Target number of segments. Default: 1.

        """

    def search(
        self,
        query: dict[str, Any],
        size: int = 10,
    ) -> SearchResults:
        """Search the index.

        Accepts either a bare query dict or a full search request with
        query, kNN, aggregations, and size.

        Args:
            query: ES-compatible query dict or full search request.
            size: Maximum number of hits to return. Default: 10.

        Returns:
            A :class:`SearchResults` dict with ``hits``, ``total_hits``,
            and optionally ``aggregations``.

        Raises:
            ValueError: If the query is invalid.

        """

    def buffered_doc_count(self) -> int:
        """Return the number of documents buffered but not yet committed."""

def set_num_threads(num_threads: int) -> None:
    """Configure the number of threads for parallel search.

    Must be called before any search operation. Can only be called once
    per process. If not called, defaults to the number of available CPUs.

    Args:
        num_threads: Number of threads to use for parallel search.

    Raises:
        ValueError: If called more than once or after a search has started.

    """

__version__: str

__all__: list[str] = ["Index", "set_num_threads"]
