"""Comprehensive tests for kokoro_chat.prompt.assembler helpers and assembly."""

from __future__ import annotations

import pytest

from kokoro_chat.character.model import Character
from kokoro_chat.lorebook.engine import LorebookEntry
from kokoro_chat.prompt.assembler import (
    _build_lorebook_engine,
    _format_lorebook_section,
    _format_memory_section,
    _get_unlocked_secrets,
    _select_post_history,
    _select_speech_pattern,
    assemble_static_system_prompt,
    assemble_static_system_prompt_blocks,
    assemble_system_prompt,
)


# ============================================================
# Fixtures
# ============================================================


@pytest.fixture
def char():
    return Character(
        id="test-char",
        name="TestChar",
        personality="friendly",
        description="A test character",
        likes=["music"],
        dislikes=["rain"],
    )


@pytest.fixture
def legacy_char():
    """Character without description -- falls back to legacy fields."""
    return Character(
        id="legacy-char",
        name="LegacyChar",
        personality="brave",
        background="a warrior from the north",
        speaking_style="rough and direct",
        mbti="ENTJ",
        appearance="tall with scars",
        likes=["swords"],
        dislikes=["cowardice"],
    )


# ============================================================
# 1. _select_speech_pattern
# ============================================================


class TestSelectSpeechPattern:
    def test_v2_mood_key_takes_priority(self):
        v2 = {"happy": "speaks cheerfully!", "default": "speaks normally"}
        result = _select_speech_pattern(v2, "legacy pattern", "happy")
        assert result == "speaks cheerfully!"

    def test_v2_default_when_mood_missing(self):
        v2 = {"default": "speaks normally"}
        result = _select_speech_pattern(v2, "legacy pattern", "angry")
        assert result == "speaks normally"

    def test_v2_empty_falls_to_legacy(self):
        result = _select_speech_pattern({}, "legacy pattern", "happy")
        assert result == "legacy pattern"

    def test_both_empty_returns_empty(self):
        result = _select_speech_pattern({}, "", "happy")
        assert result == ""

    def test_strips_whitespace(self):
        v2 = {"happy": "  cheerful!  "}
        result = _select_speech_pattern(v2, "", "happy")
        assert result == "cheerful!"

    def test_legacy_strips_whitespace(self):
        result = _select_speech_pattern({}, "  legacy  ", "happy")
        assert result == "legacy"

    def test_v2_mood_value_empty_falls_to_default(self):
        v2 = {"happy": "", "default": "fallback default"}
        result = _select_speech_pattern(v2, "legacy", "happy")
        assert result == "fallback default"


# ============================================================
# 2. _select_post_history
# ============================================================


class TestSelectPostHistory:
    def test_by_level_has_matching_level(self):
        by_level = {"friend": "be friendly", "lover": "be intimate"}
        result = _select_post_history(by_level, "legacy instructions", "friend")
        assert result == "be friendly"

    def test_by_level_no_match_returns_empty(self):
        by_level = {"friend": "be friendly"}
        result = _select_post_history(by_level, "legacy instructions", "stranger")
        assert result == ""

    def test_by_level_empty_falls_to_legacy(self):
        result = _select_post_history({}, "legacy instructions", "friend")
        assert result == "legacy instructions"

    def test_both_empty_returns_empty(self):
        result = _select_post_history({}, "", "friend")
        assert result == ""

    def test_strips_whitespace(self):
        by_level = {"friend": "  trimmed  "}
        result = _select_post_history(by_level, "", "friend")
        assert result == "trimmed"

    def test_legacy_strips_whitespace(self):
        result = _select_post_history({}, "  legacy  ", "friend")
        assert result == "legacy"


# ============================================================
# 3. _format_lorebook_section
# ============================================================


class TestFormatLorebookSection:
    def test_empty_list_returns_empty(self):
        assert _format_lorebook_section([]) == ""

    def test_entries_with_content(self):
        entries = [
            LorebookEntry(content="Lore A"),
            LorebookEntry(content="Lore B"),
        ]
        result = _format_lorebook_section(entries)
        assert result.startswith("[角色世界觀]\n")
        assert "Lore A" in result
        assert "Lore B" in result
        assert "\n---\n" in result

    def test_entries_with_empty_content_filtered(self):
        entries = [
            LorebookEntry(content="Valid lore"),
            LorebookEntry(content=""),
            LorebookEntry(content="   "),
        ]
        result = _format_lorebook_section(entries)
        assert result == "[角色世界觀]\nValid lore"

    def test_all_empty_content_returns_empty(self):
        entries = [
            LorebookEntry(content=""),
            LorebookEntry(content="   "),
        ]
        assert _format_lorebook_section(entries) == ""

    def test_single_entry_no_separator(self):
        entries = [LorebookEntry(content="Only lore")]
        result = _format_lorebook_section(entries)
        assert result == "[角色世界觀]\nOnly lore"
        assert "---" not in result


# ============================================================
# 4. _get_unlocked_secrets
# ============================================================


class TestGetUnlockedSecrets:
    def test_conditions_met(self):
        secrets = [
            {"content": "secret A", "reveal_condition": {"min_affection": 50, "min_trust": 30}},
        ]
        result = _get_unlocked_secrets(secrets, affection=60, trust=40)
        assert result == ["secret A"]

    def test_conditions_not_met_affection(self):
        secrets = [
            {"content": "secret A", "reveal_condition": {"min_affection": 80}},
        ]
        result = _get_unlocked_secrets(secrets, affection=50, trust=50)
        assert result == []

    def test_conditions_not_met_trust(self):
        secrets = [
            {"content": "secret A", "reveal_condition": {"min_trust": 80}},
        ]
        result = _get_unlocked_secrets(secrets, affection=90, trust=50)
        assert result == []

    def test_no_conditions_defaults_to_zero(self):
        secrets = [
            {"content": "always unlocked", "reveal_condition": {}},
        ]
        result = _get_unlocked_secrets(secrets, affection=0, trust=0)
        assert result == ["always unlocked"]

    def test_empty_content_filtered(self):
        secrets = [
            {"content": "", "reveal_condition": {}},
        ]
        result = _get_unlocked_secrets(secrets, affection=100, trust=100)
        assert result == []

    def test_multiple_secrets_partial_unlock(self):
        secrets = [
            {"content": "easy secret", "reveal_condition": {"min_affection": 10}},
            {"content": "hard secret", "reveal_condition": {"min_affection": 90}},
        ]
        result = _get_unlocked_secrets(secrets, affection=50, trust=0)
        assert result == ["easy secret"]

    def test_empty_secrets_list(self):
        assert _get_unlocked_secrets([], affection=100, trust=100) == []

    def test_both_conditions_must_be_met(self):
        secrets = [
            {"content": "both needed", "reveal_condition": {"min_affection": 50, "min_trust": 50}},
        ]
        # affection met but trust not
        assert _get_unlocked_secrets(secrets, affection=60, trust=30) == []
        # trust met but affection not
        assert _get_unlocked_secrets(secrets, affection=30, trust=60) == []
        # both met
        assert _get_unlocked_secrets(secrets, affection=50, trust=50) == ["both needed"]


# ============================================================
# 5. _format_memory_section
# ============================================================


class TestFormatMemorySection:
    def test_empty_memories_returns_empty(self):
        assert _format_memory_section([], "char-1") == ""

    def test_flat_format_no_source(self):
        memories = [
            {"text": "talked about music", "relevance": 0.85},
            {"text": "mentioned cooking", "relevance": 0.60},
        ]
        result = _format_memory_section(memories, "char-1")
        assert "[相關記憶]" in result
        assert "talked about music" in result
        assert "85%" in result
        assert "mentioned cooking" in result

    def test_split_format_with_source(self):
        memories = [
            {"text": "own memory", "relevance": 0.9, "source_character_id": "char-1"},
            {"text": "room memory", "relevance": 0.7, "source_character_id": "char-2"},
        ]
        result = _format_memory_section(memories, "char-1")
        assert "[你對用戶的記憶]" in result
        assert "own memory" in result
        assert "[同房間其他角色得知的資訊]" in result
        assert "room memory" in result

    def test_lookup_for_source_names(self):
        memories = [
            {"text": "room memory", "relevance": 0.7, "source_character_id": "char-2"},
        ]
        lookup = {"char-2": "Alice"}
        result = _format_memory_section(memories, "char-1", character_name_lookup=lookup)
        assert "Alice" in result

    def test_no_lookup_uses_raw_id(self):
        memories = [
            {"text": "room memory", "relevance": 0.7, "source_character_id": "char-2"},
        ]
        result = _format_memory_section(memories, "char-1")
        assert "char-2" in result

    def test_flat_format_when_no_character_id(self):
        """Falls to flat format when character_id is empty."""
        memories = [
            {"text": "some memory", "relevance": 0.5, "source_character_id": "char-1"},
        ]
        result = _format_memory_section(memories, "")
        assert "[相關記憶]" in result

    def test_only_own_memories(self):
        memories = [
            {"text": "my memory", "relevance": 0.8, "source_character_id": "char-1"},
        ]
        result = _format_memory_section(memories, "char-1")
        assert "[你對用戶的記憶]" in result
        assert "[同房間其他角色得知的資訊]" not in result

    def test_only_room_memories(self):
        memories = [
            {"text": "other memory", "relevance": 0.6, "source_character_id": "char-2"},
        ]
        result = _format_memory_section(memories, "char-1")
        assert "[你對用戶的記憶]" not in result
        assert "[同房間其他角色得知的資訊]" in result


# ============================================================
# 6. _build_lorebook_engine
# ============================================================


class TestBuildLorebookEngine:
    def test_none_returns_none(self):
        assert _build_lorebook_engine(None) is None

    def test_empty_dict_returns_none(self):
        assert _build_lorebook_engine({}) is None

    def test_empty_entries_returns_none(self):
        assert _build_lorebook_engine({"entries": []}) is None

    def test_valid_entries_returns_engine(self):
        book = {
            "entries": [
                {"keys": ["magic"], "content": "Magic is real.", "enabled": True},
            ],
            "scan_depth": 5,
            "token_budget": 200,
        }
        engine = _build_lorebook_engine(book)
        assert engine is not None
        assert len(engine.entries) == 1
        assert engine.entries[0].content == "Magic is real."
        assert engine.scan_depth == 5
        assert engine.token_budget == 200

    def test_entry_defaults(self):
        book = {
            "entries": [{"keys": ["test"]}],
        }
        engine = _build_lorebook_engine(book)
        entry = engine.entries[0]
        assert entry.content == ""
        assert entry.enabled is True
        assert entry.constant is False
        assert entry.selective is False
        assert entry.priority == 5
        assert entry.position == "after_char"
        assert entry.insertion_order == 100

    def test_multiple_entries(self):
        book = {
            "entries": [
                {"keys": ["a"], "content": "Entry A"},
                {"keys": ["b"], "content": "Entry B", "priority": 1},
            ],
        }
        engine = _build_lorebook_engine(book)
        assert len(engine.entries) == 2


# ============================================================
# 7. assemble_static_system_prompt
# ============================================================


class TestAssembleStaticSystemPrompt:
    def test_with_description_uses_description(self, char):
        result = assemble_static_system_prompt(char, "short")
        assert "[角色設定]" in result
        assert "A test character" in result
        assert "TestChar" in result

    def test_without_description_uses_legacy(self, legacy_char):
        result = assemble_static_system_prompt(legacy_char, "short")
        assert "[角色設定]" in result
        assert "brave" in result
        assert "a warrior from the north" in result
        assert "rough and direct" in result
        assert "ENTJ" in result
        assert "[外觀]" in result
        assert "tall with scars" in result

    def test_mandatory_sections_present(self, char):
        result = assemble_static_system_prompt(char, "short")
        assert "[平台規則]" in result
        assert "[角色設定]" in result
        assert "[回覆模式]" in result
        assert "[動態資訊查詢]" in result

    def test_traits_section(self, char):
        result = assemble_static_system_prompt(char, "short")
        assert "[特質]" in result
        assert "music" in result
        assert "rain" in result

    def test_speech_patterns_v2(self):
        c = Character(
            id="sp",
            name="SpeechChar",
            description="test",
            speech_patterns_v2={"happy": "yay!", "default": "hmm"},
        )
        result = assemble_static_system_prompt(c, "short", current_mood="happy")
        assert "[語癖]" in result
        assert "yay!" in result

    def test_anti_patterns(self):
        c = Character(
            id="ap",
            name="AntiChar",
            description="test",
            anti_patterns=[
                {"bad": "Hello sir", "correct": "Hey there", "why": "too formal"},
            ],
        )
        result = assemble_static_system_prompt(c, "short")
        assert "[不要這樣說]" in result
        assert "Hello sir" in result
        assert "Hey there" in result
        assert "too formal" in result

    def test_scene_from_param(self, char):
        result = assemble_static_system_prompt(char, "short", scene="a coffee shop")
        assert "[當前場景]" in result
        assert "a coffee shop" in result

    def test_scene_from_scenario(self):
        c = Character(id="sc", name="SceneChar", description="test", scenario="a park")
        result = assemble_static_system_prompt(c, "short")
        assert "[當前場景]" in result
        assert "a park" in result

    def test_no_scene(self, char):
        result = assemble_static_system_prompt(char, "short")
        assert "[當前場景]" not in result

    def test_personality_summary(self):
        c = Character(
            id="ps",
            name="SummaryChar",
            description="test",
            personality_summary="warm and caring",
        )
        result = assemble_static_system_prompt(c, "short")
        assert "[性格摘要]" in result
        assert "warm and caring" in result


# ============================================================
# 8. assemble_static_system_prompt_blocks
# ============================================================


class TestAssembleStaticSystemPromptBlocks:
    def test_returns_list_with_single_block(self, char):
        blocks = assemble_static_system_prompt_blocks(char)
        assert isinstance(blocks, list)
        assert len(blocks) == 1

    def test_block_structure(self, char):
        block = assemble_static_system_prompt_blocks(char)[0]
        assert block["type"] == "text"
        assert isinstance(block["text"], str)
        assert block["cache_control"] == {"type": "ephemeral"}

    def test_block_text_matches_static_prompt(self, char):
        block = assemble_static_system_prompt_blocks(char, "short")[0]
        expected = assemble_static_system_prompt(char, "short")
        assert block["text"] == expected


# ============================================================
# 9. assemble_system_prompt
# ============================================================


class TestAssembleSystemPrompt:
    def test_basic_sections_included(self, char):
        result = assemble_system_prompt(
            character=char,
            affection_score=10,
            chat_mode="short",
            events=[],
            memories=[],
            memos=[],
        )
        assert "[平台規則]" in result
        assert "[角色設定]" in result
        assert "[回覆模式]" in result
        assert "[關係狀態]" in result

    def test_events_included(self, char):
        events = [{"content": "Birthday party", "type": "relationship"}]
        result = assemble_system_prompt(
            character=char,
            affection_score=10,
            chat_mode="short",
            events=events,
            memories=[],
            memos=[],
        )
        assert "[重要事件]" in result
        assert "Birthday party" in result

    def test_memories_included(self, char):
        memories = [{"text": "likes jazz", "relevance": 0.9}]
        result = assemble_system_prompt(
            character=char,
            affection_score=10,
            chat_mode="short",
            events=[],
            memories=memories,
            memos=[],
        )
        assert "likes jazz" in result

    def test_memos_included(self, char):
        memos = [{"key": "name", "value": "Wade"}]
        result = assemble_system_prompt(
            character=char,
            affection_score=10,
            chat_mode="short",
            events=[],
            memories=[],
            memos=memos,
        )
        assert "[用戶資訊備忘]" in result
        assert "Wade" in result

    def test_degradation_drops_memories_first(self, char):
        """With a very tight token budget, memories should be dropped."""
        memories = [{"text": f"memory-{i}", "relevance": 0.5} for i in range(20)]
        events = [{"content": "important event", "type": "relationship"}]

        result = assemble_system_prompt(
            character=char,
            affection_score=10,
            chat_mode="short",
            events=events,
            memories=memories,
            memos=[],
            max_tokens=200,
        )
        # Core sections always present
        assert "[平台規則]" in result
        assert "[角色設定]" in result
        # With budget=200, many memories should have been dropped
        memory_count = sum(1 for i in range(20) if f"memory-{i}" in result)
        assert memory_count < 20

    def test_secrets_unlocking(self):
        c = Character(
            id="s",
            name="SecretChar",
            description="test",
            secrets=[
                {"content": "I am a robot", "reveal_condition": {"min_affection": 50}},
                {"content": "I can fly", "reveal_condition": {"min_affection": 90}},
            ],
        )
        result = assemble_system_prompt(
            character=c,
            affection_score=60,
            chat_mode="short",
            events=[],
            memories=[],
            memos=[],
        )
        assert "I am a robot" in result
        assert "I can fly" not in result

    def test_secrets_with_trust(self):
        c = Character(
            id="s",
            name="SecretChar",
            description="test",
            secrets=[
                {
                    "content": "trust secret",
                    "reveal_condition": {"min_affection": 10, "min_trust": 80},
                },
            ],
        )
        # Trust too low
        result_low = assemble_system_prompt(
            character=c,
            affection_score=50,
            chat_mode="short",
            events=[],
            memories=[],
            memos=[],
            trust_score=30,
        )
        assert "trust secret" not in result_low

        # Trust high enough
        result_high = assemble_system_prompt(
            character=c,
            affection_score=50,
            chat_mode="short",
            events=[],
            memories=[],
            memos=[],
            trust_score=80,
        )
        assert "trust secret" in result_high

    def test_lorebook_integration(self):
        c = Character(
            id="lb",
            name="LoreChar",
            description="test",
            character_book={
                "entries": [
                    {"keys": ["dragon"], "content": "Dragons are ancient beings.", "enabled": True},
                ],
                "scan_depth": 5,
                "token_budget": 500,
            },
        )
        result = assemble_system_prompt(
            character=c,
            affection_score=10,
            chat_mode="short",
            events=[],
            memories=[],
            memos=[],
            recent_messages=["I saw a dragon yesterday!"],
        )
        assert "[角色世界觀]" in result
        assert "Dragons are ancient beings." in result

    def test_lorebook_not_triggered_without_keyword(self):
        c = Character(
            id="lb",
            name="LoreChar",
            description="test",
            character_book={
                "entries": [
                    {"keys": ["dragon"], "content": "Dragons are ancient beings.", "enabled": True},
                ],
            },
        )
        result = assemble_system_prompt(
            character=c,
            affection_score=10,
            chat_mode="short",
            events=[],
            memories=[],
            memos=[],
            recent_messages=["I went to the store"],
        )
        assert "Dragons are ancient beings." not in result

    def test_affection_overrides(self):
        c = Character(
            id="ao",
            name="OverrideChar",
            description="test",
            affection_overrides={"stranger": "custom stranger behavior"},
        )
        result = assemble_system_prompt(
            character=c,
            affection_score=10,
            chat_mode="short",
            events=[],
            memories=[],
            memos=[],
        )
        assert "custom stranger behavior" in result

    def test_post_history_by_level(self):
        c = Character(
            id="ph",
            name="PostHistChar",
            description="test",
            post_history_by_level={"stranger": "remember to be formal"},
        )
        result = assemble_system_prompt(
            character=c,
            affection_score=10,
            chat_mode="short",
            events=[],
            memories=[],
            memos=[],
        )
        assert "[角色提醒]" in result
        assert "remember to be formal" in result

    def test_legacy_post_history_instructions(self):
        c = Character(
            id="ph",
            name="PostHistChar",
            description="test",
            post_history_instructions="always end with a question",
        )
        result = assemble_system_prompt(
            character=c,
            affection_score=10,
            chat_mode="short",
            events=[],
            memories=[],
            memos=[],
        )
        assert "[角色提醒]" in result
        assert "always end with a question" in result

    def test_relationship_level_changes_with_score(self):
        c = Character(id="rl", name="RelChar", description="test")
        # stranger level (0-20)
        result_stranger = assemble_system_prompt(
            character=c, affection_score=10, chat_mode="short",
            events=[], memories=[], memos=[],
        )
        assert "熟人" in result_stranger

        # friend level (21-40)
        result_friend = assemble_system_prompt(
            character=c, affection_score=30, chat_mode="short",
            events=[], memories=[], memos=[],
        )
        assert "朋友" in result_friend

    def test_attitude_section_appended(self, char):
        result = assemble_system_prompt(
            character=char,
            affection_score=10,
            chat_mode="short",
            events=[],
            memories=[],
            memos=[],
            attitude_section="Currently feeling happy",
        )
        assert "[當前態度（動態）]" in result
        assert "Currently feeling happy" in result

    def test_rolling_summary_included(self, char):
        result = assemble_system_prompt(
            character=char,
            affection_score=10,
            chat_mode="short",
            events=[],
            memories=[],
            memos=[],
            rolling_summary="Previously discussed travel plans.",
        )
        assert "[之前的對話摘要]" in result
        assert "Previously discussed travel plans." in result

    def test_story_context_included(self, char):
        result = assemble_system_prompt(
            character=char,
            affection_score=10,
            chat_mode="short",
            events=[],
            memories=[],
            memos=[],
            story_context="Chapter 3: The meeting",
        )
        assert "[故事進度]" in result
        assert "Chapter 3: The meeting" in result

    def test_creator_event_type(self, char):
        events = [{"content": "Server maintenance", "type": "creator"}]
        result = assemble_system_prompt(
            character=char,
            affection_score=10,
            chat_mode="short",
            events=events,
            memories=[],
            memos=[],
        )
        assert "【創作者事件】" in result
        assert "Server maintenance" in result
