"""
Tests for thread engine management in encrypted database.

These tests verify:
1. Thread engine tracking and reuse
2. Engine cleanup by username and thread ID
3. Connection leak prevention with multiple registrations
4. Thread isolation (different threads get different engines)
"""

import shutil
import tempfile
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

import pytest
from sqlalchemy import text

from local_deep_research.database.auth_db import (
    get_auth_db_session,
    init_auth_database,
)
from local_deep_research.database.encrypted_db import DatabaseManager
from local_deep_research.database.models.auth import User


@pytest.fixture
def temp_data_dir():
    """Create a temporary data directory for testing."""
    temp_dir = tempfile.mkdtemp()
    yield Path(temp_dir)
    shutil.rmtree(temp_dir)


@pytest.fixture
def db_manager(temp_data_dir, monkeypatch):
    """Create a DatabaseManager with test configuration."""
    monkeypatch.setenv("LDR_DATA_DIR", str(temp_data_dir))
    manager = DatabaseManager()
    manager.data_dir = temp_data_dir / "encrypted_databases"
    manager.data_dir.mkdir(parents=True, exist_ok=True)
    yield manager
    # Cleanup all engines after test
    manager.cleanup_all_thread_engines()
    for username in list(manager.connections.keys()):
        manager.close_user_database(username)


@pytest.fixture
def auth_user(temp_data_dir, monkeypatch):
    """Create a test user in auth database."""
    monkeypatch.setenv("LDR_DATA_DIR", str(temp_data_dir))
    init_auth_database()
    auth_db = get_auth_db_session()
    user = User(username="testuser")
    auth_db.add(user)
    auth_db.commit()
    auth_db.close()
    return user


@pytest.fixture
def db_manager_with_user(db_manager, auth_user):
    """Create DB manager with established user database."""
    username = "testuser"
    password = "TestPassword123!"
    db_manager.create_user_database(username, password)
    yield db_manager, username, password
    db_manager.close_user_database(username)


class TestThreadEngineTracking:
    """Test thread engine tracking and key management."""

    def test_thread_engines_dict_initialized(self, db_manager):
        """Thread engines dict should be initialized empty."""
        assert hasattr(db_manager, "_thread_engines")
        assert isinstance(db_manager._thread_engines, dict)
        assert len(db_manager._thread_engines) == 0

    def test_thread_engine_lock_exists(self, db_manager):
        """Thread engine lock should exist for thread safety."""
        assert hasattr(db_manager, "_thread_engine_lock")
        assert isinstance(
            db_manager._thread_engine_lock, type(threading.Lock())
        )

    def test_engine_key_format(self, db_manager_with_user):
        """Engine key should be a (username, thread_id) tuple."""
        db_manager, username, password = db_manager_with_user

        # Create a thread-safe session
        session = db_manager.create_thread_safe_session_for_metrics(
            username, password
        )
        session.close()

        # Check that engine was stored with correct key format (tuple)
        assert len(db_manager._thread_engines) == 1
        key = list(db_manager._thread_engines.keys())[0]
        assert isinstance(key, tuple)
        assert len(key) == 2
        assert key == (username, threading.get_ident())

    def test_engine_stored_after_creation(self, db_manager_with_user):
        """Engine should be stored in _thread_engines after creation."""
        db_manager, username, password = db_manager_with_user

        assert len(db_manager._thread_engines) == 0

        session = db_manager.create_thread_safe_session_for_metrics(
            username, password
        )
        session.close()

        assert len(db_manager._thread_engines) == 1


class TestEngineReuseWithinThread:
    """Test that same thread reuses engine instead of creating new ones."""

    def test_same_thread_reuses_engine(self, db_manager_with_user):
        """Multiple calls from same thread should reuse the same engine."""
        db_manager, username, password = db_manager_with_user

        # First call creates engine
        session1 = db_manager.create_thread_safe_session_for_metrics(
            username, password
        )
        assert len(db_manager._thread_engines) == 1
        engine_key = list(db_manager._thread_engines.keys())[0]
        engine1 = db_manager._thread_engines[engine_key]
        session1.close()

        # Second call should reuse same engine
        session2 = db_manager.create_thread_safe_session_for_metrics(
            username, password
        )
        assert len(db_manager._thread_engines) == 1  # Still only 1 engine
        engine2 = db_manager._thread_engines[engine_key]
        assert engine1 is engine2  # Same engine object
        session2.close()

        # Third call should also reuse
        session3 = db_manager.create_thread_safe_session_for_metrics(
            username, password
        )
        assert len(db_manager._thread_engines) == 1
        session3.close()

    def test_engine_validity_check(self, db_manager_with_user):
        """Engine should be validated before reuse."""
        db_manager, username, password = db_manager_with_user

        # Create initial session
        session1 = db_manager.create_thread_safe_session_for_metrics(
            username, password
        )
        session1.close()

        # Engine should still be valid and reused
        session2 = db_manager.create_thread_safe_session_for_metrics(
            username, password
        )

        # Verify session works
        result = session2.execute(text("SELECT 1"))
        assert result.scalar() == 1
        session2.close()


class TestThreadIsolation:
    """Test that different threads get different engines."""

    def test_different_threads_different_engines(self, db_manager_with_user):
        """Different threads should have different engines."""
        db_manager, username, password = db_manager_with_user
        results = {"thread_ids": [], "engine_keys": []}
        errors = []

        def worker():
            try:
                thread_id = threading.get_ident()
                session = db_manager.create_thread_safe_session_for_metrics(
                    username, password
                )
                # Verify session works
                result = session.execute(text("SELECT 1"))
                assert result.scalar() == 1
                session.close()

                results["thread_ids"].append(thread_id)
                expected_key = (username, thread_id)
                results["engine_keys"].append(expected_key)
            except Exception as e:
                errors.append(str(e))

        # Run in multiple threads
        threads = []
        for _ in range(3):
            t = threading.Thread(target=worker)
            threads.append(t)
            t.start()

        for t in threads:
            t.join(timeout=10)

        assert len(errors) == 0, f"Errors: {errors}"
        assert len(results["thread_ids"]) == 3

        # Each thread should have created its own engine
        assert len(set(results["thread_ids"])) == 3  # All unique thread IDs

        # All engine keys should exist
        for key in results["engine_keys"]:
            assert key in db_manager._thread_engines

    def test_concurrent_engine_creation(self, db_manager_with_user):
        """Multiple threads creating engines concurrently should not conflict."""
        db_manager, username, password = db_manager_with_user
        num_threads = 5
        results = []
        errors = []
        barrier = threading.Barrier(num_threads)

        def worker(worker_id):
            try:
                barrier.wait(timeout=5)  # Synchronize start
                session = db_manager.create_thread_safe_session_for_metrics(
                    username, password
                )
                result = session.execute(text("SELECT 1"))
                results.append((worker_id, result.scalar()))
                session.close()
            except Exception as e:
                errors.append((worker_id, str(e)))

        with ThreadPoolExecutor(max_workers=num_threads) as executor:
            futures = [executor.submit(worker, i) for i in range(num_threads)]
            for future in as_completed(futures, timeout=30):
                future.result()  # Raise any exceptions

        assert len(errors) == 0, f"Errors: {errors}"
        assert len(results) == num_threads
        assert all(r[1] == 1 for r in results)


class TestCleanupThreadEngines:
    """Test cleanup_thread_engines() method."""

    def test_cleanup_by_username_only(self, db_manager_with_user):
        """cleanup_thread_engines(username=x) should clean all engines for user."""
        db_manager, username, password = db_manager_with_user
        engine_keys = []
        barrier = threading.Barrier(3)

        def create_engine_in_thread():
            barrier.wait(timeout=5)  # Synchronize start
            session = db_manager.create_thread_safe_session_for_metrics(
                username, password
            )
            key = (username, threading.get_ident())
            engine_keys.append(key)
            session.close()

        # Create engines in multiple threads
        threads = []
        for _ in range(3):
            t = threading.Thread(target=create_engine_in_thread)
            threads.append(t)
            t.start()

        for t in threads:
            t.join(timeout=10)

        # Verify engines were created (at least 2 - timing can cause variations)
        assert len(db_manager._thread_engines) >= 2

        # Cleanup by username
        db_manager.cleanup_thread_engines(username=username)

        # All engines for this user should be removed
        remaining_keys = list(db_manager._thread_engines.keys())
        for key in remaining_keys:
            assert key[0] != username  # First element of tuple is username

    def test_cleanup_by_thread_id_only(self, db_manager_with_user):
        """cleanup_thread_engines(thread_id=x) should clean engines for thread."""
        db_manager, username, password = db_manager_with_user

        # Create engine in main thread
        session = db_manager.create_thread_safe_session_for_metrics(
            username, password
        )
        main_thread_id = threading.get_ident()
        main_key = (username, main_thread_id)
        session.close()

        assert main_key in db_manager._thread_engines

        # Cleanup by thread ID
        db_manager.cleanup_thread_engines(thread_id=main_thread_id)

        # Engine for this thread should be removed
        assert main_key not in db_manager._thread_engines

    def test_cleanup_by_both_username_and_thread_id(self, db_manager_with_user):
        """cleanup_thread_engines(username, thread_id) should clean specific engine."""
        db_manager, username, password = db_manager_with_user

        # Create engine
        session = db_manager.create_thread_safe_session_for_metrics(
            username, password
        )
        thread_id = threading.get_ident()
        key = (username, thread_id)
        session.close()

        assert key in db_manager._thread_engines

        # Cleanup with both parameters
        db_manager.cleanup_thread_engines(
            username=username, thread_id=thread_id
        )

        assert key not in db_manager._thread_engines

    def test_cleanup_current_thread_default(self, db_manager_with_user):
        """cleanup_thread_engines() with no args should clean current thread."""
        db_manager, username, password = db_manager_with_user

        # Create engine
        session = db_manager.create_thread_safe_session_for_metrics(
            username, password
        )
        key = (username, threading.get_ident())
        session.close()

        assert key in db_manager._thread_engines

        # Cleanup with no args (defaults to current thread)
        db_manager.cleanup_thread_engines()

        assert key not in db_manager._thread_engines


class TestCleanupAllThreadEngines:
    """Test cleanup_all_thread_engines() method."""

    def test_cleanup_all_disposes_all_engines(self, db_manager_with_user):
        """cleanup_all_thread_engines() should dispose and clear all engines."""
        db_manager, username, password = db_manager_with_user
        barrier = threading.Barrier(4)  # 3 worker threads + main thread

        def create_engine_in_thread():
            barrier.wait(timeout=5)  # Synchronize start
            session = db_manager.create_thread_safe_session_for_metrics(
                username, password
            )
            session.close()

        # Create engines in multiple threads
        threads = []
        for _ in range(3):
            t = threading.Thread(target=create_engine_in_thread)
            threads.append(t)
            t.start()

        # Main thread also participates in barrier
        barrier.wait(timeout=5)
        session = db_manager.create_thread_safe_session_for_metrics(
            username, password
        )
        session.close()

        for t in threads:
            t.join(timeout=10)

        # Verify engines were created (at least 3 - timing can cause variations)
        assert len(db_manager._thread_engines) >= 3

        # Cleanup all
        db_manager.cleanup_all_thread_engines()

        # All engines should be removed
        assert len(db_manager._thread_engines) == 0

    def test_cleanup_all_handles_errors(self, db_manager_with_user):
        """cleanup_all_thread_engines() should handle disposal errors gracefully."""
        db_manager, username, password = db_manager_with_user

        # Create engine
        session = db_manager.create_thread_safe_session_for_metrics(
            username, password
        )
        session.close()

        # Should not raise even if disposal fails
        db_manager.cleanup_all_thread_engines()
        assert len(db_manager._thread_engines) == 0


class TestCloseUserDatabaseCleansThreadEngines:
    """Test that close_user_database also cleans up thread engines."""

    def test_close_user_database_cleans_thread_engines(
        self, db_manager_with_user
    ):
        """close_user_database should also cleanup thread engines for that user."""
        db_manager, username, password = db_manager_with_user

        # Create thread engine
        session = db_manager.create_thread_safe_session_for_metrics(
            username, password
        )
        key = (username, threading.get_ident())
        session.close()

        assert key in db_manager._thread_engines

        # Close user database
        db_manager.close_user_database(username)

        # Thread engine for this user should also be cleaned
        assert key not in db_manager._thread_engines


class TestConnectionLeakPrevention:
    """Test that connection leak is prevented with multiple operations."""

    def test_multiple_session_calls_no_engine_leak(self, db_manager_with_user):
        """Multiple session calls should not create multiple engines per thread."""
        db_manager, username, password = db_manager_with_user

        # Make many session calls
        for i in range(10):
            session = db_manager.create_thread_safe_session_for_metrics(
                username, password
            )
            result = session.execute(text("SELECT 1"))
            assert result.scalar() == 1
            session.close()

        # Should still only have 1 engine for this thread
        assert len(db_manager._thread_engines) == 1

    def test_memory_usage_includes_thread_engines(self, db_manager_with_user):
        """get_memory_usage should include thread engines count."""
        db_manager, username, password = db_manager_with_user

        # Create thread engine
        session = db_manager.create_thread_safe_session_for_metrics(
            username, password
        )
        session.close()

        stats = db_manager.get_memory_usage()
        assert "thread_engines" in stats
        assert stats["thread_engines"] >= 1


class TestNonexistentDatabaseHandling:
    """Test error handling for non-existent databases."""

    def test_create_thread_safe_session_nonexistent_db(self, db_manager):
        """Should raise ValueError for non-existent database."""
        with pytest.raises(ValueError, match="No database found"):
            db_manager.create_thread_safe_session_for_metrics(
                "nonexistent_user", "password"
            )


class TestStaleEngineHandling:
    """Test handling of stale/invalid engines."""

    def test_stale_engine_recreated(self, db_manager_with_user):
        """Stale engines should be detected and recreated."""
        db_manager, username, password = db_manager_with_user

        # Create initial session
        session1 = db_manager.create_thread_safe_session_for_metrics(
            username, password
        )
        session1.close()

        key = (username, threading.get_ident())
        original_engine = db_manager._thread_engines[key]
        original_engine_id = id(original_engine)

        # Mock the engine's connect to raise an exception (simulating stale engine)
        def failing_connect(*args, **kwargs):
            raise Exception("Simulated stale engine connection failure")

        original_engine.connect = failing_connect

        # Next call should detect staleness and recreate
        session2 = db_manager.create_thread_safe_session_for_metrics(
            username, password
        )

        # Should work with new engine
        result = session2.execute(text("SELECT 1"))
        assert result.scalar() == 1
        session2.close()

        # Engine should have been replaced (different object ID)
        new_engine = db_manager._thread_engines[key]
        assert id(new_engine) != original_engine_id


class TestMultipleUsersThreadEngines:
    """Test thread engines with multiple users."""

    def test_different_users_different_engines(
        self, db_manager, temp_data_dir, monkeypatch
    ):
        """Different users should have separate engines even in same thread."""
        monkeypatch.setenv("LDR_DATA_DIR", str(temp_data_dir))
        init_auth_database()

        # Create two users
        auth_db = get_auth_db_session()
        user1 = User(username="user1")
        user2 = User(username="user2")
        auth_db.add(user1)
        auth_db.add(user2)
        auth_db.commit()
        auth_db.close()

        # Create databases for both users
        db_manager.create_user_database("user1", "pass1")
        db_manager.create_user_database("user2", "pass2")

        # Create thread engines for both
        session1 = db_manager.create_thread_safe_session_for_metrics(
            "user1", "pass1"
        )
        session2 = db_manager.create_thread_safe_session_for_metrics(
            "user2", "pass2"
        )

        session1.close()
        session2.close()

        # Should have 2 engines (one per user)
        assert len(db_manager._thread_engines) == 2

        thread_id = threading.get_ident()
        assert ("user1", thread_id) in db_manager._thread_engines
        assert ("user2", thread_id) in db_manager._thread_engines

        # Cleanup
        db_manager.close_user_database("user1")
        db_manager.close_user_database("user2")


class TestConnectionPoolExhaustionRegression:
    """
    Regression test for the connection pool exhaustion bug.

    Original bug: create_thread_safe_session_for_metrics() created a new
    orphaned SQLAlchemy engine on every call. After ~4 user registrations
    (4 registrations × 3 engines each = 12+ leaked connections), the
    connection pool was exhausted, causing critical-ui-tests failures.

    Fix: Track engines in _thread_engines dict and reuse them per thread.
    """

    def test_multiple_registrations_no_pool_exhaustion(
        self, db_manager, temp_data_dir, monkeypatch
    ):
        """
        Simulate multiple user registrations creating metric sessions.

        Before the fix, each registration would leak ~3 engines.
        After the fix, engines are reused per thread.
        """
        monkeypatch.setenv("LDR_DATA_DIR", str(temp_data_dir))
        init_auth_database()

        num_users = 6  # More than the ~4 that would exhaust the pool before
        users_created = []

        for i in range(num_users):
            username = f"reguser{i}"
            password = f"Password{i}!"

            # Create auth user
            auth_db = get_auth_db_session()
            user = User(username=username)
            auth_db.add(user)
            auth_db.commit()
            auth_db.close()

            # Create user database
            db_manager.create_user_database(username, password)
            users_created.append((username, password))

            # Simulate multiple metric session calls per registration
            # (mimics what happens during user registration flow)
            for _ in range(3):
                session = db_manager.create_thread_safe_session_for_metrics(
                    username, password
                )
                result = session.execute(text("SELECT 1"))
                assert result.scalar() == 1
                session.close()

        # Key assertion: We should have exactly num_users engines
        # (one per user in this thread), NOT num_users * 3 engines
        thread_id = threading.get_ident()
        user_engines_in_thread = [
            key for key in db_manager._thread_engines if key[1] == thread_id
        ]
        assert len(user_engines_in_thread) == num_users

        # Total engines should be bounded, not growing unbounded
        assert len(db_manager._thread_engines) == num_users

        # Cleanup
        for username, _ in users_created:
            db_manager.close_user_database(username)

        # After cleanup, no thread engines should remain
        assert len(db_manager._thread_engines) == 0

    def test_concurrent_registrations_bounded_engines(
        self, db_manager, temp_data_dir, monkeypatch
    ):
        """
        Simulate concurrent user registrations from multiple threads.

        Each thread should have its own engine per user, but engines
        should be reused within each thread.
        """
        monkeypatch.setenv("LDR_DATA_DIR", str(temp_data_dir))
        init_auth_database()

        num_threads = 4
        sessions_per_thread = 5
        username = "concurrent_user"
        password = "ConcurrentPass123!"

        # Create single user
        auth_db = get_auth_db_session()
        user = User(username=username)
        auth_db.add(user)
        auth_db.commit()
        auth_db.close()
        db_manager.create_user_database(username, password)

        errors = []
        thread_ids = []
        barrier = threading.Barrier(num_threads)

        def worker():
            try:
                barrier.wait(timeout=10)
                thread_ids.append(threading.get_ident())

                # Each thread makes multiple session calls
                for _ in range(sessions_per_thread):
                    session = db_manager.create_thread_safe_session_for_metrics(
                        username, password
                    )
                    result = session.execute(text("SELECT 1"))
                    assert result.scalar() == 1
                    session.close()
            except Exception as e:
                errors.append(str(e))

        threads = [threading.Thread(target=worker) for _ in range(num_threads)]
        for t in threads:
            t.start()
        for t in threads:
            t.join(timeout=30)

        assert len(errors) == 0, f"Errors: {errors}"

        # Key assertion: Should have exactly num_threads engines
        # (one per thread), NOT num_threads * sessions_per_thread
        assert len(db_manager._thread_engines) == num_threads

        # Each thread should have exactly one engine for this user
        for tid in thread_ids:
            assert (username, tid) in db_manager._thread_engines

        # Cleanup
        db_manager.close_user_database(username)
        assert len(db_manager._thread_engines) == 0


class TestThreadLocalSessionIntegration:
    """
    Integration tests for thread_local_session.py cleanup integration.

    These tests verify that the ThreadLocalSessionManager properly cleans up
    thread engines when sessions are cleaned up.

    Note: These tests patch the global db_manager in thread_local_session module
    to use our test fixture's db_manager instance.
    """

    def test_cleanup_thread_cleans_engine(
        self, db_manager, temp_data_dir, monkeypatch
    ):
        """
        cleanup_thread() should clean up both the session AND the thread engine.

        This tests the integration between thread_local_session.py and
        encrypted_db.py - when _cleanup_thread_session() is called, it should
        also call db_manager.cleanup_thread_engines().
        """
        import local_deep_research.database.thread_local_session as tls_module
        from local_deep_research.database.thread_local_session import (
            ThreadLocalSessionManager,
        )

        monkeypatch.setenv("LDR_DATA_DIR", str(temp_data_dir))
        # Patch the global db_manager in thread_local_session module
        monkeypatch.setattr(tls_module, "db_manager", db_manager)

        init_auth_database()

        # Create user
        auth_db = get_auth_db_session()
        user = User(username="integration_user")
        auth_db.add(user)
        auth_db.commit()
        auth_db.close()

        username = "integration_user"
        password = "IntegrationPass123!"
        db_manager.create_user_database(username, password)

        # Create a fresh ThreadLocalSessionManager that uses our db_manager
        session_manager = ThreadLocalSessionManager()

        # Get a session (this creates an engine in db_manager._thread_engines)
        session = session_manager.get_session(username, password)
        assert session is not None

        thread_id = threading.get_ident()
        engine_key = (username, thread_id)

        # Verify engine was created
        assert engine_key in db_manager._thread_engines

        # Cleanup the thread session
        session_manager.cleanup_thread()

        # Verify engine was also cleaned up
        assert engine_key not in db_manager._thread_engines

        # Cleanup
        db_manager.close_user_database(username)

    def test_cleanup_all_cleans_all_engines(
        self, db_manager, temp_data_dir, monkeypatch
    ):
        """
        cleanup_all() should clean up all thread engines via cleanup_all_thread_engines().
        """
        import local_deep_research.database.thread_local_session as tls_module
        from local_deep_research.database.thread_local_session import (
            ThreadLocalSessionManager,
        )

        monkeypatch.setenv("LDR_DATA_DIR", str(temp_data_dir))
        monkeypatch.setattr(tls_module, "db_manager", db_manager)

        init_auth_database()

        # Create user
        auth_db = get_auth_db_session()
        user = User(username="cleanup_all_user")
        auth_db.add(user)
        auth_db.commit()
        auth_db.close()

        username = "cleanup_all_user"
        password = "CleanupAllPass123!"
        db_manager.create_user_database(username, password)

        session_manager = ThreadLocalSessionManager()
        engines_created = []

        # Create sessions in multiple threads
        def create_session_in_thread():
            session = session_manager.get_session(username, password)
            if session:
                engines_created.append((username, threading.get_ident()))

        threads = []
        for _ in range(3):
            t = threading.Thread(target=create_session_in_thread)
            threads.append(t)
            t.start()

        for t in threads:
            t.join(timeout=10)

        # Also create one in main thread
        _session = session_manager.get_session(username, password)
        assert _session is not None
        engines_created.append((username, threading.get_ident()))

        # Verify engines were created
        assert len(db_manager._thread_engines) >= 1

        # Call cleanup_all
        session_manager.cleanup_all()

        # All thread engines should be cleaned
        assert len(db_manager._thread_engines) == 0

        # Cleanup
        db_manager.close_user_database(username)

    def test_get_metrics_session_creates_tracked_engine(
        self, db_manager, temp_data_dir, monkeypatch
    ):
        """
        get_metrics_session() should create an engine that is properly tracked.
        """
        import local_deep_research.database.thread_local_session as tls_module
        from local_deep_research.database.thread_local_session import (
            ThreadLocalSessionManager,
        )

        monkeypatch.setenv("LDR_DATA_DIR", str(temp_data_dir))
        monkeypatch.setattr(tls_module, "db_manager", db_manager)

        init_auth_database()

        # Create user
        auth_db = get_auth_db_session()
        user = User(username="metrics_user")
        auth_db.add(user)
        auth_db.commit()
        auth_db.close()

        username = "metrics_user"
        password = "MetricsPass123!"
        db_manager.create_user_database(username, password)

        # Create fresh session manager (uses patched db_manager)
        session_manager = ThreadLocalSessionManager()
        session = session_manager.get_session(username, password)
        assert session is not None

        thread_id = threading.get_ident()
        engine_key = (username, thread_id)

        # Verify engine is tracked
        assert engine_key in db_manager._thread_engines

        # Verify session works
        result = session.execute(text("SELECT 1"))
        assert result.scalar() == 1

        # Cleanup current thread
        session_manager.cleanup_thread()

        # Engine should be cleaned up
        assert engine_key not in db_manager._thread_engines

        # Cleanup
        db_manager.close_user_database(username)

    def test_thread_session_context_creates_tracked_engine(
        self, db_manager, temp_data_dir, monkeypatch
    ):
        """
        ThreadSessionContext should create an engine that can be cleaned up.
        """
        import local_deep_research.database.thread_local_session as tls_module
        from local_deep_research.database.thread_local_session import (
            ThreadLocalSessionManager,
            ThreadSessionContext,
        )

        monkeypatch.setenv("LDR_DATA_DIR", str(temp_data_dir))
        monkeypatch.setattr(tls_module, "db_manager", db_manager)
        # Also need to patch the global thread_session_manager used by ThreadSessionContext
        patched_session_manager = ThreadLocalSessionManager()
        monkeypatch.setattr(
            tls_module, "thread_session_manager", patched_session_manager
        )

        init_auth_database()

        # Create user
        auth_db = get_auth_db_session()
        user = User(username="context_user")
        auth_db.add(user)
        auth_db.commit()
        auth_db.close()

        username = "context_user"
        password = "ContextPass123!"
        db_manager.create_user_database(username, password)

        thread_id = threading.get_ident()
        engine_key = (username, thread_id)

        # Use context manager
        with ThreadSessionContext(username, password) as session:
            assert session is not None
            result = session.execute(text("SELECT 1"))
            assert result.scalar() == 1

            # Engine should be tracked while in context
            assert engine_key in db_manager._thread_engines

        # Engine still tracked after context (by design - cleanup on thread end)
        assert engine_key in db_manager._thread_engines

        # Manual cleanup via the patched session manager
        patched_session_manager.cleanup_thread()
        assert engine_key not in db_manager._thread_engines

        # Cleanup
        db_manager.close_user_database(username)
