"""Tests for context compression — tiered hot/warm/cold."""

from __future__ import annotations

import asyncio
from datetime import UTC, datetime
from types import SimpleNamespace

from arandu.config import MemoryConfig
from arandu.read.context_compression import (
    CompressedContext,
    _est_tokens,
    compress_broad_context,
    compress_context,
)


def run_async(coro):
    return asyncio.run(coro)


def _make_fact_entry(
    entity: str = "person:Ana",
    attribute: str = "nome",
    value: str = "Ana",
    score: float = 0.9,
    *,
    is_stale: bool = False,
    valid_from: datetime | None = None,
) -> dict:
    """Create a fact entry dict with an embedded fact-like object."""
    fact = SimpleNamespace(
        is_stale=is_stale,
        attribute_key=attribute,
        valid_from=valid_from or datetime.now(UTC),
    )
    return {
        "fact": fact,
        "entity": entity,
        "attribute": attribute,
        "value": value,
        "score": score,
        "date": "2025-01-01",
    }


# ---------------------------------------------------------------------------
# Token estimation
# ---------------------------------------------------------------------------


def test_est_tokens_basic():
    """Heuristic: ~4 chars per token."""
    assert _est_tokens("abcd") == 1
    assert _est_tokens("12345678") == 2


def test_est_tokens_empty():
    """Empty string → 0 tokens."""
    assert _est_tokens("") == 0


# ---------------------------------------------------------------------------
# CompressedContext defaults
# ---------------------------------------------------------------------------


def test_compressed_context_defaults():
    """CompressedContext should have sane defaults."""
    cc = CompressedContext()
    assert cc.context_text == ""
    assert cc.hot_count == 0
    assert cc.warm_count == 0
    assert cc.cold_count == 0
    assert cc.total_tokens == 0


# ---------------------------------------------------------------------------
# compress_context
# ---------------------------------------------------------------------------


def test_compress_empty_facts():
    """Empty facts → empty context."""
    config = MemoryConfig(context_max_tokens=2000)
    result = run_async(compress_context([], [], config))
    assert result.context_text == ""
    assert result.hot_count == 0
    assert result.total_tokens == 0


def test_compress_hot_tier_only():
    """With ≤5 facts, only hot tier is populated."""
    facts = [_make_fact_entry(score=0.9 - i * 0.1) for i in range(3)]
    config = MemoryConfig(context_max_tokens=2000)
    result = run_async(compress_context(facts, [], config))
    assert result.hot_count == 3
    assert result.warm_count == 0
    assert "CORE MEMORY" in result.context_text


def test_compress_warm_tier_populated():
    """With >5 facts, warm tier should be populated."""
    facts = [_make_fact_entry(entity=f"person:P{i}", score=0.9 - i * 0.05) for i in range(10)]
    config = MemoryConfig(context_max_tokens=5000)
    result = run_async(compress_context(facts, [], config))
    assert result.hot_count == 5
    assert result.warm_count > 0
    assert "EXTENDED CONTEXT" in result.context_text


def test_compress_token_budget_respected():
    """Total tokens should not wildly exceed the budget."""
    facts = [_make_fact_entry(entity=f"person:P{i}", value="x" * 200, score=0.9 - i * 0.01) for i in range(20)]
    config = MemoryConfig(context_max_tokens=500)
    result = run_async(compress_context(facts, [], config))
    # Allow some slack (events get +200 buffer), but should be reasonable
    assert result.total_tokens <= 800


def test_compress_stale_facts_marked():
    """Stale facts should be marked with [STALE]."""
    facts = [_make_fact_entry(is_stale=True, score=0.95)]
    config = MemoryConfig(context_max_tokens=2000)
    result = run_async(compress_context(facts, [], config))
    assert "[STALE]" in result.context_text


def test_compress_with_events():
    """Events should appear in RELEVANT_EVENTS section."""
    facts = [_make_fact_entry(score=0.9)]
    events = [{"date": "2025-01-15", "text": "User mentioned feeling happy"}]
    config = MemoryConfig(context_max_tokens=2000)
    result = run_async(compress_context(facts, events, config))
    assert "RELEVANT_EVENTS" in result.context_text
    assert result.cold_count >= 1


def test_compress_with_meta_observations():
    """Meta-observations should appear in PATTERNS section."""
    facts = [_make_fact_entry(score=0.9)]
    obs = SimpleNamespace(
        observation_type="pattern",
        title="Weekly routine detected",
        confidence=0.85,
        times_reinforced=3,
    )
    config = MemoryConfig(context_max_tokens=2000)
    result = run_async(compress_context(facts, [], config, meta_observations=[obs]))
    assert "PATTERNS & INSIGHTS" in result.context_text


# ---------------------------------------------------------------------------
# compress_broad_context
# ---------------------------------------------------------------------------


def test_broad_context_basic():
    """Broad context should produce BRAIN OVERVIEW section."""
    cluster = SimpleNamespace(label="work", summary_text="Work-related facts", fact_count=5)
    cluster_facts = {"work": [{"attribute": "role", "value": "Engineer"}]}
    config = MemoryConfig(context_max_tokens=2000)
    result = run_async(compress_broad_context(cluster_facts, [cluster], config))
    assert "BRAIN OVERVIEW" in result.context_text
    assert result.hot_count >= 1


def test_broad_context_empty():
    """Empty clusters → empty context."""
    config = MemoryConfig(context_max_tokens=2000)
    result = run_async(compress_broad_context({}, [], config))
    assert result.context_text == ""
    assert result.hot_count == 0
