"""Tests for LLM timeout enforcement in extraction."""

import asyncio

import pytest

from arandu.config import MemoryConfig
from arandu.constants import ExtractionOutput
from arandu.protocols import LLMResult
from arandu.write.extract import extract_from_message


class SlowLLMProvider:
    """Mock LLM that never returns (sleeps forever)."""

    async def complete(self, messages, temperature=0, response_format=None, max_tokens=None) -> LLMResult:
        await asyncio.sleep(999)
        return LLMResult(text="{}")  # never reached

    async def embed(self, texts: list[str]) -> list[list[float]]:
        return [[0.0] * 768 for _ in texts]

    async def embed_one(self, text: str) -> list[float] | None:
        return [0.0] * 768


class TestExtractionTimeout:
    @pytest.mark.asyncio
    async def test_extraction_timeout_returns_empty(self):
        """When LLM times out, extraction returns empty result (fail-safe)."""
        config = MemoryConfig(extraction_timeout_sec=0.01)
        llm = SlowLLMProvider()

        result, _meta = await extract_from_message(
            "I live in São Paulo with my wife Ana",
            llm=llm,
            config=config,
        )

        assert isinstance(result, ExtractionOutput)
        assert len(result.facts) == 0
        assert len(result.entities) == 0

    @pytest.mark.asyncio
    async def test_extraction_timeout_multi_pass_returns_empty(self):
        """Multi-pass extraction also respects timeout."""
        config = MemoryConfig(extraction_timeout_sec=0.01)
        llm = SlowLLMProvider()

        result, _meta = await extract_from_message(
            "I live in São Paulo with my wife Ana",
            llm=llm,
            config=config,
        )

        assert isinstance(result, ExtractionOutput)
        assert len(result.facts) == 0
        assert len(result.entities) == 0

    @pytest.mark.asyncio
    async def test_extraction_completes_within_timeout(self):
        """Verify timeout doesn't interfere when LLM responds quickly."""

        class FastLLMProvider:
            async def complete(self, messages, temperature=0, response_format=None, max_tokens=None) -> LLMResult:
                return LLMResult(text='{"entities": [], "facts": [], "relations": []}')

            async def embed(self, texts: list[str]) -> list[list[float]]:
                return [[0.0] * 768 for _ in texts]

            async def embed_one(self, text: str) -> list[float] | None:
                return [0.0] * 768

        config = MemoryConfig(extraction_timeout_sec=30.0)
        result, _meta = await extract_from_message(
            "Hello",
            llm=FastLLMProvider(),
            config=config,
        )

        assert isinstance(result, ExtractionOutput)
