"""Tests for memo filtering, config/pricing, token budget, and conversation formatting."""
from __future__ import annotations

from datetime import datetime, timedelta, timezone

import pytest

from kokoro_chat.config import MODEL_PRICING, ChatConfig, calculate_cost
from kokoro_chat.memory.memo import STALE_DAYS, STALE_KEYS, _filter_stale_memos
from kokoro_chat.prompt.budget import (
    DEFAULT_MAX_MEMO_TOKENS,
    DEFAULT_MAX_MES_EXAMPLES,
    DEFAULT_MAX_ROLLING_SUMMARY_TOKENS,
    DEFAULT_SYSTEM_PROMPT_TOKEN_BUDGET,
    _estimate_tokens,
)
from kokoro_chat.prompt.formatter import format_conversation_history

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

_NOW = datetime.now(timezone.utc)


def _memo(key: str, *, days_ago: int | None = None) -> dict:
    """Build a memo dict. If *days_ago* is None, omit updated_at."""
    m: dict = {"key": key, "value": "v"}
    if days_ago is not None:
        m["updated_at"] = _NOW - timedelta(days=days_ago)
    return m


# ===================================================================
# 1. _filter_stale_memos
# ===================================================================


class TestFilterStaleMemos:
    """Stale-memo filtering logic."""

    @pytest.mark.parametrize("days_ago", [0, 5, 100])
    def test_non_stale_key_kept_regardless_of_age(self, days_ago: int):
        memo = _memo("favourite_food", days_ago=days_ago)
        assert _filter_stale_memos([memo]) == [memo]

    @pytest.mark.parametrize("days_ago", [0, 1, 2])
    def test_stale_key_within_3_days_kept(self, days_ago: int):
        memo = _memo("current_mood", days_ago=days_ago)
        assert _filter_stale_memos([memo]) == [memo]

    def test_stale_key_at_exactly_3_days_kept(self):
        """Boundary: > 3 not >= 3, so 3-day-old memo stays."""
        memo = _memo("current_stress", days_ago=3)
        assert _filter_stale_memos([memo]) == [memo]

    def test_stale_key_at_4_days_dropped(self):
        memo = _memo("recent_event", days_ago=4)
        assert _filter_stale_memos([memo]) == []

    def test_stale_key_no_updated_at_kept(self):
        """Missing updated_at means the guard is skipped; memo stays."""
        memo = _memo("current_mood", days_ago=None)
        assert _filter_stale_memos([memo]) == [memo]

    def test_mixed_list_filtering(self):
        memos = [
            _memo("favourite_food", days_ago=100),   # non-stale key -> kept
            _memo("current_mood", days_ago=1),        # stale key, fresh -> kept
            _memo("current_stress", days_ago=10),     # stale key, old -> dropped
            _memo("recent_event", days_ago=None),     # stale key, no date -> kept
        ]
        result = _filter_stale_memos(memos)
        assert len(result) == 3
        assert result[0]["key"] == "favourite_food"
        assert result[1]["key"] == "current_mood"
        assert result[2]["key"] == "recent_event"

    def test_empty_list(self):
        assert _filter_stale_memos([]) == []


# ===================================================================
# 2. config — calculate_cost & ChatConfig
# ===================================================================


class TestCalculateCost:
    """Pricing calculation for various models."""

    @pytest.mark.parametrize(
        "model, inp, out, expected",
        [
            # Sonnet: (3.0, 15.0)
            ("claude-sonnet-4-5-20250929", 1000, 500, (1000 * 3.0 + 500 * 15.0) / 1_000_000),
            # Haiku: (0.80, 4.0)
            ("claude-haiku-4-5-20251001", 1000, 500, (1000 * 0.80 + 500 * 4.0) / 1_000_000),
            # MiniMax: (0.30, 1.20)
            ("MiniMax-M1", 1000, 500, (1000 * 0.30 + 500 * 1.20) / 1_000_000),
            # Unknown model falls back to Sonnet rates
            ("unknown-model-xyz", 1000, 500, (1000 * 3.0 + 500 * 15.0) / 1_000_000),
        ],
        ids=["sonnet", "haiku", "minimax", "unknown_default"],
    )
    def test_cost_calculation(self, model: str, inp: int, out: int, expected: float):
        assert calculate_cost(model, inp, out) == pytest.approx(expected)

    def test_zero_tokens(self):
        assert calculate_cost("claude-sonnet-4-5-20250929", 0, 0) == 0.0


class TestChatConfigDefaults:
    """ChatConfig default values."""

    def test_defaults(self):
        cfg = ChatConfig()
        assert cfg.chat_model == "claude-sonnet-4-5-20250929"
        assert cfg.chat_provider == "anthropic"
        assert cfg.memory_model == "MiniMax-M1"
        assert cfg.memory_provider == "minimax"
        assert cfg.max_tokens == 4096
        assert cfg.recent_turns == 6
        assert cfg.prompt_cache_enabled is False
        assert cfg.bg_task_concurrency == 5
        assert cfg.summarize_every == 8


# ===================================================================
# 3. prompt/budget — _estimate_tokens & constants
# ===================================================================


class TestEstimateTokens:
    """Rough token estimation via len//2."""

    @pytest.mark.parametrize(
        "text, expected",
        [
            ("", 0),
            ("a", 0),          # 1 // 2
            ("ab", 1),         # 2 // 2
            ("a" * 100, 50),   # long ascii
        ],
        ids=["empty", "single_char", "two_chars", "long_text"],
    )
    def test_ascii(self, text: str, expected: int):
        assert _estimate_tokens(text) == expected

    def test_chinese_text(self):
        # Python len counts code points, not bytes.
        assert _estimate_tokens("\u4f60\u597d") == 1  # 2 chars // 2


class TestBudgetConstants:
    """Budget default constants have expected values."""

    def test_defaults(self):
        assert DEFAULT_SYSTEM_PROMPT_TOKEN_BUDGET == 12000
        assert DEFAULT_MAX_MEMO_TOKENS == 500
        assert DEFAULT_MAX_MES_EXAMPLES == 1
        assert DEFAULT_MAX_ROLLING_SUMMARY_TOKENS == 80


# ===================================================================
# 4. prompt/formatter — format_conversation_history
# ===================================================================


def _turn(idx: int) -> list[dict]:
    """Return a user+assistant pair for turn *idx*."""
    return [
        {"role": "user", "content": f"u{idx}"},
        {"role": "assistant", "content": f"a{idx}"},
    ]


class TestFormatConversationHistory:
    """Conversation history truncation and formatting."""

    def test_empty_history(self):
        assert format_conversation_history([]) == []

    def test_single_turn(self):
        history = _turn(1)
        result = format_conversation_history(history)
        assert len(result) == 2
        assert result[0] == {"role": "user", "content": "u1"}
        assert result[1] == {"role": "assistant", "content": "a1"}

    @pytest.mark.parametrize("max_turns", [1, 2, 5])
    def test_exceeds_max_turns_truncated(self, max_turns: int):
        # Build 10 turns (20 messages)
        history: list[dict] = []
        for i in range(10):
            history.extend(_turn(i))
        result = format_conversation_history(history, max_turns=max_turns)
        assert len(result) == max_turns * 2

    def test_max_turns_zero(self):
        history = _turn(1)
        # max_turns=0 -> slice [-0:] keeps everything? Actually -0 == 0 in Python,
        # so history[0:] returns the full list. Let's verify actual behaviour.
        result = format_conversation_history(history, max_turns=0)
        # history[-0:] is history[0:] which is the full list
        assert result == [
            {"role": "user", "content": "u1"},
            {"role": "assistant", "content": "a1"},
        ]

    def test_preserves_role_and_content(self):
        history = [
            {"role": "user", "content": "hello", "extra": "ignored"},
            {"role": "assistant", "content": "hi", "meta": 123},
        ]
        result = format_conversation_history(history)
        for msg in result:
            assert set(msg.keys()) == {"role", "content"}
        assert result[0]["content"] == "hello"
        assert result[1]["content"] == "hi"
