"""
tests/test_storage_robustness.py — Tests for storage robustness fixes.

Covers:
- Fix 1: Write lock on all write methods (verified via lock acquisition)
- Fix 2: NaN/zero-vector embedding validation
- Fix 3: Brute-force vector search cap
- Fix 4: Sessions endpoint query (list_sessions)
- Fix 5: SQLite integrity check on startup
- Fix 6: get_active_session_by_id implementation
"""
from __future__ import annotations

import asyncio
import logging
import math

import pytest
import pytest_asyncio

from cortex.storage import SQLiteAdapter


# ---------------------------------------------------------------------------
# Fixture
# ---------------------------------------------------------------------------

@pytest_asyncio.fixture
async def adapter():
    """Fresh in-memory SQLiteAdapter for each test."""
    a = SQLiteAdapter(db_path=":memory:", busy_timeout_ms=1000, enable_wal=False)
    await a.initialize()
    yield a
    await a.close()


# ---------------------------------------------------------------------------
# Fix 2: NaN/zero-vector embedding validation
# ---------------------------------------------------------------------------

class TestEmbeddingValidation:
    """Ensure store_embedding rejects invalid vectors."""

    async def test_nan_embedding_rejected(self, adapter: SQLiteAdapter):
        """NaN values in embedding should be silently skipped (not stored)."""
        nan_vec = [0.1, float('nan'), 0.3]
        await adapter.store_embedding("claim", "test-id", nan_vec, "test-model")

        # Verify nothing was stored
        conn = await adapter._get_conn()
        async with conn.execute(
            "SELECT COUNT(*) FROM embedding_index WHERE item_id = ?", ("test-id",)
        ) as cur:
            row = await cur.fetchone()
        assert row[0] == 0, "NaN embedding should not be stored"

    async def test_inf_embedding_rejected(self, adapter: SQLiteAdapter):
        """Inf values in embedding should be silently skipped."""
        inf_vec = [0.1, float('inf'), 0.3]
        await adapter.store_embedding("claim", "test-id", inf_vec, "test-model")

        conn = await adapter._get_conn()
        async with conn.execute(
            "SELECT COUNT(*) FROM embedding_index WHERE item_id = ?", ("test-id",)
        ) as cur:
            row = await cur.fetchone()
        assert row[0] == 0, "Inf embedding should not be stored"

    async def test_zero_vector_rejected(self, adapter: SQLiteAdapter):
        """All-zero embeddings should be silently skipped."""
        zero_vec = [0.0, 0.0, 0.0]
        await adapter.store_embedding("claim", "test-id", zero_vec, "test-model")

        conn = await adapter._get_conn()
        async with conn.execute(
            "SELECT COUNT(*) FROM embedding_index WHERE item_id = ?", ("test-id",)
        ) as cur:
            row = await cur.fetchone()
        assert row[0] == 0, "Zero-vector embedding should not be stored"

    async def test_empty_embedding_rejected(self, adapter: SQLiteAdapter):
        """Empty embedding list should be silently skipped."""
        await adapter.store_embedding("claim", "test-id", [], "test-model")

        conn = await adapter._get_conn()
        async with conn.execute(
            "SELECT COUNT(*) FROM embedding_index WHERE item_id = ?", ("test-id",)
        ) as cur:
            row = await cur.fetchone()
        assert row[0] == 0, "Empty embedding should not be stored"

    async def test_valid_embedding_stored(self, adapter: SQLiteAdapter):
        """Normal valid embeddings should still be stored correctly."""
        valid_vec = [0.1, 0.2, 0.3]
        await adapter.store_embedding("claim", "test-id", valid_vec, "test-model")

        conn = await adapter._get_conn()
        async with conn.execute(
            "SELECT COUNT(*) FROM embedding_index WHERE item_id = ?", ("test-id",)
        ) as cur:
            row = await cur.fetchone()
        assert row[0] == 1, "Valid embedding should be stored"


# ---------------------------------------------------------------------------
# Fix 3: Brute-force vector search cap
# ---------------------------------------------------------------------------

class TestBruteForceSearchCap:
    """Ensure brute-force search respects the embedding count cap."""

    async def test_brute_force_cap_returns_empty_when_exceeded(self, adapter: SQLiteAdapter):
        """When embeddings exceed MAX_BRUTE_FORCE_EMBEDDINGS, return empty list."""
        # Set a very low cap for testing
        original_cap = adapter.MAX_BRUTE_FORCE_EMBEDDINGS
        adapter.MAX_BRUTE_FORCE_EMBEDDINGS = 5

        try:
            # Insert 10 embeddings (exceeding cap of 5)
            for i in range(10):
                vec = [float(i) * 0.1 + 0.01, 0.5, 0.3]  # non-zero valid vectors
                await adapter.store_embedding("claim", f"item-{i}", vec, "test-model")

            conn = await adapter._get_conn()
            query_vec = [0.1, 0.5, 0.3]
            results = await adapter._vec_search_brute_force(
                conn, query_vec, None, 5
            )
            assert results == [], "Should return empty when count exceeds cap"
        finally:
            adapter.MAX_BRUTE_FORCE_EMBEDDINGS = original_cap

    async def test_brute_force_works_under_cap(self, adapter: SQLiteAdapter):
        """When embeddings are under cap, brute-force search works normally."""
        original_cap = adapter.MAX_BRUTE_FORCE_EMBEDDINGS
        adapter.MAX_BRUTE_FORCE_EMBEDDINGS = 100

        try:
            # Insert 3 embeddings
            for i in range(3):
                vec = [float(i) * 0.1 + 0.01, 0.5, 0.3]
                await adapter.store_embedding("claim", f"item-{i}", vec, "test-model")

            conn = await adapter._get_conn()
            query_vec = [0.01, 0.5, 0.3]
            results = await adapter._vec_search_brute_force(
                conn, query_vec, None, 5
            )
            assert len(results) == 3, f"Expected 3 results, got {len(results)}"
        finally:
            adapter.MAX_BRUTE_FORCE_EMBEDDINGS = original_cap


# ---------------------------------------------------------------------------
# Fix 4: Sessions endpoint (list_sessions uses correct column)
# ---------------------------------------------------------------------------

class TestSessionsEndpoint:
    """Ensure list_sessions works without 'no such column' error."""

    async def test_list_sessions_returns_data(self, adapter: SQLiteAdapter):
        """list_sessions should not raise 'no such column: c.session_id'."""
        # Create a session first
        session = await adapter.create_session("test-session-1", owner_id="default")
        assert session is not None

        # Create a claim linked to that session
        claim = await adapter.create_claim(
            subject="test",
            predicate="is",
            object_value="a test",
            owner_id="default",
            source_session=session.id,  # Use the PK, not session_id
        )

        # This was the failing query before the fix
        sessions = await adapter.list_sessions(owner_id="default")
        assert len(sessions) >= 1
        # Check that claims_created_count is present and correct
        found = False
        for s in sessions:
            if s.get("session_id") == "test-session-1":
                found = True
                assert s.get("claims_created_count", 0) >= 1
                break
        assert found, "Session test-session-1 should appear in list_sessions"

    async def test_list_sessions_empty(self, adapter: SQLiteAdapter):
        """list_sessions should work even with no sessions."""
        sessions = await adapter.list_sessions(owner_id="default")
        assert sessions == []


# ---------------------------------------------------------------------------
# Fix 5: SQLite integrity check on startup
# ---------------------------------------------------------------------------

class TestIntegrityCheck:
    """Ensure integrity check runs on initialize."""

    async def test_integrity_check_runs_on_startup(self, caplog):
        """Initialize should run PRAGMA quick_check."""
        with caplog.at_level(logging.DEBUG, logger="cortex.storage.sqlite_adapter"):
            a = SQLiteAdapter(db_path=":memory:", busy_timeout_ms=1000, enable_wal=False)
            await a.initialize()
            await a.close()

        # Check that the integrity check log message appeared
        integrity_msgs = [r for r in caplog.records if "integrity check" in r.message.lower()]
        assert len(integrity_msgs) > 0, "Integrity check should log a message on startup"
        # For a fresh in-memory DB it should pass
        passed_msgs = [r for r in integrity_msgs if "passed" in r.message.lower()]
        assert len(passed_msgs) > 0, "Fresh DB should pass integrity check"


# ---------------------------------------------------------------------------
# Fix 6: get_active_session_by_id
# ---------------------------------------------------------------------------

class TestGetActiveSessionById:
    """Ensure get_active_session_by_id works correctly."""

    async def test_returns_active_session(self, adapter: SQLiteAdapter):
        """Should return the session when it exists and is awake."""
        session = await adapter.create_session("my-session-42", owner_id="default")
        assert session is not None
        assert session.session_id == "my-session-42"

        found = await adapter.get_active_session_by_id("my-session-42")
        assert found is not None
        assert found.session_id == "my-session-42"
        assert found.status == "awake"

    async def test_returns_none_for_nonexistent(self, adapter: SQLiteAdapter):
        """Should return None when session_id doesn't exist."""
        found = await adapter.get_active_session_by_id("nonexistent-session")
        assert found is None

    async def test_returns_none_for_closed_session(self, adapter: SQLiteAdapter):
        """Should return None when session exists but is closed (not awake)."""
        session = await adapter.create_session("closed-session", owner_id="default")
        await adapter.close_session("closed-session")

        found = await adapter.get_active_session_by_id("closed-session")
        assert found is None, "Closed sessions should not be returned"


# ---------------------------------------------------------------------------
# Fix 1: Write lock coverage (smoke test)
# ---------------------------------------------------------------------------

class TestWriteLockCoverage:
    """Verify write methods acquire the write lock."""

    async def test_write_lock_exists(self, adapter: SQLiteAdapter):
        """SQLiteAdapter should have a _write_lock attribute."""
        assert hasattr(adapter, "_write_lock")
        assert isinstance(adapter._write_lock, asyncio.Lock)

    async def test_concurrent_writes_dont_corrupt(self, adapter: SQLiteAdapter):
        """Run concurrent write operations to verify lock prevents corruption."""
        async def create_claims(n: int):
            for i in range(n):
                await adapter.create_claim(
                    subject=f"subj-{n}-{i}",
                    predicate="is",
                    object_value=f"obj-{n}-{i}",
                    owner_id="default",
                )

        # Run multiple concurrent writers
        await asyncio.gather(
            create_claims(5),
            create_claims(5),
            create_claims(5),
        )

        # All 15 claims should exist
        count = await adapter.count_claims(owner_id="default")
        assert count == 15, f"Expected 15 claims from concurrent writes, got {count}"
