"""Tests for kokoro_chat.engine.message_manager — MessageManager."""
from __future__ import annotations

from typing import Any
from unittest.mock import AsyncMock

import pytest

from kokoro_chat.character.model import Character
from kokoro_chat.config import ChatConfig
from kokoro_chat.engine.message_manager import MessageManager
from kokoro_chat.storage.protocol import ChatMessage, RollingSummaryRecord


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

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 MockChatMessageWriter:
    """In-memory ChatMessageWriter for testing."""

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


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

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

    # Satisfy remaining CacheStore protocol methods
    async def get_attitude(self, user_id, character_id):
        return None

    async def set_attitude(self, user_id, character_id, data):
        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()


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

def make_manager(
    history_messages: list[ChatMessage] | None = None,
    cached_messages: list[dict[str, str]] | None = None,
    cached_summary: dict[str, Any] | None = None,
    summary_record: RollingSummaryRecord | None = None,
    config: ChatConfig | None = None,
) -> tuple[MessageManager, MockChatHistoryStore, MockChatMessageWriter, MockCacheStore, MockRollingSummaryStore]:
    history = MockChatHistoryStore(history_messages)
    writer = MockChatMessageWriter()
    cache = MockCacheStore(cached_messages, cached_summary)
    summary_store = MockRollingSummaryStore(summary_record)
    cfg = config or ChatConfig(recent_turns=6)
    mgr = MessageManager(history, writer, cache, summary_store, cfg)
    return mgr, history, writer, cache, summary_store


# ---------------------------------------------------------------------------
# get_for_prompt — cache hit (messages)
# ---------------------------------------------------------------------------

class TestGetForPromptCacheHit:

    @pytest.mark.asyncio
    async def test_returns_cached_messages(self) -> None:
        cached = [{"role": "user", "content": "hello"}, {"role": "assistant", "content": "hi"}]
        mgr, history, _, cache, _ = make_manager(cached_messages=cached)

        _, recent = await mgr.get_for_prompt("char1", "user1")

        assert recent == cached

    @pytest.mark.asyncio
    async def test_does_not_call_history_store_on_cache_hit(self) -> None:
        cached = [{"role": "user", "content": "hello"}]
        mgr, history, _, cache, _ = make_manager(cached_messages=cached)

        await mgr.get_for_prompt("char1", "user1")

        history.get_recent.assert_not_called()

    @pytest.mark.asyncio
    async def test_does_not_overwrite_cache_on_hit(self) -> None:
        cached = [{"role": "user", "content": "hello"}]
        mgr, _, _, cache, _ = make_manager(cached_messages=cached)

        await mgr.get_for_prompt("char1", "user1")

        cache.set_recent_messages.assert_not_called()


# ---------------------------------------------------------------------------
# get_for_prompt — cache miss (messages)
# ---------------------------------------------------------------------------

class TestGetForPromptCacheMiss:

    @pytest.mark.asyncio
    async def test_calls_history_store_on_cache_miss(self) -> None:
        msgs = [ChatMessage(role="user", content="hello")]
        mgr, history, _, cache, _ = make_manager(history_messages=msgs)

        await mgr.get_for_prompt("char1", "user1")

        history.get_recent.assert_called_once()

    @pytest.mark.asyncio
    async def test_returns_store_messages_formatted(self) -> None:
        msgs = [
            ChatMessage(role="user", content="hello"),
            ChatMessage(role="assistant", content="world"),
        ]
        mgr, _, _, _, _ = make_manager(history_messages=msgs)

        _, recent = await mgr.get_for_prompt("char1", "user1")

        assert recent == [
            {"role": "user", "content": "hello"},
            {"role": "assistant", "content": "world"},
        ]

    @pytest.mark.asyncio
    async def test_caches_messages_after_store_load(self) -> None:
        msgs = [ChatMessage(role="user", content="hello")]
        mgr, _, _, cache, _ = make_manager(history_messages=msgs)

        await mgr.get_for_prompt("char1", "user1")

        cache.set_recent_messages.assert_called_once()
        call_args = cache.set_recent_messages.call_args
        assert call_args[0][0] == "user1"
        assert call_args[0][1] == "char1"

    @pytest.mark.asyncio
    async def test_does_not_cache_empty_messages(self) -> None:
        mgr, _, _, cache, _ = make_manager(history_messages=[])

        await mgr.get_for_prompt("char1", "user1")

        cache.set_recent_messages.assert_not_called()


# ---------------------------------------------------------------------------
# get_for_prompt — summary from cache
# ---------------------------------------------------------------------------

class TestGetForPromptSummaryCacheHit:

    @pytest.mark.asyncio
    async def test_returns_cached_summary(self) -> None:
        cached_summary = {"summary": "previous conversation summary", "turn_count": 5}
        mgr, _, _, _, summary_store = make_manager(cached_summary=cached_summary)

        summary, _ = await mgr.get_for_prompt("char1", "user1")

        assert summary == "previous conversation summary"

    @pytest.mark.asyncio
    async def test_does_not_call_summary_store_on_cache_hit(self) -> None:
        cached_summary = {"summary": "cached summary", "turn_count": 3}
        mgr, _, _, _, summary_store = make_manager(cached_summary=cached_summary)

        await mgr.get_for_prompt("char1", "user1")

        summary_store.get_or_create.assert_not_called()

    @pytest.mark.asyncio
    async def test_missing_summary_key_returns_empty_string(self) -> None:
        cached_summary: dict[str, Any] = {"turn_count": 2}
        mgr, _, _, _, _ = make_manager(cached_summary=cached_summary)

        summary, _ = await mgr.get_for_prompt("char1", "user1")

        assert summary == ""


# ---------------------------------------------------------------------------
# get_for_prompt — summary from store
# ---------------------------------------------------------------------------

class TestGetForPromptSummaryCacheMiss:

    @pytest.mark.asyncio
    async def test_calls_summary_store_on_cache_miss(self) -> None:
        record = RollingSummaryRecord(summary="stored summary", turn_count=4)
        mgr, _, _, _, summary_store = make_manager(summary_record=record)

        await mgr.get_for_prompt("char1", "user1")

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

    @pytest.mark.asyncio
    async def test_returns_store_summary(self) -> None:
        record = RollingSummaryRecord(summary="stored summary", turn_count=4)
        mgr, _, _, _, _ = make_manager(summary_record=record)

        summary, _ = await mgr.get_for_prompt("char1", "user1")

        assert summary == "stored summary"

    @pytest.mark.asyncio
    async def test_caches_non_empty_summary_from_store(self) -> None:
        record = RollingSummaryRecord(summary="stored summary", turn_count=4)
        mgr, _, _, cache, _ = make_manager(summary_record=record)

        await mgr.get_for_prompt("char1", "user1")

        cache.set_rolling_summary.assert_called_once_with(
            "user1", "char1", "stored summary", 4,
        )

    @pytest.mark.asyncio
    async def test_does_not_cache_empty_summary_from_store(self) -> None:
        record = RollingSummaryRecord(summary="", turn_count=0)
        mgr, _, _, cache, _ = make_manager(summary_record=record)

        await mgr.get_for_prompt("char1", "user1")

        cache.set_rolling_summary.assert_not_called()

    @pytest.mark.asyncio
    async def test_none_summary_in_record_returns_empty_string(self) -> None:
        record = RollingSummaryRecord(summary=None, turn_count=0)  # type: ignore[arg-type]
        mgr, _, _, _, _ = make_manager(summary_record=record)

        summary, _ = await mgr.get_for_prompt("char1", "user1")

        assert summary == ""

    @pytest.mark.asyncio
    async def test_respects_recent_turns_config(self) -> None:
        config = ChatConfig(recent_turns=3)
        msgs = [ChatMessage(role="user", content=f"msg{i}") for i in range(6)]
        mgr, history, _, _, _ = make_manager(history_messages=msgs, config=config)

        await mgr.get_for_prompt("char1", "user1")

        history.get_recent.assert_called_once()
        call_args = history.get_recent.call_args
        assert call_args[1]["limit"] == 6  # recent_turns * 2


# ---------------------------------------------------------------------------
# add_message
# ---------------------------------------------------------------------------

class TestAddMessage:

    @pytest.mark.asyncio
    async def test_calls_writer_add_message(self) -> None:
        mgr, _, writer, _, _ = make_manager()

        await mgr.add_message("char1", "user1", "user", "hello", "short")

        writer.add_message.assert_called_once_with(
            user_id="user1",
            character_id="char1",
            role="user",
            content="hello",
            chat_mode="short",
        )

    @pytest.mark.asyncio
    async def test_calls_cache_append_message(self) -> None:
        mgr, _, _, cache, _ = make_manager()

        await mgr.add_message("char1", "user1", "assistant", "hi there", "long")

        cache.append_message.assert_called_once_with("user1", "char1", "assistant", "hi there")

    @pytest.mark.asyncio
    async def test_calls_both_writer_and_cache(self) -> None:
        mgr, _, writer, cache, _ = make_manager()

        await mgr.add_message("char1", "user1", "user", "test", "short")

        writer.add_message.assert_called_once()
        cache.append_message.assert_called_once()

    @pytest.mark.asyncio
    async def test_default_chat_mode_is_short(self) -> None:
        mgr, _, writer, _, _ = make_manager()

        await mgr.add_message("char1", "user1", "user", "hello")

        call_args = writer.add_message.call_args
        assert call_args[1]["chat_mode"] == "short"


# ---------------------------------------------------------------------------
# build_few_shot_prefix
# ---------------------------------------------------------------------------

class TestBuildFewShotPrefix:

    def test_empty_mes_example_returns_empty_list(self) -> None:
        char = Character(mes_example=[])

        result = MessageManager.build_few_shot_prefix(char)

        assert result == []

    def test_none_mes_example_returns_empty_list(self) -> None:
        char = Character()
        char.mes_example = []

        result = MessageManager.build_few_shot_prefix(char)

        assert result == []

    def test_valid_example_produces_correct_format(self) -> None:
        char = Character(mes_example=[{"user": "hello", "char": "hi there"}])

        result = MessageManager.build_few_shot_prefix(char)

        assert result == [
            {"role": "user", "content": "hello"},
            {"role": "assistant", "content": "hi there"},
        ]

    def test_multiple_examples_all_included(self) -> None:
        char = Character(mes_example=[
            {"user": "first question", "char": "first answer"},
            {"user": "second question", "char": "second answer"},
        ])

        result = MessageManager.build_few_shot_prefix(char)

        assert len(result) == 4
        assert result[0] == {"role": "user", "content": "first question"}
        assert result[1] == {"role": "assistant", "content": "first answer"}
        assert result[2] == {"role": "user", "content": "second question"}
        assert result[3] == {"role": "assistant", "content": "second answer"}

    def test_strips_whitespace_from_user(self) -> None:
        char = Character(mes_example=[{"user": "  hello  ", "char": "hi"}])

        result = MessageManager.build_few_shot_prefix(char)

        assert result[0]["content"] == "hello"

    def test_strips_whitespace_from_char(self) -> None:
        char = Character(mes_example=[{"user": "hello", "char": "  hi there  "}])

        result = MessageManager.build_few_shot_prefix(char)

        assert result[1]["content"] == "hi there"

    def test_empty_user_text_skipped(self) -> None:
        char = Character(mes_example=[{"user": "", "char": "hi"}])

        result = MessageManager.build_few_shot_prefix(char)

        assert len(result) == 1
        assert result[0] == {"role": "assistant", "content": "hi"}

    def test_whitespace_only_user_text_skipped(self) -> None:
        char = Character(mes_example=[{"user": "   ", "char": "hi"}])

        result = MessageManager.build_few_shot_prefix(char)

        assert len(result) == 1
        assert result[0] == {"role": "assistant", "content": "hi"}

    def test_empty_char_text_skipped(self) -> None:
        char = Character(mes_example=[{"user": "hello", "char": ""}])

        result = MessageManager.build_few_shot_prefix(char)

        assert len(result) == 1
        assert result[0] == {"role": "user", "content": "hello"}

    def test_missing_keys_produce_empty_stripped_to_nothing(self) -> None:
        char = Character(mes_example=[{}])

        result = MessageManager.build_few_shot_prefix(char)

        assert result == []

    def test_is_static_method(self) -> None:
        # Can be called without an instance
        char = Character(mes_example=[{"user": "a", "char": "b"}])

        result = MessageManager.build_few_shot_prefix(char)

        assert len(result) == 2
