"""Tests for the BatchCollector micro-batching layer."""

from __future__ import annotations

import asyncio
from unittest.mock import AsyncMock

import pytest

from llmhosts.config import BatchConfig, LLMHostsConfig
from llmhosts.proxy.batch import BatchCollector
from llmhosts.proxy.models import UnifiedMessage, UnifiedRequest, UnifiedResponse

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_request(
    model: str = "llama3",
    content: str = "Hello",
    stream: bool = False,
) -> UnifiedRequest:
    """Build a minimal UnifiedRequest for testing."""
    return UnifiedRequest(
        model=model,
        messages=[UnifiedMessage(role="user", content=content)],
        stream=stream,
    )


def _make_response(content: str = "Hi there!", model: str = "llama3") -> UnifiedResponse:
    """Build a minimal UnifiedResponse for testing."""
    return UnifiedResponse(
        content=content,
        finish_reason="stop",
        model=model,
        prompt_tokens=10,
        completion_tokens=5,
        total_tokens=15,
    )


def _make_dispatch_fn(
    response: UnifiedResponse | None = None,
    delay: float = 0.0,
) -> AsyncMock:
    """Create a mock dispatch function that returns a canned response.

    Parameters
    ----------
    response:
        The response to return.  Defaults to a standard test response.
    delay:
        Artificial delay in seconds before returning the response.
    """
    resp = response or _make_response()

    async def _dispatch(req: UnifiedRequest) -> UnifiedResponse:
        if delay > 0:
            await asyncio.sleep(delay)
        return resp

    mock = AsyncMock(side_effect=_dispatch)
    return mock


# ---------------------------------------------------------------------------
# Config tests
# ---------------------------------------------------------------------------


@pytest.mark.unit
def test_batch_config_defaults() -> None:
    """BatchConfig defaults: disabled, 50ms window, max_size 16."""
    cfg = BatchConfig()
    assert cfg.enabled is False
    assert cfg.window_ms == 50.0
    assert cfg.max_size == 16


@pytest.mark.unit
def test_batch_config_on_root_config() -> None:
    """BatchConfig is accessible as a field on LLMHostsConfig."""
    cfg = LLMHostsConfig()
    assert isinstance(cfg.batch, BatchConfig)
    assert cfg.batch.enabled is False


@pytest.mark.unit
def test_batch_config_custom_values() -> None:
    """BatchConfig accepts custom values."""
    cfg = BatchConfig(enabled=True, window_ms=100.0, max_size=32)
    assert cfg.enabled is True
    assert cfg.window_ms == 100.0
    assert cfg.max_size == 32


# ---------------------------------------------------------------------------
# Passthrough / disabled mode
# ---------------------------------------------------------------------------


@pytest.mark.unit
async def test_submit_passthrough_when_disabled() -> None:
    """When batching is disabled, submit dispatches immediately (no queuing)."""
    mock_fn = _make_dispatch_fn()
    collector = BatchCollector(mock_fn, enabled=False)

    request = _make_request()
    response = await collector.submit(request)

    assert response.content == "Hi there!"
    mock_fn.assert_awaited_once()
    assert collector.pending_count == 0


@pytest.mark.unit
async def test_submit_passthrough_returns_correct_response() -> None:
    """Passthrough mode returns the exact response from the dispatch function."""
    custom_resp = _make_response(content="Custom answer", model="gpt-4")
    mock_fn = _make_dispatch_fn(response=custom_resp)
    collector = BatchCollector(mock_fn, enabled=False)

    response = await collector.submit(_make_request())

    assert response.content == "Custom answer"
    assert response.model == "gpt-4"


# ---------------------------------------------------------------------------
# Streaming bypass
# ---------------------------------------------------------------------------


@pytest.mark.unit
async def test_streaming_requests_bypass_batch() -> None:
    """Streaming requests are dispatched immediately, never batched."""
    mock_fn = _make_dispatch_fn()
    collector = BatchCollector(mock_fn, enabled=True, window_ms=5000, max_size=100)

    request = _make_request(stream=True)
    response = await collector.submit(request)

    assert response.content == "Hi there!"
    mock_fn.assert_awaited_once()
    # Nothing should be pending
    assert collector.pending_count == 0
    await collector.close()


@pytest.mark.unit
async def test_streaming_bypass_with_disabled_batch() -> None:
    """Streaming requests work correctly even when batching is disabled."""
    mock_fn = _make_dispatch_fn()
    collector = BatchCollector(mock_fn, enabled=False)

    response = await collector.submit(_make_request(stream=True))
    assert response.content == "Hi there!"
    mock_fn.assert_awaited_once()


# ---------------------------------------------------------------------------
# Batch collection and flushing
# ---------------------------------------------------------------------------


@pytest.mark.unit
async def test_multiple_requests_within_window_are_collected() -> None:
    """Multiple requests submitted within the window are batched together."""
    call_count = 0

    async def _counting_dispatch(req: UnifiedRequest) -> UnifiedResponse:
        nonlocal call_count
        call_count += 1
        return _make_response(content=f"resp-{call_count}")

    collector = BatchCollector(
        _counting_dispatch,
        enabled=True,
        window_ms=200,  # 200ms window -- long enough to collect requests
        max_size=100,  # high limit so window triggers flush
    )

    # Submit 3 requests concurrently -- they should all land in one batch
    tasks = [asyncio.create_task(collector.submit(_make_request(content=f"msg-{i}"))) for i in range(3)]

    responses = await asyncio.gather(*tasks)

    assert len(responses) == 3
    # All 3 dispatch calls should have been made
    assert call_count == 3
    assert collector.pending_count == 0
    await collector.close()


@pytest.mark.unit
async def test_batch_flushes_on_max_size() -> None:
    """Batch flushes immediately when max_size is reached."""
    mock_fn = _make_dispatch_fn()
    collector = BatchCollector(
        mock_fn,
        enabled=True,
        window_ms=10_000,  # very long window -- should not be the trigger
        max_size=3,
    )

    # Submit exactly max_size requests concurrently
    tasks = [asyncio.create_task(collector.submit(_make_request(content=f"msg-{i}"))) for i in range(3)]

    responses = await asyncio.gather(*tasks)

    assert len(responses) == 3
    assert mock_fn.await_count == 3
    assert collector.pending_count == 0
    await collector.close()


@pytest.mark.unit
async def test_batch_flushes_on_window_timeout() -> None:
    """Batch flushes when the window timer expires, even with fewer than max_size requests."""
    mock_fn = _make_dispatch_fn()
    collector = BatchCollector(
        mock_fn,
        enabled=True,
        window_ms=50,  # 50ms window
        max_size=100,  # very high -- timer should trigger first
    )

    # Submit a single request -- should be flushed after ~50ms
    response = await collector.submit(_make_request())

    assert response.content == "Hi there!"
    mock_fn.assert_awaited_once()
    assert collector.pending_count == 0
    await collector.close()


@pytest.mark.unit
async def test_batch_flushes_partial_on_timeout() -> None:
    """A partial batch (< max_size) flushes after the window timeout."""
    call_times: list[float] = []

    async def _timed_dispatch(req: UnifiedRequest) -> UnifiedResponse:
        call_times.append(asyncio.get_event_loop().time())
        return _make_response()

    collector = BatchCollector(
        _timed_dispatch,
        enabled=True,
        window_ms=30,
        max_size=50,
    )

    # Submit 2 requests -- well below max_size of 50
    tasks = [asyncio.create_task(collector.submit(_make_request(content=f"msg-{i}"))) for i in range(2)]

    responses = await asyncio.gather(*tasks)
    assert len(responses) == 2
    assert len(call_times) == 2
    await collector.close()


# ---------------------------------------------------------------------------
# Error handling
# ---------------------------------------------------------------------------


@pytest.mark.unit
async def test_batch_handles_per_request_errors() -> None:
    """One failing request in a batch does not affect other requests."""
    call_idx = 0

    async def _sometimes_fail(req: UnifiedRequest) -> UnifiedResponse:
        nonlocal call_idx
        call_idx += 1
        if call_idx == 2:
            raise ValueError("Backend error on request 2")
        return _make_response(content=f"ok-{call_idx}")

    collector = BatchCollector(
        _sometimes_fail,
        enabled=True,
        window_ms=200,
        max_size=3,
    )

    tasks = [asyncio.create_task(collector.submit(_make_request(content=f"msg-{i}"))) for i in range(3)]

    # Gather with return_exceptions so we can inspect individual results
    results = await asyncio.gather(*tasks, return_exceptions=True)

    # 2 should succeed, 1 should be a ValueError
    successes = [r for r in results if isinstance(r, UnifiedResponse)]
    failures = [r for r in results if isinstance(r, ValueError)]

    assert len(successes) == 2
    assert len(failures) == 1
    assert "Backend error on request 2" in str(failures[0])
    await collector.close()


@pytest.mark.unit
async def test_dispatch_fn_exception_propagated_to_caller() -> None:
    """When dispatch_fn raises, the exception propagates to the awaiting submit call."""

    async def _always_fail(req: UnifiedRequest) -> UnifiedResponse:
        raise RuntimeError("Backend down")

    collector = BatchCollector(
        _always_fail,
        enabled=True,
        window_ms=30,
        max_size=100,
    )

    with pytest.raises(RuntimeError, match="Backend down"):
        await collector.submit(_make_request())

    await collector.close()


@pytest.mark.unit
async def test_passthrough_exception_propagated() -> None:
    """When disabled, dispatch_fn exceptions propagate immediately."""

    async def _always_fail(req: UnifiedRequest) -> UnifiedResponse:
        raise ConnectionError("no connection")

    collector = BatchCollector(_always_fail, enabled=False)

    with pytest.raises(ConnectionError, match="no connection"):
        await collector.submit(_make_request())


# ---------------------------------------------------------------------------
# Empty batch
# ---------------------------------------------------------------------------


@pytest.mark.unit
async def test_empty_batch_does_not_dispatch() -> None:
    """An empty batch (no pending requests) does not call dispatch_fn."""
    mock_fn = _make_dispatch_fn()
    collector = BatchCollector(mock_fn, enabled=True, window_ms=10, max_size=10)

    # Manually trigger flush with nothing pending
    await collector._drain_and_flush()

    mock_fn.assert_not_awaited()
    await collector.close()


# ---------------------------------------------------------------------------
# Close / lifecycle
# ---------------------------------------------------------------------------


@pytest.mark.unit
async def test_close_rejects_new_submissions() -> None:
    """After close(), submit() raises RuntimeError."""
    mock_fn = _make_dispatch_fn()
    collector = BatchCollector(mock_fn, enabled=True, window_ms=50, max_size=10)

    await collector.close()

    with pytest.raises(RuntimeError, match="closed"):
        await collector.submit(_make_request())


@pytest.mark.unit
async def test_close_cancels_pending_requests() -> None:
    """Closing the collector rejects any pending requests."""
    # Use a dispatch that takes a long time so we can close while pending
    mock_fn = _make_dispatch_fn(delay=10.0)
    collector = BatchCollector(
        mock_fn,
        enabled=True,
        window_ms=5_000,  # long window
        max_size=100,
    )

    # Submit a request but don't await it yet -- it should get queued
    submit_task = asyncio.create_task(collector.submit(_make_request()))

    # Give the event loop a moment to process the submit
    await asyncio.sleep(0.01)

    # Close the collector while the request is pending
    await collector.close()

    # The submit should raise RuntimeError (batch closed)
    with pytest.raises(RuntimeError, match="closed"):
        await submit_task


# ---------------------------------------------------------------------------
# Properties
# ---------------------------------------------------------------------------


@pytest.mark.unit
async def test_properties_reflect_init_params() -> None:
    """Public properties match constructor arguments."""
    mock_fn = _make_dispatch_fn()
    collector = BatchCollector(
        mock_fn,
        enabled=True,
        window_ms=75.0,
        max_size=8,
    )

    assert collector.enabled is True
    assert collector.window_ms == 75.0
    assert collector.max_size == 8
    assert collector.pending_count == 0
    await collector.close()


@pytest.mark.unit
def test_repr_includes_state() -> None:
    """__repr__ includes enabled state, window, max_size, and pending count."""
    mock_fn = _make_dispatch_fn()
    collector = BatchCollector(mock_fn, enabled=True, window_ms=42.0, max_size=10)

    r = repr(collector)
    assert "enabled=True" in r
    assert "window_ms=42.0" in r
    assert "max_size=10" in r
    assert "pending=0" in r


# ---------------------------------------------------------------------------
# Multiple sequential batches
# ---------------------------------------------------------------------------


@pytest.mark.unit
async def test_multiple_sequential_batches() -> None:
    """After one batch flushes, subsequent requests form a new batch."""
    call_count = 0

    async def _counting_dispatch(req: UnifiedRequest) -> UnifiedResponse:
        nonlocal call_count
        call_count += 1
        return _make_response(content=f"resp-{call_count}")

    collector = BatchCollector(
        _counting_dispatch,
        enabled=True,
        window_ms=30,
        max_size=2,
    )

    # First batch of 2 (triggers on max_size)
    tasks1 = [asyncio.create_task(collector.submit(_make_request(content=f"batch1-{i}"))) for i in range(2)]
    responses1 = await asyncio.gather(*tasks1)
    assert len(responses1) == 2

    # Second batch of 2 (triggers on max_size)
    tasks2 = [asyncio.create_task(collector.submit(_make_request(content=f"batch2-{i}"))) for i in range(2)]
    responses2 = await asyncio.gather(*tasks2)
    assert len(responses2) == 2

    assert call_count == 4
    await collector.close()


# ---------------------------------------------------------------------------
# Concurrency stress
# ---------------------------------------------------------------------------


@pytest.mark.unit
async def test_concurrent_burst_of_requests() -> None:
    """A burst of many concurrent requests is dispatched correctly."""
    mock_fn = _make_dispatch_fn(delay=0.001)
    collector = BatchCollector(
        mock_fn,
        enabled=True,
        window_ms=100,
        max_size=5,
    )

    # Submit 12 requests concurrently.  With max_size=5, expect
    # at least 2 full batches (5+5) plus a partial (2).
    tasks = [asyncio.create_task(collector.submit(_make_request(content=f"msg-{i}"))) for i in range(12)]

    responses = await asyncio.gather(*tasks)
    assert len(responses) == 12
    assert mock_fn.await_count == 12
    await collector.close()
