"""
Elasticsearch Bulk Operations.

Provides utilities for efficient bulk indexing with chunking,
partial failure handling, and progress reporting.
"""

from __future__ import annotations

from collections.abc import AsyncIterator, Callable, Iterator
from dataclasses import dataclass, field
from typing import (
    TYPE_CHECKING,
    Any,
)

from smf.exceptions import ElasticsearchBulkError
from smf.smf_logging import get_logger

if TYPE_CHECKING:
    from elasticsearch import AsyncElasticsearch, Elasticsearch

logger = get_logger("elasticsearch.bulk")


@dataclass
class BulkStats:
    """Statistics from a bulk operation.

    Attributes:
        total: Total number of documents processed.
        indexed: Number of successfully indexed documents.
        created: Number of newly created documents.
        updated: Number of updated documents.
        deleted: Number of deleted documents.
        failed: Number of failed operations.
        errors: List of error details for failed operations.
    """

    total: int = 0
    indexed: int = 0
    created: int = 0
    updated: int = 0
    deleted: int = 0
    failed: int = 0
    errors: list[dict[str, Any]] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        """Convert stats to dictionary."""
        return {
            "total": self.total,
            "indexed": self.indexed,
            "created": self.created,
            "updated": self.updated,
            "deleted": self.deleted,
            "failed": self.failed,
            "success_rate": self.success_rate,
            "error_count": len(self.errors),
        }

    @property
    def success_rate(self) -> float:
        """Calculate success rate as percentage."""
        if self.total == 0:
            return 100.0
        succeeded = self.indexed + self.created + self.updated + self.deleted
        return round((succeeded / self.total) * 100, 2)


def chunked(
    iterable: Iterator[Any],
    chunk_size: int,
) -> Iterator[list[Any]]:
    """Split an iterator into chunks of specified size.

    Args:
        iterable: Source iterator.
        chunk_size: Maximum items per chunk.

    Yields:
        Lists of items, each up to chunk_size in length.
    """
    chunk: list[Any] = []
    for item in iterable:
        chunk.append(item)
        if len(chunk) >= chunk_size:
            yield chunk
            chunk = []
    if chunk:
        yield chunk


async def chunked_async(
    iterable: AsyncIterator[Any],
    chunk_size: int,
) -> AsyncIterator[list[Any]]:
    """Split an async iterator into chunks of specified size.

    Args:
        iterable: Source async iterator.
        chunk_size: Maximum items per chunk.

    Yields:
        Lists of items, each up to chunk_size in length.
    """
    chunk: list[Any] = []
    async for item in iterable:
        chunk.append(item)
        if len(chunk) >= chunk_size:
            yield chunk
            chunk = []
    if chunk:
        yield chunk


def _process_bulk_response(
    response: dict[str, Any],
    chunk: list[dict[str, Any]],
    stats: BulkStats,
    on_error: Callable[[dict[str, Any], dict[str, Any]], None] | None,
) -> None:
    """Process a bulk API response and update statistics.

    Args:
        response: Elasticsearch bulk response.
        chunk: The documents that were sent.
        stats: Statistics object to update.
        on_error: Optional callback for individual errors.
    """
    items = response.get("items", [])

    for i, item in enumerate(items):
        # Each item has a single key indicating the action type
        for action_type in ["index", "create", "update", "delete"]:
            if action_type in item:
                action_result = item[action_type]
                status = action_result.get("status", 0)

                if 200 <= status < 300:
                    # Success
                    if action_type == "index":
                        stats.indexed += 1
                    elif action_type == "create":
                        stats.created += 1
                    elif action_type == "update":
                        stats.updated += 1
                    elif action_type == "delete":
                        stats.deleted += 1
                else:
                    # Failure
                    stats.failed += 1
                    error_info = {
                        "index": i,
                        "action": action_type,
                        "status": status,
                        "error": action_result.get("error", {}),
                    }
                    stats.errors.append(error_info)

                    if on_error and i < len(chunk):
                        on_error(chunk[i], error_info)

                break  # Only one action type per item


def bulk_index(
    client: Elasticsearch,
    index: str,
    documents: Iterator[dict[str, Any]],
    *,
    id_field: str | None = None,
    chunk_size: int = 500,
    refresh: bool = False,
    raise_on_error: bool = False,
    on_error: Callable[[dict[str, Any], dict[str, Any]], None] | None = None,
    on_chunk_complete: Callable[[BulkStats], None] | None = None,
) -> BulkStats:
    """Bulk index documents with chunking and error handling.

    Args:
        client: Elasticsearch client.
        index: Target index name.
        documents: Iterator of documents to index.
        id_field: Field to use as document ID (optional).
        chunk_size: Number of documents per bulk request (default: 500).
        refresh: Whether to refresh index after each chunk (default: False).
        raise_on_error: If True, raise ElasticsearchBulkError on any failure.
        on_error: Callback for individual document errors (doc, error_info).
        on_chunk_complete: Callback after each chunk completes (stats).

    Returns:
        BulkStats with operation statistics.

    Raises:
        ElasticsearchBulkError: If raise_on_error=True and any operations fail.

    Example:
        >>> def on_error(doc, error):
        ...     print(f"Failed to index: {doc['id']}: {error}")
        >>> stats = bulk_index(client, "my-index", documents, on_error=on_error)
        >>> print(f"Indexed {stats.indexed}/{stats.total} documents")
    """
    stats = BulkStats()

    for chunk in chunked(documents, chunk_size):
        operations: list[dict[str, Any]] = []

        for doc in chunk:
            stats.total += 1

            # Build the action metadata
            action: dict[str, Any] = {"_index": index}
            if id_field and id_field in doc:
                action["_id"] = doc[id_field]

            operations.append({"index": action})
            operations.append(doc)

        try:
            # Execute bulk request
            response = client.bulk(
                operations=operations,
                refresh=refresh,
            )

            # Process response
            if isinstance(response, dict):
                _process_bulk_response(response, chunk, stats, on_error)
            else:
                # Handle ObjectApiResponse wrapper
                response_dict = (
                    response.body
                    if hasattr(response, "body")
                    else dict(response)
                    if hasattr(response, "__iter__")
                    else {}
                )
                _process_bulk_response(response_dict, chunk, stats, on_error)

        except Exception as exc:
            logger.error(
                f"Bulk operation failed: {exc}",
                chunk_size=len(chunk),
                index=index,
                error=str(exc),
            )
            # Mark all documents in chunk as failed
            for i, doc in enumerate(chunk):
                stats.failed += 1
                error_info = {
                    "index": i,
                    "action": "index",
                    "error": {"type": "bulk_failure", "reason": str(exc)},
                }
                stats.errors.append(error_info)
                if on_error:
                    on_error(doc, error_info)

        if on_chunk_complete:
            on_chunk_complete(stats)

        logger.debug(
            f"Bulk chunk complete: {len(chunk)} documents",
            indexed=stats.indexed,
            failed=stats.failed,
        )

    if raise_on_error and stats.failed > 0:
        raise ElasticsearchBulkError(
            f"Bulk operation had {stats.failed} failures out of {stats.total} documents",
            succeeded=stats.indexed + stats.created + stats.updated,
            failed=stats.failed,
            errors=stats.errors,
        )

    logger.info(
        f"Bulk indexing complete: {stats.indexed}/{stats.total} documents",
        stats=stats.to_dict(),
    )

    return stats


async def bulk_index_async(
    client: AsyncElasticsearch,
    index: str,
    documents: AsyncIterator[dict[str, Any]],
    *,
    id_field: str | None = None,
    chunk_size: int = 500,
    refresh: bool = False,
    raise_on_error: bool = False,
    on_error: Callable[[dict[str, Any], dict[str, Any]], None] | None = None,
    on_chunk_complete: Callable[[BulkStats], None] | None = None,
) -> BulkStats:
    """Async bulk index documents with chunking and error handling.

    Args:
        client: Async Elasticsearch client.
        index: Target index name.
        documents: Async iterator of documents to index.
        id_field: Field to use as document ID (optional).
        chunk_size: Number of documents per bulk request (default: 500).
        refresh: Whether to refresh index after each chunk (default: False).
        raise_on_error: If True, raise ElasticsearchBulkError on any failure.
        on_error: Callback for individual document errors (doc, error_info).
        on_chunk_complete: Callback after each chunk completes (stats).

    Returns:
        BulkStats with operation statistics.

    Raises:
        ElasticsearchBulkError: If raise_on_error=True and any operations fail.

    Example:
        >>> async def generate_docs():
        ...     for i in range(1000):
        ...         yield {"id": i, "text": f"Document {i}"}
        >>> stats = await bulk_index_async(client, "my-index", generate_docs())
    """
    stats = BulkStats()

    async for chunk in chunked_async(documents, chunk_size):
        operations: list[dict[str, Any]] = []

        for doc in chunk:
            stats.total += 1

            action: dict[str, Any] = {"_index": index}
            if id_field and id_field in doc:
                action["_id"] = doc[id_field]

            operations.append({"index": action})
            operations.append(doc)

        try:
            response = await client.bulk(
                operations=operations,
                refresh=refresh,
            )

            if isinstance(response, dict):
                _process_bulk_response(response, chunk, stats, on_error)
            else:
                response_dict = (
                    response.body
                    if hasattr(response, "body")
                    else dict(response)
                    if hasattr(response, "__iter__")
                    else {}
                )
                _process_bulk_response(response_dict, chunk, stats, on_error)

        except Exception as exc:
            logger.error(
                f"Async bulk operation failed: {exc}",
                chunk_size=len(chunk),
                index=index,
                error=str(exc),
            )
            for i, doc in enumerate(chunk):
                stats.failed += 1
                error_info = {
                    "index": i,
                    "action": "index",
                    "error": {"type": "bulk_failure", "reason": str(exc)},
                }
                stats.errors.append(error_info)
                if on_error:
                    on_error(doc, error_info)

        if on_chunk_complete:
            on_chunk_complete(stats)

    if raise_on_error and stats.failed > 0:
        raise ElasticsearchBulkError(
            f"Bulk operation had {stats.failed} failures out of {stats.total} documents",
            succeeded=stats.indexed + stats.created + stats.updated,
            failed=stats.failed,
            errors=stats.errors,
        )

    logger.info(
        f"Async bulk indexing complete: {stats.indexed}/{stats.total} documents",
        stats=stats.to_dict(),
    )

    return stats


__all__ = [
    "BulkStats",
    "bulk_index",
    "bulk_index_async",
    "chunked",
    "chunked_async",
]
