"""Tests for Engram-inspired improvements to OMEGA's retrieval pipeline.

Covers all 6 improvements:
1. Hash-based fast-path lookup (trigram fingerprints)
2. Tiered memory cache (hot/warm/session)
3. Content canonicalization (NFKC, markdown strip)
4. Context-dependent gating (SurfacingContext)
5. Adaptive retrieval budget (QueryIntent)
6. Predictive prefetching
"""

import hashlib
import json
import os
import time

import pytest

from omega.sqlite_store import (
    SCHEMA_VERSION,
    SQLiteStore,
    SurfacingContext,
    QueryIntent,
    MemoryResult,
    _canonicalize,
    _trigram_fingerprint,
    _trigram_jaccard,
    _SURFACING_THRESHOLDS,
    _INTENT_WEIGHTS,
    _HOT_CACHE_SIZE,
    _HOT_CACHE_REFRESH_S,
    _FAST_PATH_MIN_OVERLAP,
)


# ============================================================================
# 1. Hash-Based Fast-Path Lookup
# ============================================================================


class TestFastPathLookup:
    """Tests for trigram fingerprint fast-path (#1)."""

    def test_trigram_fingerprint_basic(self):
        """Trigram fingerprint produces non-empty frozenset for normal text."""
        fp = _trigram_fingerprint("hello world")
        assert isinstance(fp, frozenset)
        assert len(fp) > 0

    def test_trigram_fingerprint_short_text(self):
        """Text shorter than 3 chars produces empty fingerprint."""
        assert _trigram_fingerprint("ab") == frozenset()
        assert _trigram_fingerprint("") == frozenset()

    def test_trigram_fingerprint_case_insensitive(self):
        """Fingerprints are case-insensitive (via canonicalization)."""
        fp1 = _trigram_fingerprint("Hello World")
        fp2 = _trigram_fingerprint("hello world")
        assert fp1 == fp2

    def test_trigram_jaccard_identical(self):
        """Identical fingerprints have Jaccard similarity of 1.0."""
        fp = _trigram_fingerprint("exact match text")
        assert _trigram_jaccard(fp, fp) == 1.0

    def test_trigram_jaccard_disjoint(self):
        """Completely different texts have low Jaccard similarity."""
        fp1 = _trigram_fingerprint("python programming language")
        fp2 = _trigram_fingerprint("xyz abc qrs tuv")
        sim = _trigram_jaccard(fp1, fp2)
        assert sim < 0.3

    def test_trigram_jaccard_empty(self):
        """Empty fingerprints return 0.0."""
        assert _trigram_jaccard(frozenset(), frozenset()) == 0.0
        fp = _trigram_fingerprint("hello")
        assert _trigram_jaccard(fp, frozenset()) == 0.0

    def test_fast_path_returns_results(self, store):
        """Fast-path should return results for keyword-heavy queries."""
        store.store(
            content="Error: sqlite3.OperationalError: database is locked",
            session_id="s1",
            metadata={"event_type": "error_pattern"},
        )
        # Touch to populate hot cache
        store._refresh_hot_cache()
        results = store._fast_path_lookup("sqlite3.OperationalError database locked")
        # May or may not match depending on keyword detection — just verify no crash
        assert isinstance(results, list)

    def test_fast_path_empty_for_conceptual(self, store):
        """Fast-path should return empty for non-keyword queries."""
        store.store(content="How authentication works in the system", session_id="s1")
        store._refresh_hot_cache()
        results = store._fast_path_lookup("how does authentication work")
        assert results == []


# ============================================================================
# 2. Tiered Memory Cache
# ============================================================================


class TestTieredCache:
    """Tests for tiered cache with hot set (#2)."""

    def test_hot_cache_populated_on_init(self, store):
        """Hot cache should be initialized (possibly empty in fresh store)."""
        assert hasattr(store, "_hot_memories")
        assert isinstance(store._hot_memories, dict)

    def test_hot_cache_refresh(self, store):
        """_refresh_hot_cache populates from high-access memories."""
        nid = store.store(content="Frequently accessed memory about testing patterns")
        # Bump access_count directly
        store._conn.execute(
            "UPDATE memories SET access_count = 50 WHERE node_id = ?", (nid,)
        )
        store._conn.commit()
        store._refresh_hot_cache()
        assert nid in store._hot_memories

    def test_hot_cache_invalidated_on_write(self, store):
        """Writing should mark hot cache for refresh."""
        store._hot_cache_ts = time.monotonic()  # Set fresh
        store.store(content="New memory triggers cache invalidation")
        assert store._hot_cache_ts == 0.0  # Reset by _invalidate_query_cache

    def test_session_cache_cleared(self, store):
        """clear_session_cache removes entries for the session."""
        store._session_cache["test-session"] = [
            MemoryResult(id="mem-test", content="cached")
        ]
        store.clear_session_cache("test-session")
        assert "test-session" not in store._session_cache

    def test_session_cache_populated_after_query(self, store):
        """Query with session_id should populate session cache."""
        store.store(content="Memory about Python testing patterns for cache test")
        results = store.query("Python testing patterns", session_id="sess-cache-test")
        # Session cache should be populated if there were results
        if results:
            assert "sess-cache-test" in store._session_cache

    def test_warm_cache_extended_ttl(self, store):
        """High-confidence results should use warm TTL (300s vs 60s)."""
        store.store(content="High confidence memory for warm cache test")
        # First query populates cache
        store.query("High confidence memory warm cache", use_cache=True)
        # Check cache entry has 3-tuple (ts, results, confidence)
        for key, val in store._query_cache.items():
            assert len(val) == 3  # (timestamp, results, confidence)
            break

    def test_query_cache_key_includes_surfacing_context(self, store):
        """Cache key should vary by surfacing context."""
        store.store(content="Context-sensitive caching test memory")
        store.query("caching test", surfacing_context=SurfacingContext.GENERAL)
        store.query("caching test", surfacing_context=SurfacingContext.ERROR_DEBUG)
        # Both should be cached separately
        keys = list(store._query_cache.keys())
        assert len(keys) >= 1  # At least one cached

    def test_hot_cache_size_limit(self, store):
        """Hot cache should respect _HOT_CACHE_SIZE limit."""
        # Store more than HOT_CACHE_SIZE memories
        for i in range(60):
            nid = store.store(content=f"Memory {i}: unique content for hot cache limit test item number {i}")
            store._conn.execute(
                "UPDATE memories SET access_count = ? WHERE node_id = ?",
                (100 - i, nid),
            )
        store._conn.commit()
        store._refresh_hot_cache()
        assert len(store._hot_memories) <= _HOT_CACHE_SIZE


# ============================================================================
# 3. Adaptive Retrieval Budget
# ============================================================================


class TestAdaptiveRetrievalBudget:
    """Tests for query intent classification (#5)."""

    def test_navigational_intent(self, store):
        """File paths and code identifiers should classify as NAVIGATIONAL."""
        assert store._classify_query_intent("/src/omega/bridge.py") == QueryIntent.NAVIGATIONAL
        assert store._classify_query_intent("`MemoryResult`") == QueryIntent.NAVIGATIONAL

    def test_factual_intent(self, store):
        """Factual questions should classify as FACTUAL."""
        assert store._classify_query_intent("what was the decision about caching") == QueryIntent.FACTUAL
        assert store._classify_query_intent("when did we add the entity system") == QueryIntent.FACTUAL

    def test_conceptual_intent(self, store):
        """Conceptual questions should classify as CONCEPTUAL."""
        assert store._classify_query_intent("how does the query pipeline work") == QueryIntent.CONCEPTUAL
        assert store._classify_query_intent("explain the architecture of omega") == QueryIntent.CONCEPTUAL

    def test_ambiguous_returns_none(self, store):
        """Ambiguous queries return None (use default weights)."""
        result = store._classify_query_intent("deployment rollback steps")
        assert result is None

    def test_intent_weights_structure(self):
        """All intent weights should be 5-tuples."""
        for intent, weights in _INTENT_WEIGHTS.items():
            assert len(weights) == 5
            assert all(isinstance(w, (int, float)) for w in weights)


# ============================================================================
# 4. Context-Dependent Gating
# ============================================================================


class TestContextDependentGating:
    """Tests for surfacing context thresholds (#4)."""

    def test_all_contexts_have_thresholds(self):
        """Every SurfacingContext enum value should have thresholds defined."""
        for ctx in SurfacingContext:
            assert ctx in _SURFACING_THRESHOLDS
            thresholds = _SURFACING_THRESHOLDS[ctx]
            assert len(thresholds) == 4
            assert all(isinstance(t, (int, float)) for t in thresholds)

    def test_error_debug_lower_vec_threshold(self):
        """ERROR_DEBUG should have lower vec threshold than GENERAL."""
        gen = _SURFACING_THRESHOLDS[SurfacingContext.GENERAL]
        err = _SURFACING_THRESHOLDS[SurfacingContext.ERROR_DEBUG]
        assert err[0] < gen[0]  # Lower min_vec_similarity

    def test_session_start_relaxed_thresholds(self):
        """SESSION_START should have relaxed thresholds to surface more context at session start."""
        start = _SURFACING_THRESHOLDS[SurfacingContext.SESSION_START]
        # Thresholds should be reasonable but not overly strict
        assert start[0] <= 0.50  # min_vec_similarity not too high
        assert start[2] <= 0.12  # min_composite not too high

    def test_file_edit_boosted_context_weight(self):
        """FILE_EDIT should have boosted context weight."""
        gen = _SURFACING_THRESHOLDS[SurfacingContext.GENERAL]
        edit = _SURFACING_THRESHOLDS[SurfacingContext.FILE_EDIT]
        assert edit[3] > gen[3]  # Higher context_weight_boost

    def test_query_with_surfacing_context(self, store):
        """Query should accept surfacing_context parameter without error."""
        store.store(content="Test memory for surfacing context query integration")
        for ctx in SurfacingContext:
            results = store.query(
                "surfacing context test",
                surfacing_context=ctx,
                limit=5,
            )
            assert isinstance(results, list)


# ============================================================================
# 5. Predictive Prefetching
# ============================================================================


class TestPredictivePrefetching:
    """Tests for predictive prefetching (#6)."""

    def test_prefetch_with_file_stems(self, store):
        """prefetch_for_project with explicit file stems."""
        store.store(
            content="The bridge.py module handles all MCP tool routing",
            metadata={"project": "/proj/omega"},
        )
        count = store.prefetch_for_project("/proj/omega", file_stems=["bridge"])
        assert count >= 0  # May be 0 if content doesn't match LIKE

    def test_prefetch_auto_discovers_stems(self, store):
        """prefetch_for_project auto-discovers file stems from content."""
        for i in range(5):
            store.store(
                content=f"Change {i} to sqlite_store.py for performance improvement",
                metadata={"project": "/proj/omega"},
            )
            store._conn.execute(
                "UPDATE memories SET access_count = ? WHERE node_id = (SELECT node_id FROM memories ORDER BY id DESC LIMIT 1)",
                (10 + i,),
            )
        store._conn.commit()
        count = store.prefetch_for_project("/proj/omega")
        assert isinstance(count, int)

    def test_prefetch_empty_project(self, store):
        """Prefetching for non-existent project returns 0."""
        count = store.prefetch_for_project("/nonexistent/project")
        assert count == 0

    def test_prefetch_populates_cache(self, store):
        """Prefetched results should appear in _prefetch_cache."""
        store.store(
            content="Important pattern in sqlite_store for query optimization",
            metadata={"project": "/proj/test"},
        )
        store._conn.execute(
            "UPDATE memories SET access_count = 20 WHERE node_id = (SELECT node_id FROM memories ORDER BY id DESC LIMIT 1)"
        )
        store._conn.commit()
        store.prefetch_for_project("/proj/test", file_stems=["sqlite_store"])
        # Cache may or may not have entries depending on LIKE match
        assert isinstance(store._prefetch_cache, dict)


# ============================================================================
# 6. Content Canonicalization
# ============================================================================


class TestContentCanonicalization:
    """Tests for content canonicalization (#6)."""

    def test_canonicalize_basic(self):
        """Basic canonicalization: lowercase, strip whitespace."""
        assert _canonicalize("Hello World") == "hello world"

    def test_canonicalize_markdown(self):
        """Markdown formatting should be stripped."""
        result = _canonicalize("**bold** and `code` and ## heading")
        assert "**" not in result
        assert "`" not in result
        assert "##" not in result
        assert "bold" in result
        assert "code" in result

    def test_canonicalize_whitespace_collapse(self):
        """Multiple whitespace chars should collapse to single space."""
        result = _canonicalize("hello    world\n\nnew  paragraph")
        assert "  " not in result
        assert "\n" not in result

    def test_canonicalize_nfkc(self):
        """NFKC normalization should handle Unicode equivalences."""
        # Fullwidth 'A' (U+FF21) should normalize to regular 'a'
        result = _canonicalize("\uff21\uff22\uff23")
        assert result == "abc"

    def test_canonical_hash_stored(self, store):
        """Stored memories should have canonical_hash computed."""
        nid = store.store(content="Test canonical hash storage")
        row = store._conn.execute(
            "SELECT canonical_hash FROM memories WHERE node_id = ?", (nid,)
        ).fetchone()
        assert row is not None
        assert row[0] is not None
        expected = hashlib.sha256(
            _canonicalize("Test canonical hash storage").encode()
        ).hexdigest()
        assert row[0] == expected

    def test_canonical_dedup(self, store):
        """Content with different markdown but same canonical form should dedup."""
        nid1 = store.store(content="Important lesson about testing")
        nid2 = store.store(content="**Important** lesson about testing")
        assert nid1 == nid2
        assert store.node_count() == 1

    def test_canonical_dedup_whitespace(self, store):
        """Content differing only in whitespace should dedup."""
        nid1 = store.store(content="lesson about   testing patterns")
        nid2 = store.store(content="lesson about testing patterns")
        assert nid1 == nid2

    def test_canonical_hash_updated_on_edit(self, store):
        """update_node with new content should recompute canonical_hash."""
        nid = store.store(content="Original content for hash test")
        store.update_node(nid, content="Updated content for hash test")
        row = store._conn.execute(
            "SELECT canonical_hash FROM memories WHERE node_id = ?", (nid,)
        ).fetchone()
        expected = hashlib.sha256(
            _canonicalize("Updated content for hash test").encode()
        ).hexdigest()
        assert row[0] == expected

    def test_word_overlap_uses_canonicalization(self, store):
        """_word_overlap should canonicalize both query and searchable text."""
        # Markdown in content should not prevent matching
        overlap = store._word_overlap(
            ["important", "lesson"],
            "**Important** lesson about `testing`",
        )
        assert overlap > 0

    def test_canonicalize_preserves_content(self, store):
        """Original content should be preserved despite canonical dedup."""
        nid = store.store(content="**Bold** lesson about `code`")
        node = store.get_node(nid)
        assert node.content == "**Bold** lesson about `code`"


# ============================================================================
# Integration Tests
# ============================================================================


class TestEngramIntegration:
    """Integration tests combining multiple improvements."""

    def test_full_query_pipeline(self, store):
        """Full query through the enhanced pipeline."""
        store.store(
            content="Decision: use SQLite for persistence instead of Redis",
            metadata={"event_type": "decision", "project": "/proj/omega"},
        )
        store.store(
            content="Lesson: always validate input at API boundaries",
            metadata={"event_type": "lesson_learned", "project": "/proj/omega"},
        )
        store.store(
            content="Error: connection timeout when Redis is down",
            metadata={"event_type": "error_pattern", "project": "/proj/omega"},
        )

        # General query
        results = store.query(
            "SQLite persistence decision",
            surfacing_context=SurfacingContext.GENERAL,
            limit=5,
        )
        assert isinstance(results, list)

        # Error debug query
        results = store.query(
            "connection timeout error",
            surfacing_context=SurfacingContext.ERROR_DEBUG,
            limit=5,
        )
        assert isinstance(results, list)

    def test_schema_version_is_5(self, store):
        """SCHEMA_VERSION should be 8 after retrieval improvements."""
        assert SCHEMA_VERSION == 14
        row = store._conn.execute(
            "SELECT version FROM schema_version LIMIT 1"
        ).fetchone()
        assert row[0] == 14

    def test_stats_tracking(self, store):
        """Engram stats should be tracked."""
        store.store(content="Memory for stats tracking test")
        store._refresh_hot_cache()
        store.query("stats tracking test")
        # Stats dict should exist and have standard keys
        assert "queries" in store.stats
        assert "stores" in store.stats
