"""Integration tests for Agent cache integration."""
import pytest

from logos_adk.agent import Agent
from logos_adk.cache import InMemoryCacheBackend
from logos_adk.providers.mock import MockProvider
from logos_adk.types import ContentPart, MediaChunk, Response, StreamDone, TextChunk, ToolCall, Usage


def _mock_provider(*responses):
    return MockProvider(responses=list(responses))


class TestAgentCacheRun:
    @pytest.mark.asyncio
    async def test_cache_hit_skips_llm(self):
        provider = _mock_provider("first call", "second call")
        backend = InMemoryCacheBackend()
        agent = Agent("test", model=provider).cache(backend)

        r1 = await agent.run("hello")
        assert r1.content == "first call"
        assert len(provider.call_history) == 1

        # Same prompt should hit cache
        r2 = await agent.run("hello", session_id="s2")
        assert r2.content == "first call"
        assert len(provider.call_history) == 1  # no new LLM call

    @pytest.mark.asyncio
    async def test_different_prompts_no_cache_hit(self):
        provider = _mock_provider("response A", "response B")
        backend = InMemoryCacheBackend()
        agent = Agent("test", model=provider).cache(backend)

        r1 = await agent.run("hello")
        assert r1.content == "response A"

        r2 = await agent.run("goodbye", session_id="s2")
        assert r2.content == "response B"
        assert len(provider.call_history) == 2

    @pytest.mark.asyncio
    async def test_tool_responses_not_cached(self):
        """Responses with tool calls should not be stored in cache."""
        from logos_adk.tools import tool

        @tool(name="search", description="Search")
        async def search(q: str) -> str:
            return "found it"

        tool_response = Response(
            content="",
            tool_calls=[ToolCall(id="1", name="search", arguments={"q": "x"})],
            usage=Usage(10, 5),
            raw={},
        )
        final_response = Response(
            content="search result",
            tool_calls=[],
            usage=Usage(10, 5),
            raw={},
        )
        provider = _mock_provider(tool_response, final_response, "second call final")

        backend = InMemoryCacheBackend()
        agent = Agent("test", model=provider).tools([search]).cache(backend)

        r1 = await agent.run("find something")
        assert r1.content == "search result"

        # tool_response + final_response = 2 LLM calls
        assert len(provider.call_history) == 2

    @pytest.mark.asyncio
    async def test_cache_respects_ttl(self):
        import time

        provider = _mock_provider("cached", "fresh")
        backend = InMemoryCacheBackend()
        agent = Agent("test", model=provider).cache(backend, ttl=0)

        await agent.run("hello")
        time.sleep(0.01)

        r2 = await agent.run("hello", session_id="s2")
        assert r2.content == "fresh"
        assert len(provider.call_history) == 2

    @pytest.mark.asyncio
    async def test_semantic_cache_hit(self):
        """Semantic cache should find similar prompts."""

        async def mock_embed(text):
            # Return same vector for similar texts
            if "hello" in text.lower() or "hi" in text.lower():
                return [1.0, 0.0, 0.0]
            return [0.0, 1.0, 0.0]

        provider = _mock_provider("greeting response", "other response")
        backend = InMemoryCacheBackend()
        agent = Agent("test", model=provider).cache(
            backend, semantic=True, embed_fn=mock_embed,
            similarity_threshold=0.99,
        )

        r1 = await agent.run("hello there")
        assert r1.content == "greeting response"
        assert len(provider.call_history) == 1

        # "hi" maps to same vector, should hit semantic cache
        r2 = await agent.run("hi there", session_id="s2")
        assert r2.content == "greeting response"
        assert len(provider.call_history) == 1

    @pytest.mark.asyncio
    async def test_semantic_cache_miss(self):
        """Different embeddings should miss semantic cache."""

        async def mock_embed(text):
            if "hello" in text.lower():
                return [1.0, 0.0, 0.0]
            return [0.0, 1.0, 0.0]

        provider = _mock_provider("hello resp", "goodbye resp")
        backend = InMemoryCacheBackend()
        agent = Agent("test", model=provider).cache(
            backend, semantic=True, embed_fn=mock_embed,
            similarity_threshold=0.99,
        )

        await agent.run("hello")
        r2 = await agent.run("goodbye", session_id="s2")
        assert r2.content == "goodbye resp"
        assert len(provider.call_history) == 2

    @pytest.mark.asyncio
    async def test_cache_with_no_backend_is_noop(self):
        """Agent without cache should work normally."""
        provider = _mock_provider("response")
        agent = Agent("test", model=provider)

        r = await agent.run("hello")
        assert r.content == "response"
        assert len(provider.call_history) == 1

    @pytest.mark.asyncio
    async def test_semantic_without_embed_fn_raises(self):
        from logos_adk.errors import ConfigError

        with pytest.raises(ConfigError, match="embed_fn"):
            Agent("test").cache(
                InMemoryCacheBackend(), semantic=True,
            )


class TestAgentCacheStream:
    @pytest.mark.asyncio
    async def test_stream_cache_hit(self):
        provider = MockProvider(
            responses=["streamed response"],
            stream_responses=["streamed response"],
        )
        backend = InMemoryCacheBackend()
        agent = Agent("test", model=provider).cache(backend)

        # First call — populates cache via stream
        chunks = []
        async for chunk in agent.stream("hello"):
            chunks.append(chunk)

        assert any(isinstance(c, TextChunk) for c in chunks)
        assert any(isinstance(c, StreamDone) for c in chunks)
        assert len(provider.call_history) == 1

        # Second call — should hit cache, yield TextChunk + StreamDone
        chunks2 = []
        async for chunk in agent.stream("hello", session_id="s2"):
            chunks2.append(chunk)

        text_chunks = [c for c in chunks2 if isinstance(c, TextChunk)]
        done_chunks = [c for c in chunks2 if isinstance(c, StreamDone)]

        assert len(text_chunks) == 1
        assert text_chunks[0].text == "streamed response"
        assert len(done_chunks) == 1
        assert done_chunks[0].response.content == "streamed response"
        # No new LLM call
        assert len(provider.call_history) == 1

    @pytest.mark.asyncio
    async def test_stream_cache_miss(self):
        provider = MockProvider(
            responses=["a", "b"],
            stream_responses=["a", "b"],
        )
        backend = InMemoryCacheBackend()
        agent = Agent("test", model=provider).cache(backend)

        async for _ in agent.stream("hello"):
            pass

        async for _ in agent.stream("goodbye", session_id="s2"):
            pass

        assert len(provider.call_history) == 2

    @pytest.mark.asyncio
    async def test_stream_cache_hit_replays_media_parts(self):
        provider = MockProvider(
            responses=[
                Response(
                    content="See attachment",
                    tool_calls=[],
                    usage=Usage(5, 3),
                    raw={"source": "provider"},
                    content_parts=[
                        ContentPart(type="text", text="See attachment"),
                        ContentPart(
                            type="image",
                            uri="https://example.com/generated.png",
                            mime_type="image/png",
                            name="generated",
                        ),
                    ],
                )
            ],
            stream_responses=[["See attachment"]],
        )
        backend = InMemoryCacheBackend()
        agent = Agent("test", model=provider).cache(backend)

        response = await agent.run("show me the result")
        assert response.content == "See attachment"
        assert response.content_parts[1].type == "image"
        assert len(provider.call_history) == 1

        chunks2 = [
            chunk
            async for chunk in agent.stream("show me the result", session_id="s2")
        ]

        text_chunks = [chunk for chunk in chunks2 if isinstance(chunk, TextChunk)]
        media_chunks = [chunk for chunk in chunks2 if isinstance(chunk, MediaChunk)]
        done_chunks = [chunk for chunk in chunks2 if isinstance(chunk, StreamDone)]

        assert [chunk.text for chunk in text_chunks] == ["See attachment"]
        assert len(media_chunks) == 1
        assert media_chunks[0].part.type == "image"
        assert media_chunks[0].part.uri == "https://example.com/generated.png"
        assert len(done_chunks) == 1
        assert done_chunks[0].response.content == "See attachment"
        assert done_chunks[0].response.content_parts[1].type == "image"
        assert len(provider.call_history) == 1
