# SPDX-License-Identifier: AGPL-3.0-only
# Copyright (c) 2026 sol pbc

"""Tests for the Batch async batch processor."""

import asyncio
from unittest.mock import AsyncMock, patch

import pytest

from solstone.think.batch import Batch, BatchRequest
from solstone.think.models import GEMINI_FLASH, SchemaValidationError


def _result(text: str = "Response", finish_reason: str = "stop", **extra):
    return {"text": text, "finish_reason": finish_reason, **extra}


def test_batch_request_creation():
    """Test BatchRequest can be created with required and custom params."""
    # Required params only
    req = BatchRequest(contents="Test prompt", context="test.context")
    assert req.contents == "Test prompt"
    assert req.context == "test.context"
    assert req.temperature == 0.3
    assert req.response is None
    assert req.error is None

    # With generation options
    req2 = BatchRequest(
        contents=["Part 1", "Part 2"],
        context="test.context",
        temperature=0.7,
        json_output=True,
    )
    assert req2.contents == ["Part 1", "Part 2"]
    assert req2.temperature == 0.7
    assert req2.json_output is True


def test_batch_request_custom_attributes():
    """Test that arbitrary attributes can be added to BatchRequest."""
    req = BatchRequest(contents="Test", context="test.context")
    req.frame_id = 123
    req.stage = "initial"
    req.custom_data = {"foo": "bar"}

    assert req.frame_id == 123
    assert req.stage == "initial"
    assert req.custom_data == {"foo": "bar"}


@pytest.mark.asyncio
@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
async def test_batch_basic(mock_agenerate):
    """Test basic batch execution with single request."""
    mock_agenerate.return_value = _result("Response 1")

    # Create batch and add request
    batch = Batch(max_concurrent=5)
    req = batch.create(contents="What is 2+2?", context="test.calc")
    req.task_id = "calc1"
    batch.add(req)

    # Iterate and verify
    results = []
    async for completed_req in batch.drain_batch():
        results.append(completed_req)

    assert len(results) == 1
    assert results[0].task_id == "calc1"
    assert results[0].response == "Response 1"
    assert results[0].error is None
    assert results[0].duration > 0

    # Verify API was called correctly
    mock_agenerate.assert_called_once()
    call_kwargs = mock_agenerate.call_args[1]
    assert call_kwargs["contents"] == "What is 2+2?"
    assert call_kwargs["context"] == "test.calc"


@pytest.mark.asyncio
@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
async def test_batch_records_resolved_model(mock_agenerate):
    """The transport-reported model is recorded without a request override."""
    mock_agenerate.return_value = _result("Response", model=GEMINI_FLASH)

    batch = Batch(max_concurrent=5)
    req = batch.create(contents="Test", context="test.context")
    batch.add(req)

    results = []
    async for completed_req in batch.drain_batch():
        results.append(completed_req)

    assert len(results) == 1
    assert results[0].model_used == GEMINI_FLASH

    call_kwargs = mock_agenerate.call_args[1]
    assert "model" not in call_kwargs


@pytest.mark.asyncio
@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
async def test_batch_multiple_requests(mock_agenerate):
    """Test batch with multiple requests."""
    mock_agenerate.side_effect = [
        _result("Response 1"),
        _result("Response 2"),
        _result("Response 3"),
    ]

    batch = Batch(max_concurrent=2)

    # Add multiple requests
    req1 = batch.create(contents="Prompt 1", context="test.context")
    req1.id = 1
    batch.add(req1)

    req2 = batch.create(contents="Prompt 2", context="test.context")
    req2.id = 2
    batch.add(req2)

    req3 = batch.create(contents="Prompt 3", context="test.context")
    req3.id = 3
    batch.add(req3)

    # Collect results
    results = []
    async for req in batch.drain_batch():
        results.append(req)

    # Should have all 3 results
    assert len(results) == 3

    # Results may come in any order (concurrent execution)
    result_ids = {r.id for r in results}
    assert result_ids == {1, 2, 3}

    # All should have responses
    for r in results:
        assert r.response is not None
        assert r.error is None


@pytest.mark.asyncio
@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
async def test_batch_error_handling(mock_agenerate):
    """Test that errors are captured in request.error."""
    mock_agenerate.side_effect = ValueError("API error")

    batch = Batch(max_concurrent=5)
    req = batch.create(contents="Bad prompt", context="test.context")
    req.id = "error_test"
    batch.add(req)

    # Get result
    results = []
    async for r in batch.drain_batch():
        results.append(r)

    assert len(results) == 1
    assert results[0].id == "error_test"
    assert results[0].response is None
    assert results[0].error == "API error"
    assert results[0].duration > 0


@pytest.mark.asyncio
@patch("solstone.think.batch.resolve_provider")
@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
async def test_batch_preserves_exception_metadata(
    mock_agenerate, mock_resolve_provider
):
    """Structured provider metadata should survive the batch error path."""

    class ProviderError(RuntimeError):
        reason_code = "provider_key_invalid"
        reset_at_ms = 12345
        provider = "anthropic"

    mock_agenerate.side_effect = ProviderError("bad key")
    mock_resolve_provider.return_value = ("google", "unused")

    batch = Batch(max_concurrent=5)
    req = batch.create(contents="Bad prompt", context="test.context")
    batch.add(req)

    results = []
    async for r in batch.drain_batch():
        results.append(r)

    assert len(results) == 1
    assert results[0].error == "bad key"
    assert results[0].reason_code == "provider_key_invalid"
    assert results[0].reset_at_ms == 12345
    assert results[0].provider == "anthropic"
    mock_resolve_provider.assert_not_called()


@pytest.mark.asyncio
@patch("solstone.think.batch.resolve_provider")
@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
async def test_batch_classifies_exception_metadata_when_attrs_missing(
    mock_agenerate, mock_resolve_provider
):
    """Batch should classify generic exceptions and best-effort provider."""
    mock_agenerate.side_effect = ConnectionError("network down")
    mock_resolve_provider.return_value = ("google", "gemini-test")

    batch = Batch(max_concurrent=5)
    req = batch.create(contents="Bad prompt", context="test.context")
    batch.add(req)

    results = []
    async for r in batch.drain_batch():
        results.append(r)

    assert len(results) == 1
    assert results[0].reason_code == "network_unreachable"
    assert results[0].reset_at_ms is None
    assert results[0].provider == "google"
    mock_resolve_provider.assert_called_once_with("generate")


def test_batch_update_clears_error_metadata():
    batch = Batch(max_concurrent=5)
    req = batch.create(contents="Retry prompt", context="test.context")
    req.reason_code = "provider_key_invalid"
    req.provider = "anthropic"
    req.reset_at_ms = 12345
    req.error = "bad key"

    with patch.object(batch, "add") as mock_add:
        batch.update(req)

    assert req.reason_code is None
    assert req.provider is None
    assert req.reset_at_ms is None
    assert req.error is None
    mock_add.assert_called_once_with(req)


@pytest.mark.asyncio
@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
async def test_batch_dynamic_adding(mock_agenerate):
    """Test adding requests dynamically during iteration."""
    mock_agenerate.return_value = _result("Response")

    batch = Batch(max_concurrent=5)

    # Add initial request
    req1 = batch.create(contents="Initial", context="test.context")
    req1.stage = "initial"
    batch.add(req1)

    # Process and add more during iteration
    results = []
    added_followup = False

    async for req in batch.drain_batch():
        results.append(req)

        # After first result, add a follow-up
        if req.stage == "initial" and not added_followup:
            req2 = batch.create(contents="Follow-up", context="test.context")
            req2.stage = "followup"
            batch.add(req2)
            added_followup = True

    # Should have both results
    assert len(results) == 2
    stages = {r.stage for r in results}
    assert stages == {"initial", "followup"}


@pytest.mark.asyncio
@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
async def test_batch_retry_pattern(mock_agenerate):
    """Test retry pattern without bypassing the active model profile."""
    # First call fails, second succeeds
    call_count = 0

    async def mock_response(*args, **kwargs):
        nonlocal call_count
        call_count += 1
        if call_count == 1:
            raise ValueError("Transient error")
        return _result("Success on retry")

    mock_agenerate.side_effect = mock_response

    batch = Batch(max_concurrent=5)

    # Add initial request
    req1 = batch.create(contents="Test", context="test.context")
    req1.attempt = 1
    batch.add(req1)

    results = []
    async for req in batch.drain_batch():
        results.append(req)

        # If error, retry through the same active profile.
        if req.error and req.attempt == 1:
            retry = batch.create(
                contents=req.contents,
                context="test.context",
            )
            retry.attempt = 2
            batch.add(retry)

    # Should have both attempts
    assert len(results) == 2
    assert results[0].error is not None
    assert results[0].attempt == 1
    assert results[1].response == "Success on retry"
    assert results[1].attempt == 2


@pytest.mark.asyncio
@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
async def test_batch_factory_method(mock_agenerate):
    """Test that batch.create() factory method works correctly."""
    mock_agenerate.return_value = _result("Response")

    batch = Batch()

    # Use factory method
    req = batch.create(
        contents="Test",
        context="test.context",
        temperature=0.8,
        json_output=True,
    )

    assert isinstance(req, BatchRequest)
    assert req.contents == "Test"
    assert req.context == "test.context"
    assert req.temperature == 0.8
    assert req.json_output is True


@pytest.mark.asyncio
@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
async def test_batch_can_add_after_draining(mock_agenerate):
    """Test that adding after draining works (reusable batch)."""
    mock_agenerate.side_effect = [_result("Response 1"), _result("Response 2")]

    batch = Batch()

    # First batch
    req1 = batch.create(contents="First", context="test.context")
    req1.id = 1
    batch.add(req1)

    results = []
    async for req in batch.drain_batch():
        results.append(req)

    assert len(results) == 1
    assert results[0].id == 1

    # Add more work after draining
    req2 = batch.create(contents="Second", context="test.context")
    req2.id = 2
    batch.add(req2)

    async for req in batch.drain_batch():
        results.append(req)

    # Should have both results
    assert len(results) == 2
    assert {r.id for r in results} == {1, 2}


@pytest.mark.asyncio
@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
async def test_batch_empty_batch(mock_agenerate):
    """Test that empty batch (no requests) completes immediately."""
    batch = Batch()

    results = []
    async for req in batch.drain_batch():
        results.append(req)

    assert len(results) == 0


@pytest.mark.asyncio
@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
async def test_batch_concurrency_limit(mock_agenerate):
    """Test that semaphore limits concurrent requests."""
    # Track concurrent calls
    concurrent_calls = 0
    max_concurrent_seen = 0
    lock = asyncio.Lock()

    async def mock_with_tracking(*args, **kwargs):
        nonlocal concurrent_calls, max_concurrent_seen
        async with lock:
            concurrent_calls += 1
            max_concurrent_seen = max(max_concurrent_seen, concurrent_calls)

        await asyncio.sleep(0.1)  # Simulate API call

        async with lock:
            concurrent_calls -= 1

        return _result("Response")

    mock_agenerate.side_effect = mock_with_tracking

    # Create batch with max_concurrent=2
    batch = Batch(max_concurrent=2)

    # Add 5 requests
    for i in range(5):
        req = batch.create(contents=f"Request {i}", context="test.context")
        batch.add(req)

    results = []
    async for req in batch.drain_batch():
        results.append(req)

    assert len(results) == 5
    # Should never exceed max_concurrent=2
    assert max_concurrent_seen <= 2


@pytest.mark.asyncio
@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
async def test_batch_update_method(mock_agenerate):
    """Test batch.update() method for modifying and re-adding requests."""
    call_temperatures = []

    async def mock_track_options(*args, **kwargs):
        call_temperatures.append(kwargs["temperature"])
        return _result(f"Response at {kwargs['temperature']}")

    mock_agenerate.side_effect = mock_track_options

    batch = Batch(max_concurrent=5)

    # Create initial request
    req = batch.create(contents="Initial prompt", context="test.context")
    req.stage = "initial"
    batch.add(req)

    results = []
    result_count = 0
    async for completed_req in batch.drain_batch():
        result_count += 1
        # Capture the response at this point
        results.append((result_count, completed_req.response, completed_req.stage))

        # After first result, update generation options and re-add.
        if result_count == 1:
            batch.update(
                completed_req,
                contents="Updated prompt",
                temperature=0.8,
                stage="updated",  # Update custom attribute too
                custom_field="test_value",  # Add new custom attribute
            )

    # Should have both results
    assert len(results) == 2
    assert results[0][2] == "initial"  # First result was initial stage
    assert results[1][2] == "updated"  # Second result was updated stage

    assert call_temperatures == [0.3, 0.8]

    # Verify correct responses at each stage
    assert results[0][1] == "Response at 0.3"
    assert results[1][1] == "Response at 0.8"

    # Verify custom attribute was set
    assert req.custom_field == "test_value"


def test_batch_request_with_timeout():
    """Test BatchRequest can be created with timeout_s parameter."""
    req = BatchRequest(contents="Test prompt", context="test.context", timeout_s=30)
    assert req.timeout_s == 30

    req2 = BatchRequest(contents="Test prompt", context="test.context", timeout_s=60.5)
    assert req2.timeout_s == 60.5

    # Default is None
    req3 = BatchRequest(contents="Test prompt", context="test.context")
    assert req3.timeout_s is None


@pytest.mark.asyncio
@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
async def test_batch_timeout_passthrough(mock_agenerate):
    """Test that timeout_s is passed through to agenerate_with_result."""
    mock_agenerate.return_value = _result("Response")

    batch = Batch(max_concurrent=5)

    # Create request with timeout_s
    req = batch.create(contents="Test prompt", context="test.context", timeout_s=45)
    batch.add(req)

    results = []
    async for completed_req in batch.drain_batch():
        results.append(completed_req)

    assert len(results) == 1

    # Verify timeout_s was passed to agenerate_with_result
    mock_agenerate.assert_called_once()
    call_kwargs = mock_agenerate.call_args[1]
    assert call_kwargs["timeout_s"] == 45


def test_batch_rejects_provider_specific_client_override():
    with pytest.raises(TypeError, match="unexpected keyword argument 'client'"):
        Batch(max_concurrent=5, client=object())


@pytest.mark.asyncio
@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
async def test_batch_passes_json_schema_to_agenerate(mock_agenerate):
    """Test that json_schema is passed through to agenerate_with_result."""
    mock_agenerate.return_value = _result("Response")

    batch = Batch(max_concurrent=5)
    req = batch.create(
        contents="Test prompt",
        context="test.context",
        json_schema={"type": "object"},
    )
    batch.add(req)

    results = []
    async for completed_req in batch.drain_batch():
        results.append(completed_req)

    assert len(results) == 1

    call_kwargs = mock_agenerate.call_args[1]
    assert call_kwargs["json_schema"] == {"type": "object"}


@pytest.mark.asyncio
@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
async def test_batch_schema_validation_error_populates_request_error(mock_agenerate):
    mock_agenerate.side_effect = SchemaValidationError(
        [{"path": "", "constraint": "json_parse", "message": "empty"}],
        "",
    )

    batch = Batch(max_concurrent=5)
    req = batch.create(
        contents="Test prompt",
        context="test.context",
        json_schema={"type": "object"},
    )
    batch.add(req)

    results = []
    async for completed_req in batch.drain_batch():
        results.append(completed_req)

    assert len(results) == 1
    assert results[0].response is None
    assert "schema validation" in results[0].error


@pytest.mark.asyncio
@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
async def test_batch_non_json_length_finish_populates_text_length_error(mock_agenerate):
    mock_agenerate.return_value = _result("partial text", finish_reason="max_tokens")

    batch = Batch(max_concurrent=5)
    req = batch.create(contents="Test prompt", context="test.context")
    batch.add(req)

    results = []
    async for completed_req in batch.drain_batch():
        results.append(completed_req)

    assert len(results) == 1
    assert results[0].response is None
    assert results[0].reason_code == "incomplete_text_length"
    assert "Text response incomplete" in results[0].error


@pytest.mark.asyncio
@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
async def test_batch_non_json_stop_finish_succeeds(mock_agenerate):
    mock_agenerate.return_value = _result("complete text", finish_reason="stop")

    batch = Batch(max_concurrent=1)
    req = batch.create(
        contents="Test",
        context="test.context",
        json_output=False,
    )
    batch.add(req)

    results = []
    async for completed_req in batch.drain_batch():
        results.append(completed_req)

    assert len(results) == 1
    assert req.response == "complete text"
    assert req.error is None
    assert req.reason_code is None


@pytest.mark.asyncio
@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
async def test_batch_non_json_non_length_finish_is_provider_response_invalid(
    mock_agenerate,
):
    mock_agenerate.return_value = _result("", finish_reason="content_filter")

    batch = Batch(max_concurrent=5)
    req = batch.create(contents="Test prompt", context="test.context")
    batch.add(req)

    results = []
    async for completed_req in batch.drain_batch():
        results.append(completed_req)

    assert len(results) == 1
    assert results[0].response is None
    assert results[0].reason_code == "provider_response_invalid"


@pytest.mark.asyncio
@patch("solstone.think.batch.agenerate_with_result", new_callable=AsyncMock)
async def test_batch_full_result_schema_invalid_stays_hard_failure(mock_agenerate):
    mock_agenerate.return_value = _result(
        '{"field": "bad"}',
        schema_validation={
            "valid": False,
            "errors": [{"path": "/field", "constraint": "type", "message": "bad"}],
        },
    )

    batch = Batch(max_concurrent=5)
    req = batch.create(
        contents="Test prompt",
        context="test.context",
        json_schema={"type": "object"},
    )
    batch.add(req)

    results = []
    async for completed_req in batch.drain_batch():
        results.append(completed_req)

    assert len(results) == 1
    assert results[0].response is None
    assert "schema validation" in results[0].error
