"""
Elasticsearch Pagination Helpers.

Provides search_after pagination with Point-in-Time (PIT) for
efficient traversal of large result sets.
"""

from __future__ import annotations

from collections.abc import AsyncIterator, Iterator
from typing import (
    TYPE_CHECKING,
    Any,
)

from smf.smf_logging import get_logger

if TYPE_CHECKING:
    from elasticsearch import AsyncElasticsearch, Elasticsearch

logger = get_logger("elasticsearch.pagination")


def search_after_iterator(
    client: Elasticsearch,
    index: str,
    query: dict[str, Any],
    sort: list[dict[str, Any]],
    *,
    page_size: int = 1000,
    max_results: int | None = None,
    pit_keep_alive: str = "5m",
    source: list[str] | None = None,
) -> Iterator[dict[str, Any]]:
    """Iterate through search results using search_after with Point-in-Time.

    Uses Point-in-Time (PIT) to maintain a consistent view of the data
    throughout pagination, avoiding issues with index changes during iteration.

    Args:
        client: Elasticsearch client.
        index: Index to search.
        query: Elasticsearch query dict (e.g., {"match_all": {}}).
        sort: Sort specification (required for search_after).
               Must include a tiebreaker like _shard_doc or _id.
        page_size: Number of documents per page (default: 1000).
        max_results: Maximum total results to return (None for unlimited).
        pit_keep_alive: Point-in-Time keep-alive duration (default: "5m").
        source: Fields to include in _source (None for all fields).

    Yields:
        Individual hit documents (dict with _id, _source, etc.).

    Example:
        >>> query = {"match": {"status": "active"}}
        >>> sort = [{"created_at": "desc"}, {"_id": "asc"}]
        >>> for hit in search_after_iterator(client, "my-index", query, sort):
        ...     print(hit["_source"]["title"])
    """
    # Open Point-in-Time for consistent pagination
    pit_response = client.open_point_in_time(
        index=index,
        keep_alive=pit_keep_alive,
    )

    # Handle ObjectApiResponse wrapper
    if hasattr(pit_response, "body"):
        pit_id = pit_response.body.get("id") or pit_response.body.get("pit_id")
    elif isinstance(pit_response, dict):
        pit_id = pit_response.get("id") or pit_response.get("pit_id")
    else:
        pit_id = str(pit_response)

    logger.debug(f"Opened PIT: {pit_id[:20]}...", index=index)

    try:
        search_after: list[Any] | None = None
        total_returned = 0

        while True:
            # Build search request
            body: dict[str, Any] = {
                "query": query,
                "sort": sort,
                "size": page_size,
                "pit": {"id": pit_id, "keep_alive": pit_keep_alive},
                "track_total_hits": False,  # Performance optimization
            }

            if search_after:
                body["search_after"] = search_after

            if source:
                body["_source"] = source

            # Execute search
            try:
                response = client.search(body=body)
            except TypeError:
                # Some ES versions don't accept body parameter directly
                response = client.search(**body)

            # Extract hits from response
            if hasattr(response, "body"):
                response_dict = response.body
            elif isinstance(response, dict):
                response_dict = response
            else:
                response_dict = dict(response)

            hits = response_dict.get("hits", {}).get("hits", [])

            if not hits:
                logger.debug(
                    "Pagination complete - no more results",
                    total_returned=total_returned,
                )
                break

            for hit in hits:
                yield hit
                total_returned += 1

                if max_results and total_returned >= max_results:
                    logger.debug(
                        f"Reached max_results limit ({max_results})",
                        total_returned=total_returned,
                    )
                    return

            # Get search_after value from last hit's sort
            search_after = hits[-1].get("sort")

            if not search_after:
                logger.warning("Last hit has no sort value, ending pagination")
                break

            logger.debug(
                f"Fetched page with {len(hits)} documents",
                total_returned=total_returned,
            )

    finally:
        # Clean up Point-in-Time
        try:
            client.close_point_in_time(body={"id": pit_id})
            logger.debug(f"Closed PIT: {pit_id[:20]}...")
        except Exception as exc:
            logger.warning(f"Failed to close PIT: {exc}")


async def search_after_iterator_async(
    client: AsyncElasticsearch,
    index: str,
    query: dict[str, Any],
    sort: list[dict[str, Any]],
    *,
    page_size: int = 1000,
    max_results: int | None = None,
    pit_keep_alive: str = "5m",
    source: list[str] | None = None,
) -> AsyncIterator[dict[str, Any]]:
    """Async iterate through search results using search_after with Point-in-Time.

    Uses Point-in-Time (PIT) to maintain a consistent view of the data
    throughout pagination, avoiding issues with index changes during iteration.

    Args:
        client: Async Elasticsearch client.
        index: Index to search.
        query: Elasticsearch query dict.
        sort: Sort specification (required for search_after).
        page_size: Number of documents per page (default: 1000).
        max_results: Maximum total results to return (None for unlimited).
        pit_keep_alive: Point-in-Time keep-alive duration (default: "5m").
        source: Fields to include in _source (None for all fields).

    Yields:
        Individual hit documents (dict with _id, _source, etc.).

    Example:
        >>> query = {"match": {"status": "active"}}
        >>> sort = [{"created_at": "desc"}, {"_id": "asc"}]
        >>> async for hit in search_after_iterator_async(client, "my-index", query, sort):
        ...     print(hit["_source"]["title"])
    """
    pit_response = await client.open_point_in_time(
        index=index,
        keep_alive=pit_keep_alive,
    )

    if hasattr(pit_response, "body"):
        pit_id = pit_response.body.get("id") or pit_response.body.get("pit_id")
    elif isinstance(pit_response, dict):
        pit_id = pit_response.get("id") or pit_response.get("pit_id")
    else:
        pit_id = str(pit_response)

    logger.debug(f"Opened async PIT: {pit_id[:20]}...", index=index)

    try:
        search_after: list[Any] | None = None
        total_returned = 0

        while True:
            body: dict[str, Any] = {
                "query": query,
                "sort": sort,
                "size": page_size,
                "pit": {"id": pit_id, "keep_alive": pit_keep_alive},
                "track_total_hits": False,
            }

            if search_after:
                body["search_after"] = search_after

            if source:
                body["_source"] = source

            try:
                response = await client.search(body=body)
            except TypeError:
                response = await client.search(**body)

            if hasattr(response, "body"):
                response_dict = response.body
            elif isinstance(response, dict):
                response_dict = response
            else:
                response_dict = dict(response)

            hits = response_dict.get("hits", {}).get("hits", [])

            if not hits:
                logger.debug(
                    "Async pagination complete - no more results",
                    total_returned=total_returned,
                )
                break

            for hit in hits:
                yield hit
                total_returned += 1

                if max_results and total_returned >= max_results:
                    logger.debug(
                        f"Reached max_results limit ({max_results})",
                        total_returned=total_returned,
                    )
                    return

            search_after = hits[-1].get("sort")

            if not search_after:
                logger.warning("Last hit has no sort value, ending pagination")
                break

            logger.debug(
                f"Async fetched page with {len(hits)} documents",
                total_returned=total_returned,
            )

    finally:
        try:
            await client.close_point_in_time(body={"id": pit_id})
            logger.debug(f"Closed async PIT: {pit_id[:20]}...")
        except Exception as exc:
            logger.warning(f"Failed to close async PIT: {exc}")


def scroll_iterator(
    client: Elasticsearch,
    index: str,
    query: dict[str, Any],
    *,
    page_size: int = 1000,
    max_results: int | None = None,
    scroll_timeout: str = "5m",
    source: list[str] | None = None,
) -> Iterator[dict[str, Any]]:
    """Iterate through search results using scroll API (legacy).

    Note: For new applications, prefer search_after_iterator which is
    more efficient and doesn't consume server resources for maintaining scroll context.

    Args:
        client: Elasticsearch client.
        index: Index to search.
        query: Elasticsearch query dict.
        page_size: Number of documents per page (default: 1000).
        max_results: Maximum total results to return (None for unlimited).
        scroll_timeout: Scroll context timeout (default: "5m").
        source: Fields to include in _source (None for all fields).

    Yields:
        Individual hit documents.
    """
    logger.warning("Using scroll API - consider search_after_iterator for better performance")

    # Initial search with scroll
    body: dict[str, Any] = {
        "query": query,
        "size": page_size,
    }

    if source:
        body["_source"] = source

    try:
        response = client.search(
            index=index,
            body=body,
            scroll=scroll_timeout,
        )
    except TypeError:
        response = client.search(
            index=index,
            scroll=scroll_timeout,
            **body,
        )

    if hasattr(response, "body"):
        response_dict = response.body
    elif isinstance(response, dict):
        response_dict = response
    else:
        response_dict = dict(response)

    scroll_id = response_dict.get("_scroll_id")
    total_returned = 0

    try:
        while True:
            hits = response_dict.get("hits", {}).get("hits", [])

            if not hits:
                break

            for hit in hits:
                yield hit
                total_returned += 1

                if max_results and total_returned >= max_results:
                    return

            # Fetch next batch
            response = client.scroll(scroll_id=scroll_id, scroll=scroll_timeout)

            if hasattr(response, "body"):
                response_dict = response.body
            elif isinstance(response, dict):
                response_dict = response
            else:
                response_dict = dict(response)

            scroll_id = response_dict.get("_scroll_id")

    finally:
        # Clean up scroll context
        if scroll_id:
            try:
                client.clear_scroll(scroll_id=scroll_id)
            except Exception as exc:
                logger.warning(f"Failed to clear scroll: {exc}")


__all__ = [
    "scroll_iterator",
    "search_after_iterator",
    "search_after_iterator_async",
]
