"""Unit tests for the batch execution engine and convenience methods."""

from __future__ import annotations

from typing import Any
from unittest.mock import patch

import pytest
import respx

from httpx import Response

from pinecone._internal.batch import (
    _chunk,
    _create_progress_bar,
    _NoOpProgressBar,
    async_batch_execute,
    batch_execute,
)
from pinecone.async_index import AsyncIndex
from pinecone.index import Index
from pinecone.models.batch import BatchError, BatchResult


pytestmark = pytest.mark.unit


HOST = "my-index-abc123.svc.pinecone.io"
API_KEY = "test-key"
BASE_URL = f"https://{HOST}"


# ---------------------------------------------------------------------------
# _chunk utility
# ---------------------------------------------------------------------------


class TestChunk:
    def test_empty_list(self) -> None:
        assert _chunk([], 5) == []

    def test_exact_multiple(self) -> None:
        result = _chunk([1, 2, 3, 4], 2)
        assert result == [[1, 2], [3, 4]]

    def test_with_remainder(self) -> None:
        result = _chunk([1, 2, 3, 4, 5], 2)
        assert result == [[1, 2], [3, 4], [5]]

    def test_single_item(self) -> None:
        assert _chunk(["a"], 10) == [["a"]]

    def test_chunk_size_larger_than_list(self) -> None:
        result = _chunk([1, 2], 100)
        assert result == [[1, 2]]


# ---------------------------------------------------------------------------
# Progress bar helpers
# ---------------------------------------------------------------------------


class TestProgressBar:
    def test_noop_progress_bar(self) -> None:
        bar = _NoOpProgressBar()
        bar.update(1)
        bar.set_postfix_str("test")
        bar.close()

    def test_noop_as_context_manager(self) -> None:
        with _NoOpProgressBar() as bar:
            bar.update(1)

    def test_show_false_returns_noop(self) -> None:
        bar = _create_progress_bar(10, "test", show=False)
        assert isinstance(bar, _NoOpProgressBar)

    def test_missing_tqdm_returns_noop(self) -> None:
        with patch.dict("sys.modules", {"tqdm": None, "tqdm.auto": None}):
            bar = _create_progress_bar(10, "test", show=True)
            assert isinstance(bar, _NoOpProgressBar)


# ---------------------------------------------------------------------------
# BatchResult / BatchError models
# ---------------------------------------------------------------------------


class TestBatchResultModel:
    def test_no_errors(self) -> None:
        result = BatchResult(
            total_item_count=100,
            successful_item_count=100,
            failed_item_count=0,
            total_batch_count=10,
            successful_batch_count=10,
            failed_batch_count=0,
            errors=[],
        )
        assert not result.has_errors
        assert result.failed_items == []
        assert "SUCCESS" in repr(result)

    def test_with_errors(self) -> None:
        err = BatchError(
            batch_index=2,
            items=[{"id": "a"}, {"id": "b"}],
            error=RuntimeError("fail"),
            error_message="fail",
        )
        result = BatchResult(
            total_item_count=10,
            successful_item_count=8,
            failed_item_count=2,
            total_batch_count=5,
            successful_batch_count=4,
            failed_batch_count=1,
            errors=[err],
        )
        assert result.has_errors
        assert result.failed_items == [{"id": "a"}, {"id": "b"}]
        assert "PARTIAL FAILURE" in repr(result)

    def test_failed_items_flattened_from_multiple_errors(self) -> None:
        errors = [
            BatchError(
                batch_index=0,
                items=[{"id": "1"}],
                error=RuntimeError("e1"),
                error_message="e1",
            ),
            BatchError(
                batch_index=3,
                items=[{"id": "2"}, {"id": "3"}],
                error=RuntimeError("e2"),
                error_message="e2",
            ),
        ]
        result = BatchResult(
            total_item_count=10,
            successful_item_count=7,
            failed_item_count=3,
            total_batch_count=5,
            successful_batch_count=3,
            failed_batch_count=2,
            errors=errors,
        )
        assert len(result.failed_items) == 3

    def test_repr_html(self) -> None:
        result = BatchResult(
            total_item_count=5,
            successful_item_count=5,
            failed_item_count=0,
            total_batch_count=1,
            successful_batch_count=1,
            failed_batch_count=0,
            errors=[],
        )
        html = result._repr_html_()
        assert "BatchResult" in html
        assert "5" in html


class TestBatchErrorModel:
    def test_repr(self) -> None:
        err = BatchError(
            batch_index=1,
            items=[{"id": "x"}],
            error=ValueError("bad"),
            error_message="bad",
        )
        assert "batch_index=1" in repr(err)
        assert "item_count=1" in repr(err)


# ---------------------------------------------------------------------------
# Sync batch_execute
# ---------------------------------------------------------------------------


class TestBatchValidation:
    """Input validation for batch_size and concurrency parameters."""

    def test_batch_size_zero_raises(self) -> None:
        with pytest.raises(ValueError, match="batch_size must be >= 1"):
            batch_execute(
                items=[{"id": "1"}],
                operation=lambda _: None,
                batch_size=0,
                show_progress=False,
            )

    def test_batch_size_negative_raises(self) -> None:
        with pytest.raises(ValueError, match="batch_size must be >= 1"):
            batch_execute(
                items=[{"id": "1"}],
                operation=lambda _: None,
                batch_size=-5,
                show_progress=False,
            )

    def test_max_workers_zero_raises(self) -> None:
        with pytest.raises(ValueError, match="concurrency must be between 1 and 64"):
            batch_execute(
                items=[{"id": "1"}],
                operation=lambda _: None,
                batch_size=10,
                max_workers=0,
                show_progress=False,
            )

    def test_max_workers_too_high_raises(self) -> None:
        with pytest.raises(ValueError, match="concurrency must be between 1 and 64"):
            batch_execute(
                items=[{"id": "1"}],
                operation=lambda _: None,
                batch_size=10,
                max_workers=100,
                show_progress=False,
            )

    @pytest.mark.asyncio()
    async def test_async_batch_size_zero_raises(self) -> None:
        async def op(batch: list[dict[str, Any]]) -> None:
            pass

        with pytest.raises(ValueError, match="batch_size must be >= 1"):
            await async_batch_execute(
                items=[{"id": "1"}],
                operation=op,
                batch_size=0,
                show_progress=False,
            )

    @pytest.mark.asyncio()
    async def test_async_max_concurrency_zero_raises(self) -> None:
        async def op(batch: list[dict[str, Any]]) -> None:
            pass

        with pytest.raises(ValueError, match="concurrency must be between 1 and 64"):
            await async_batch_execute(
                items=[{"id": "1"}],
                operation=op,
                batch_size=10,
                max_concurrency=0,
                show_progress=False,
            )

    def test_validation_with_empty_items_batch_size(self) -> None:
        """Validation should occur even when items is empty."""
        with pytest.raises(ValueError, match="batch_size must be >= 1"):
            batch_execute(
                items=[],
                operation=lambda _: None,
                batch_size=0,
                show_progress=False,
            )

    def test_validation_with_empty_items_max_workers(self) -> None:
        """Validation should occur even when items is empty."""
        with pytest.raises(ValueError, match="concurrency must be between 1 and 64"):
            batch_execute(
                items=[],
                operation=lambda _: None,
                batch_size=10,
                max_workers=100,
                show_progress=False,
            )

    @pytest.mark.asyncio()
    async def test_async_validation_with_empty_items_batch_size(self) -> None:
        """Validation should occur even when items is empty."""

        async def op(batch: list[dict[str, Any]]) -> None:
            pass

        with pytest.raises(ValueError, match="batch_size must be >= 1"):
            await async_batch_execute(
                items=[],
                operation=op,
                batch_size=0,
                show_progress=False,
            )

    @pytest.mark.asyncio()
    async def test_async_validation_with_empty_items_max_concurrency(self) -> None:
        """Validation should occur even when items is empty."""

        async def op(batch: list[dict[str, Any]]) -> None:
            pass

        with pytest.raises(ValueError, match="concurrency must be between 1 and 64"):
            await async_batch_execute(
                items=[],
                operation=op,
                batch_size=10,
                max_concurrency=100,
                show_progress=False,
            )


class TestBatchExecute:
    def test_empty_items(self) -> None:
        result = batch_execute(
            items=[],
            operation=lambda _: None,
            batch_size=10,
            show_progress=False,
        )
        assert result.total_item_count == 0
        assert result.total_batch_count == 0

    def test_all_succeed(self) -> None:
        calls: list[list[dict[str, Any]]] = []

        def op(batch: list[dict[str, Any]]) -> None:
            calls.append(batch)

        items = [{"id": str(i)} for i in range(5)]
        result = batch_execute(
            items=items,
            operation=op,
            batch_size=2,
            show_progress=False,
        )

        assert result.total_item_count == 5
        assert result.successful_item_count == 5
        assert result.failed_item_count == 0
        assert result.total_batch_count == 3  # [2, 2, 1]
        assert result.successful_batch_count == 3
        assert result.failed_batch_count == 0
        assert not result.has_errors

    def test_partial_failure(self) -> None:
        call_count = 0

        def op(batch: list[dict[str, Any]]) -> None:
            nonlocal call_count
            call_count += 1
            # Fail every other batch
            if call_count % 2 == 0:
                raise RuntimeError("boom")

        items = [{"id": str(i)} for i in range(10)]
        result = batch_execute(
            items=items,
            operation=op,
            batch_size=2,
            max_workers=1,  # Sequential for deterministic ordering
            show_progress=False,
        )

        assert result.total_item_count == 10
        assert result.successful_item_count + result.failed_item_count == 10
        assert result.has_errors
        assert len(result.errors) > 0
        # Every error should have the original exception
        for err in result.errors:
            assert isinstance(err.error, RuntimeError)

    def test_all_fail(self) -> None:
        def op(batch: list[dict[str, Any]]) -> None:
            raise ValueError("nope")

        items = [{"id": "1"}, {"id": "2"}, {"id": "3"}]
        result = batch_execute(
            items=items,
            operation=op,
            batch_size=1,
            show_progress=False,
        )

        assert result.successful_item_count == 0
        assert result.failed_item_count == 3
        assert result.failed_batch_count == 3

    def test_single_batch(self) -> None:
        called = False

        def op(batch: list[dict[str, Any]]) -> None:
            nonlocal called
            called = True

        result = batch_execute(
            items=[{"id": "1"}],
            operation=op,
            batch_size=100,
            show_progress=False,
        )

        assert called
        assert result.total_batch_count == 1
        assert result.successful_batch_count == 1

    def test_failed_items_for_retry(self) -> None:
        def op(batch: list[dict[str, Any]]) -> None:
            if batch[0]["id"] == "bad":
                raise RuntimeError("fail")

        items = [{"id": "good"}, {"id": "bad"}]
        result = batch_execute(
            items=items,
            operation=op,
            batch_size=1,
            show_progress=False,
        )

        assert result.failed_items == [{"id": "bad"}]


# ---------------------------------------------------------------------------
# Async batch_execute
# ---------------------------------------------------------------------------


class TestAsyncBatchExecute:
    @pytest.mark.asyncio()
    async def test_empty_items(self) -> None:
        async def op(batch: list[dict[str, Any]]) -> None:
            pass

        result = await async_batch_execute(
            items=[],
            operation=op,
            batch_size=10,
            show_progress=False,
        )
        assert result.total_item_count == 0

    @pytest.mark.asyncio()
    async def test_all_succeed(self) -> None:
        calls: list[list[dict[str, Any]]] = []

        async def op(batch: list[dict[str, Any]]) -> None:
            calls.append(batch)

        items = [{"id": str(i)} for i in range(5)]
        result = await async_batch_execute(
            items=items,
            operation=op,
            batch_size=2,
            show_progress=False,
        )

        assert result.total_item_count == 5
        assert result.successful_item_count == 5
        assert result.total_batch_count == 3

    @pytest.mark.asyncio()
    async def test_partial_failure(self) -> None:
        async def op(batch: list[dict[str, Any]]) -> None:
            if batch[0]["id"] == "0":
                raise RuntimeError("boom")

        items = [{"id": str(i)} for i in range(4)]
        result = await async_batch_execute(
            items=items,
            operation=op,
            batch_size=2,
            show_progress=False,
        )

        # Batch [{"id":"0"},{"id":"1"}] should fail, batch [{"id":"2"},{"id":"3"}] succeeds
        assert result.has_errors
        assert result.successful_item_count == 2
        assert result.failed_item_count == 2

    @pytest.mark.asyncio()
    async def test_all_fail(self) -> None:
        async def op(batch: list[dict[str, Any]]) -> None:
            raise ValueError("nope")

        result = await async_batch_execute(
            items=[{"id": "1"}, {"id": "2"}],
            operation=op,
            batch_size=1,
            show_progress=False,
        )

        assert result.successful_item_count == 0
        assert result.failed_item_count == 2


# ---------------------------------------------------------------------------
# Sync convenience methods
# ---------------------------------------------------------------------------


class TestIndexBatchUpsert:
    @respx.mock
    def test_splits_into_batches(self) -> None:
        route = respx.post(f"{BASE_URL}/vectors/upsert").mock(
            return_value=Response(200, json={"upsertedCount": 2})
        )

        index = Index(host=HOST, api_key=API_KEY)
        vectors = [{"id": f"v{i}", "values": [0.1]} for i in range(5)]
        result = index.batch_upsert(
            vectors=vectors,
            batch_size=2,
            max_workers=1,
            show_progress=False,
        )

        assert route.call_count == 3  # 2 + 2 + 1
        assert result.total_item_count == 5
        assert result.successful_item_count == 5
        assert result.total_batch_count == 3

    @respx.mock
    def test_passes_namespace(self) -> None:
        route = respx.post(f"{BASE_URL}/vectors/upsert").mock(
            return_value=Response(200, json={"upsertedCount": 1})
        )

        index = Index(host=HOST, api_key=API_KEY)
        index.batch_upsert(
            vectors=[{"id": "v1", "values": [0.1]}],
            namespace="test-ns",
            batch_size=10,
            show_progress=False,
        )

        request = route.calls.last.request
        assert b'"namespace":"test-ns"' in request.content

    @respx.mock
    def test_partial_failure_captured(self) -> None:
        call_count = 0

        def side_effect(request: Any) -> Response:
            nonlocal call_count
            call_count += 1
            # Use 400 (not retried) so the failure is deterministic
            if call_count == 2:
                return Response(400, json={"message": "bad request"})
            return Response(200, json={"upsertedCount": 1})

        respx.post(f"{BASE_URL}/vectors/upsert").mock(side_effect=side_effect)

        index = Index(host=HOST, api_key=API_KEY)
        result = index.batch_upsert(
            vectors=[{"id": f"v{i}", "values": [0.1]} for i in range(3)],
            batch_size=1,
            max_workers=1,
            show_progress=False,
        )

        assert result.has_errors
        assert result.successful_batch_count == 2
        assert result.failed_batch_count == 1
        assert len(result.failed_items) == 1


class TestRecordsBatchUpsert:
    @respx.mock
    def test_splits_into_batches(self) -> None:
        route = respx.post(f"{BASE_URL}/records/namespaces/ns/upsert").mock(
            return_value=Response(200)
        )

        index = Index(host=HOST, api_key=API_KEY)
        records = [{"_id": f"r{i}", "text": f"t{i}"} for i in range(5)]
        result = index.records.batch_upsert(
            namespace="ns",
            records=records,
            batch_size=2,
            max_workers=1,
            show_progress=False,
        )

        assert route.call_count == 3
        assert result.total_item_count == 5
        assert result.successful_item_count == 5


class TestDocumentsBatchUpsert:
    @respx.mock
    def test_splits_into_batches(self) -> None:
        route = respx.post(f"{BASE_URL}/namespaces/ns/documents/upsert").mock(
            return_value=Response(200, json={"upserted_count": 2})
        )

        index = Index(host=HOST, api_key=API_KEY)
        docs = [{"_id": f"d{i}", "title": f"t{i}"} for i in range(5)]
        result = index.documents.batch_upsert(
            namespace="ns",
            documents=docs,
            batch_size=2,
            max_workers=1,
            show_progress=False,
        )

        assert route.call_count == 3
        assert result.total_item_count == 5
        assert result.successful_item_count == 5


# ---------------------------------------------------------------------------
# Async convenience methods
# ---------------------------------------------------------------------------


class TestAsyncIndexBatchUpsert:
    @pytest.mark.asyncio()
    @respx.mock
    async def test_splits_into_batches(self) -> None:
        route = respx.post(f"{BASE_URL}/vectors/upsert").mock(
            return_value=Response(200, json={"upsertedCount": 2})
        )

        async with AsyncIndex(host=HOST, api_key=API_KEY) as index:
            vectors = [{"id": f"v{i}", "values": [0.1]} for i in range(5)]
            result = await index.batch_upsert(
                vectors=vectors,
                batch_size=2,
                max_concurrency=1,
                show_progress=False,
            )

        assert route.call_count == 3
        assert result.total_item_count == 5
        assert result.successful_item_count == 5


class TestAsyncRecordsBatchUpsert:
    @pytest.mark.asyncio()
    @respx.mock
    async def test_splits_into_batches(self) -> None:
        route = respx.post(f"{BASE_URL}/records/namespaces/ns/upsert").mock(
            return_value=Response(200)
        )

        async with AsyncIndex(host=HOST, api_key=API_KEY) as index:
            records = [{"_id": f"r{i}", "text": f"t{i}"} for i in range(5)]
            result = await index.records.batch_upsert(
                namespace="ns",
                records=records,
                batch_size=2,
                max_concurrency=1,
                show_progress=False,
            )

        assert route.call_count == 3
        assert result.total_item_count == 5
        assert result.successful_item_count == 5


class TestAsyncDocumentsBatchUpsert:
    @pytest.mark.asyncio()
    @respx.mock
    async def test_splits_into_batches(self) -> None:
        route = respx.post(f"{BASE_URL}/namespaces/ns/documents/upsert").mock(
            return_value=Response(200, json={"upserted_count": 2})
        )

        async with AsyncIndex(host=HOST, api_key=API_KEY) as index:
            docs = [{"_id": f"d{i}", "title": f"t{i}"} for i in range(5)]
            result = await index.documents.batch_upsert(
                namespace="ns",
                documents=docs,
                batch_size=2,
                max_concurrency=1,
                show_progress=False,
            )

        assert route.call_count == 3
        assert result.total_item_count == 5
        assert result.successful_item_count == 5
