"""Tests for kokoro_chat.engine.summarizer — Summarizer."""
from __future__ import annotations

from typing import Any
from unittest.mock import AsyncMock, MagicMock

import pytest

from kokoro_chat.config import ChatConfig
from kokoro_chat.engine.summarizer import Summarizer
from kokoro_chat.storage.protocol import ChatMessage, RollingSummaryRecord


# ---------------------------------------------------------------------------
# Mock implementations
# ---------------------------------------------------------------------------

class MockCacheStore:
    """In-memory CacheStore for testing."""

    def __init__(
        self,
        cached_summary: dict[str, Any] | None = None,
    ) -> None:
        self._cached_summary = cached_summary
        self.get_rolling_summary = AsyncMock(return_value=cached_summary)
        self.set_rolling_summary = AsyncMock()
        self.set_recent_messages = AsyncMock()

    async def get_attitude(self, user_id, character_id):
        return None

    async def set_attitude(self, user_id, character_id, data):
        pass

    async def get_recent_messages(self, user_id, character_id, limit=20):
        return None

    async def append_message(self, user_id, character_id, role, content):
        pass


class MockRollingSummaryStore:
    """In-memory RollingSummaryStore for testing."""

    def __init__(self, record: RollingSummaryRecord | None = None) -> None:
        self._record = record or RollingSummaryRecord(summary="", turn_count=0)
        self.get_or_create = AsyncMock(return_value=self._record)
        self.update = AsyncMock()


class MockChatHistoryStore:
    """In-memory ChatHistoryStore for testing."""

    def __init__(self, messages: list[ChatMessage] | None = None) -> None:
        self._messages = messages or []
        self.get_recent = AsyncMock(return_value=self._messages)


class MockUsageLogStore:
    """In-memory UsageLogStore for testing."""

    def __init__(self) -> None:
        self.log = AsyncMock()


def make_ai_response(content: str = "summary text", input_tokens: int = 100, output_tokens: int = 50) -> Any:
    """Build a mock AI response object."""
    resp = MagicMock()
    resp.content = content
    resp.usage.input_tokens = input_tokens
    resp.usage.output_tokens = output_tokens
    return resp


def make_summarizer(
    summary_record: RollingSummaryRecord | None = None,
    cached_summary: dict[str, Any] | None = None,
    history_messages: list[ChatMessage] | None = None,
    ai_response: Any = None,
    ai_side_effect: Exception | None = None,
    config: ChatConfig | None = None,
    usage_log: MockUsageLogStore | None = None,
) -> tuple[Summarizer, MockRollingSummaryStore, MockCacheStore, MockChatHistoryStore, AsyncMock]:
    summary_store = MockRollingSummaryStore(summary_record)
    cache = MockCacheStore(cached_summary)
    history = MockChatHistoryStore(history_messages)
    cfg = config or ChatConfig(summarize_every=8, memory_llm_model="claude-haiku-4-5-20251001")

    ai_chat_fn = AsyncMock()
    if ai_side_effect is not None:
        ai_chat_fn.side_effect = ai_side_effect
    else:
        ai_chat_fn.return_value = ai_response or make_ai_response()

    summarizer = Summarizer(
        summary_store=summary_store,
        history_store=history,
        cache=cache,
        ai_chat_fn=ai_chat_fn,
        config=cfg,
        usage_log=usage_log,
    )
    return summarizer, summary_store, cache, history, ai_chat_fn


# ---------------------------------------------------------------------------
# maybe_summarize — below threshold: no AI call
# ---------------------------------------------------------------------------

class TestMaybeSummarizebelowThreshold:

    @pytest.mark.asyncio
    async def test_below_threshold_does_not_call_ai(self) -> None:
        # turn_count starts at 0, increments to 1; threshold is 8
        record = RollingSummaryRecord(summary="", turn_count=0)
        summarizer, _, _, _, ai_chat_fn = make_summarizer(summary_record=record)

        await summarizer.maybe_summarize("char1", "user1", "Sakura")

        ai_chat_fn.assert_not_called()

    @pytest.mark.asyncio
    async def test_below_threshold_increments_turn_count(self) -> None:
        record = RollingSummaryRecord(summary="existing", turn_count=3)
        summarizer, summary_store, cache, _, _ = make_summarizer(summary_record=record)

        await summarizer.maybe_summarize("char1", "user1", "Sakura")

        summary_store.update.assert_called_once()
        call_args = summary_store.update.call_args
        # turn_count should be 3 + 1 = 4
        assert call_args[0][3] == 4

    @pytest.mark.asyncio
    async def test_below_threshold_updates_cache_with_new_count(self) -> None:
        record = RollingSummaryRecord(summary="existing", turn_count=5)
        summarizer, _, cache, _, _ = make_summarizer(summary_record=record)

        await summarizer.maybe_summarize("char1", "user1", "Sakura")

        cache.set_rolling_summary.assert_called_once()
        call_args = cache.set_rolling_summary.call_args
        # turn_count should be 5 + 1 = 6
        assert call_args[0][3] == 6


# ---------------------------------------------------------------------------
# maybe_summarize — at threshold: triggers summarization
# ---------------------------------------------------------------------------

class TestMaybeSummarizeAtThreshold:

    @pytest.mark.asyncio
    async def test_at_threshold_calls_ai_chat(self) -> None:
        # turn_count=7, summarize_every=8 → after increment becomes 8 → threshold hit
        record = RollingSummaryRecord(summary="prev summary", turn_count=7)
        msgs = [ChatMessage(role="user", content=f"msg{i}") for i in range(20)]
        cfg = ChatConfig(summarize_every=8)
        summarizer, _, _, _, ai_chat_fn = make_summarizer(
            summary_record=record, history_messages=msgs, config=cfg,
        )

        await summarizer.maybe_summarize("char1", "user1", "Sakura")

        ai_chat_fn.assert_called_once()

    @pytest.mark.asyncio
    async def test_at_threshold_updates_summary_with_ai_result(self) -> None:
        record = RollingSummaryRecord(summary="old summary", turn_count=7)
        msgs = [ChatMessage(role="user", content=f"msg{i}") for i in range(20)]
        cfg = ChatConfig(summarize_every=8)
        ai_resp = make_ai_response(content="  new compressed summary  ")
        summarizer, summary_store, _, _, _ = make_summarizer(
            summary_record=record,
            history_messages=msgs,
            ai_response=ai_resp,
            config=cfg,
        )

        await summarizer.maybe_summarize("char1", "user1", "Sakura")

        # The final update call should have the new summary
        last_call = summary_store.update.call_args_list[-1]
        assert last_call[0][2] == "new compressed summary"


# ---------------------------------------------------------------------------
# maybe_summarize — cache hit vs miss
# ---------------------------------------------------------------------------

class TestMaybeSummarizeCacheHandling:

    @pytest.mark.asyncio
    async def test_cache_hit_uses_cached_turn_count(self) -> None:
        cached = {"summary": "cached summary", "turn_count": 3}
        summarizer, summary_store, _, _, ai_chat_fn = make_summarizer(
            cached_summary=cached,
        )

        await summarizer.maybe_summarize("char1", "user1", "Sakura")

        # Should not call store.get_or_create since cache hit
        summary_store.get_or_create.assert_not_called()
        # turn_count incremented to 4
        summary_store.update.assert_called_once()
        call_args = summary_store.update.call_args
        assert call_args[0][3] == 4

    @pytest.mark.asyncio
    async def test_cache_miss_falls_back_to_summary_store(self) -> None:
        record = RollingSummaryRecord(summary="stored summary", turn_count=2)
        summarizer, summary_store, _, _, _ = make_summarizer(
            cached_summary=None,
            summary_record=record,
        )

        await summarizer.maybe_summarize("char1", "user1", "Sakura")

        summary_store.get_or_create.assert_called_once_with("user1", "char1")

    @pytest.mark.asyncio
    async def test_cache_hit_with_missing_turn_count_defaults_to_zero(self) -> None:
        cached: dict[str, Any] = {"summary": "some summary"}
        summarizer, summary_store, _, _, _ = make_summarizer(cached_summary=cached)

        await summarizer.maybe_summarize("char1", "user1", "Sakura")

        # turn_count defaults to 0, increments to 1
        call_args = summary_store.update.call_args
        assert call_args[0][3] == 1

    @pytest.mark.asyncio
    async def test_cache_hit_with_missing_summary_defaults_to_empty(self) -> None:
        cached: dict[str, Any] = {"turn_count": 2}
        summarizer, summary_store, _, _, _ = make_summarizer(cached_summary=cached)

        await summarizer.maybe_summarize("char1", "user1", "Sakura")

        call_args = summary_store.update.call_args
        assert call_args[0][2] == ""


# ---------------------------------------------------------------------------
# _do_summarize — fewer than 10 messages
# ---------------------------------------------------------------------------

class TestDoSummarizeFewerThanTenMessages:

    @pytest.mark.asyncio
    async def test_fewer_than_10_returns_existing_summary_unchanged(self) -> None:
        msgs = [ChatMessage(role="user", content=f"msg{i}") for i in range(8)]
        summarizer, _, _, _, ai_chat_fn = make_summarizer(history_messages=msgs)

        result = await summarizer._do_summarize("char1", "user1", "Sakura", "existing summary", 8)

        assert result == "existing summary"
        ai_chat_fn.assert_not_called()

    @pytest.mark.asyncio
    async def test_exactly_9_messages_returns_existing_summary(self) -> None:
        msgs = [ChatMessage(role="user", content=f"msg{i}") for i in range(9)]
        summarizer, _, _, _, ai_chat_fn = make_summarizer(history_messages=msgs)

        result = await summarizer._do_summarize("char1", "user1", "Sakura", "prev", 8)

        assert result == "prev"
        ai_chat_fn.assert_not_called()


# ---------------------------------------------------------------------------
# _do_summarize — success path
# ---------------------------------------------------------------------------

class TestDoSummarizeSuccess:

    @pytest.mark.asyncio
    async def test_calls_ai_chat_with_correct_model(self) -> None:
        msgs = [ChatMessage(role="user", content=f"msg{i}") for i in range(20)]
        cfg = ChatConfig(memory_llm_model="claude-haiku-4-5-20251001")
        summarizer, _, _, _, ai_chat_fn = make_summarizer(
            history_messages=msgs, config=cfg,
        )

        await summarizer._do_summarize("char1", "user1", "Sakura", "", 8)

        ai_chat_fn.assert_called_once()
        call_kwargs = ai_chat_fn.call_args[1]
        assert call_kwargs["model"] == "claude-haiku-4-5-20251001"
        assert call_kwargs["max_tokens"] == 300

    @pytest.mark.asyncio
    async def test_updates_summary_store_with_new_summary(self) -> None:
        msgs = [ChatMessage(role="user", content=f"msg{i}") for i in range(20)]
        ai_resp = make_ai_response(content="new summary content")
        summarizer, summary_store, _, _, _ = make_summarizer(
            history_messages=msgs,
            ai_response=ai_resp,
        )

        result = await summarizer._do_summarize("char1", "user1", "Sakura", "old", 8)

        assert result == "new summary content"
        summary_store.update.assert_called_once_with("user1", "char1", "new summary content", 8)

    @pytest.mark.asyncio
    async def test_updates_cache_with_new_summary(self) -> None:
        msgs = [ChatMessage(role="user", content=f"msg{i}") for i in range(20)]
        ai_resp = make_ai_response(content="new summary")
        summarizer, _, cache, _, _ = make_summarizer(
            history_messages=msgs,
            ai_response=ai_resp,
        )

        await summarizer._do_summarize("char1", "user1", "Sakura", "", 8)

        # set_rolling_summary should be called with the new summary
        cache.set_rolling_summary.assert_called_once_with("user1", "char1", "new summary", 8)

    @pytest.mark.asyncio
    async def test_updates_recent_messages_cache(self) -> None:
        msgs = [ChatMessage(role="user", content=f"msg{i}") for i in range(20)]
        summarizer, _, cache, _, _ = make_summarizer(history_messages=msgs)

        await summarizer._do_summarize("char1", "user1", "Sakura", "", 8)

        cache.set_recent_messages.assert_called_once()
        call_args = cache.set_recent_messages.call_args
        # Should pass user_id, char_id, and last 10 messages
        assert call_args[0][0] == "user1"
        assert call_args[0][1] == "char1"
        assert len(call_args[0][2]) == 10

    @pytest.mark.asyncio
    async def test_logs_usage_when_usage_log_provided(self) -> None:
        msgs = [ChatMessage(role="user", content=f"msg{i}") for i in range(20)]
        ai_resp = make_ai_response(input_tokens=150, output_tokens=75)
        usage_log = MockUsageLogStore()
        summarizer, _, _, _, _ = make_summarizer(
            history_messages=msgs,
            ai_response=ai_resp,
            usage_log=usage_log,
        )

        await summarizer._do_summarize("char1", "user1", "Sakura", "", 8)

        usage_log.log.assert_called_once()
        call_kwargs = usage_log.log.call_args[1]
        assert call_kwargs["user_id"] == "user1"
        assert call_kwargs["input_tokens"] == 150
        assert call_kwargs["output_tokens"] == 75
        assert call_kwargs["task_type"] == "summarize"
        assert call_kwargs["character_id"] == "char1"

    @pytest.mark.asyncio
    async def test_does_not_log_usage_when_usage_log_not_provided(self) -> None:
        msgs = [ChatMessage(role="user", content=f"msg{i}") for i in range(20)]
        summarizer, _, _, _, _ = make_summarizer(history_messages=msgs, usage_log=None)

        # Should not raise
        await summarizer._do_summarize("char1", "user1", "Sakura", "", 8)

    @pytest.mark.asyncio
    async def test_strips_whitespace_from_ai_response(self) -> None:
        msgs = [ChatMessage(role="user", content=f"msg{i}") for i in range(20)]
        ai_resp = make_ai_response(content="  padded summary  ")
        summarizer, _, _, _, _ = make_summarizer(
            history_messages=msgs, ai_response=ai_resp,
        )

        result = await summarizer._do_summarize("char1", "user1", "Sakura", "", 8)

        assert result == "padded summary"


# ---------------------------------------------------------------------------
# _do_summarize — exception path
# ---------------------------------------------------------------------------

class TestDoSummarizeException:

    @pytest.mark.asyncio
    async def test_ai_exception_returns_existing_summary(self) -> None:
        msgs = [ChatMessage(role="user", content=f"msg{i}") for i in range(20)]
        summarizer, _, _, _, _ = make_summarizer(
            history_messages=msgs,
            ai_side_effect=RuntimeError("LLM unavailable"),
        )

        result = await summarizer._do_summarize("char1", "user1", "Sakura", "fallback summary", 8)

        assert result == "fallback summary"

    @pytest.mark.asyncio
    async def test_ai_exception_does_not_update_store(self) -> None:
        msgs = [ChatMessage(role="user", content=f"msg{i}") for i in range(20)]
        summarizer, summary_store, _, _, _ = make_summarizer(
            history_messages=msgs,
            ai_side_effect=ValueError("bad response"),
        )

        await summarizer._do_summarize("char1", "user1", "Sakura", "old", 8)

        summary_store.update.assert_not_called()

    @pytest.mark.asyncio
    async def test_ai_exception_does_not_update_cache(self) -> None:
        msgs = [ChatMessage(role="user", content=f"msg{i}") for i in range(20)]
        summarizer, _, cache, _, _ = make_summarizer(
            history_messages=msgs,
            ai_side_effect=ConnectionError("network error"),
        )

        await summarizer._do_summarize("char1", "user1", "Sakura", "old", 8)

        cache.set_rolling_summary.assert_not_called()
        cache.set_recent_messages.assert_not_called()
