"""Models for batch operations."""

from __future__ import annotations

from typing import Any

from msgspec import Struct

from pinecone.models._display import render_table


class BatchError(Struct, kw_only=True, frozen=True):
    """Details about a single failed batch within a batch operation.

    Attributes:
        batch_index: Zero-based position of this batch in the original sequence.
        items: The items that were in this batch (for retry).
        error: The exception that caused the failure.
        error_message: Human-readable description of the error.
    """

    batch_index: int
    items: list[dict[str, Any]]
    error: Exception
    error_message: str

    def __repr__(self) -> str:
        return (
            f"BatchError(batch_index={self.batch_index}, "
            f"item_count={len(self.items)}, "
            f"error_message={self.error_message!r})"
        )


class BatchResult(Struct, kw_only=True):
    """Aggregated result of a batch operation.

    Tracks how many items and batches succeeded or failed, and provides
    access to the failed items for easy retry.

    Attributes:
        total_item_count: Total number of items submitted.
        successful_item_count: Number of items in successful batches.
        failed_item_count: Number of items in failed batches.
        total_batch_count: Total number of batches executed.
        successful_batch_count: Number of batches that succeeded.
        failed_batch_count: Number of batches that failed.
        errors: List of ``BatchError`` objects describing each failure.

    Examples:
        Check for partial failure and retry:

        >>> vectors = [{"id": f"v{i}", "values": [0.1] * 128} for i in range(1000)]
        >>> result = index.batch_upsert(vectors=vectors, namespace="ns")
        >>> if result.has_errors:
        ...     print(f"{result.failed_item_count} items failed, retrying...")
        ...     retry = index.batch_upsert(vectors=result.failed_items, namespace="ns")
    """

    total_item_count: int
    successful_item_count: int
    failed_item_count: int
    total_batch_count: int
    successful_batch_count: int
    failed_batch_count: int
    errors: list[BatchError]

    @property
    def has_errors(self) -> bool:
        """Whether any batches failed."""
        return len(self.errors) > 0

    @property
    def failed_items(self) -> list[dict[str, Any]]:
        """All items from failed batches, flattened for retry.

        Pass directly back to the same batch method to retry only the
        items that failed on the previous attempt.

        Returns:
            list[dict[str, Any]]: Flat list of all items from failed batches.
        """
        items: list[dict[str, Any]] = []
        for error in self.errors:
            items.extend(error.items)
        return items

    def __repr__(self) -> str:
        status = "PARTIAL FAILURE" if self.has_errors else "SUCCESS"
        return (
            f"BatchResult({status}: "
            f"{self.successful_item_count}/{self.total_item_count} items, "
            f"{self.successful_batch_count}/{self.total_batch_count} batches)"
        )

    def _repr_html_(self) -> str:
        """Jupyter notebook HTML representation."""
        rows: list[tuple[str, str | int | float]] = [
            ("Total items:", self.total_item_count),
            ("Successful items:", self.successful_item_count),
            ("Failed items:", self.failed_item_count),
            ("Total batches:", self.total_batch_count),
            ("Successful batches:", self.successful_batch_count),
            ("Failed batches:", self.failed_batch_count),
        ]
        return render_table("BatchResult", rows)
