"""Regression tests for thread-local leak prevention in ParallelSearchEngine.

These tests pin the executor to ``max_workers=1`` so two back-to-back
submissions are guaranteed to land on the same worker thread, which is
the scenario where leaks would be observable.
"""

import concurrent.futures
from unittest.mock import Mock, patch

from local_deep_research.database.thread_metrics import metrics_writer
from local_deep_research.utilities.thread_context import (
    get_search_context,
    set_search_context,
    clear_search_context,
)


class TestParallelSearchEngineThreadLocalIsolation:
    """Verify the submit-site wrapper in ``_get_previews`` cleans up
    thread-local state (search_context + metrics_writer passwords) so that
    the next task on the same worker starts with a clean slate."""

    def _make_pinned_executor(self):
        return concurrent.futures.ThreadPoolExecutor(
            max_workers=1, thread_name_prefix="test_pinned_"
        )

    def _make_engine(self):
        from local_deep_research.web_search_engines.engines.parallel_search_engine import (
            ParallelSearchEngine,
        )

        mock_llm = Mock()
        settings = {"search.max_results": {"value": 10}}
        return ParallelSearchEngine(
            llm=mock_llm,
            settings_snapshot=settings,
            programmatic_mode=True,
        )

    def test_worker_thread_has_no_stale_context_between_tasks(self):
        """First task contaminates the worker's thread-local state; the
        second task's pre-run snapshot must show that state cleared."""
        from local_deep_research.web_search_engines.engines import (
            parallel_search_engine as pse_mod,
        )

        engine = self._make_engine()

        # Record what the worker thread sees at task entry.
        entry_snapshots: list[dict] = []
        dirty_writes = {"userA_pw": "secret_A"}

        def fake_execute(self_engine, engine_name, query):
            entry_snapshots.append(
                {
                    "engine_name": engine_name,
                    "search_context": get_search_context(),
                    "passwords_copy": dict(
                        getattr(metrics_writer._thread_local, "passwords", {})
                    ),
                }
            )
            # Simulate a metrics call on the worker thread that writes a
            # password into thread-local storage mid-task. The submit-site
            # wrapper's finally must clear this before the next task runs.
            metrics_writer.set_user_password("userA", dirty_writes["userA_pw"])
            return {
                "engine": engine_name,
                "success": True,
                "results": [{"title": f"r-{engine_name}"}],
                "count": 1,
            }

        pinned = self._make_pinned_executor()
        try:
            with (
                patch.object(
                    pse_mod, "_get_global_executor", return_value=pinned
                ),
                patch.object(
                    engine,
                    "select_engines",
                    side_effect=[["engine_first"], ["engine_second"]],
                ),
                patch.object(
                    pse_mod.ParallelSearchEngine,
                    "_execute_single_engine",
                    new=fake_execute,
                ),
            ):
                # First "request": submitter has a research context.
                set_search_context({"username": "userA", "research_id": "r-1"})
                try:
                    engine._get_previews("q1")
                finally:
                    clear_search_context()

                # Second "request": submitter has NO context (simulating a
                # different user whose code path didn't populate one, or a
                # background call).
                assert get_search_context() is None
                engine._get_previews("q2")
        finally:
            pinned.shutdown(wait=True)

        assert len(entry_snapshots) == 2, entry_snapshots

        # First task saw the submitter's context applied by _run_with_context.
        assert entry_snapshots[0]["search_context"] == {
            "username": "userA",
            "research_id": "r-1",
        }

        # Second task entered with a CLEAN worker thread: no context and no
        # lingering password from the first task's mid-run write.
        assert entry_snapshots[1]["search_context"] is None, (
            "Worker thread leaked search_context across tasks: "
            f"{entry_snapshots[1]['search_context']!r}"
        )
        assert entry_snapshots[1]["passwords_copy"] == {}, (
            "Worker thread leaked metrics_writer passwords across tasks: "
            f"{entry_snapshots[1]['passwords_copy']!r}"
        )

    def test_class_level_preserve_research_context_is_removed(self):
        """The class-def-time @preserve_research_context wrap was a no-op
        and has been removed. Its import should also be gone."""
        import inspect

        from local_deep_research.web_search_engines.engines import (
            parallel_search_engine as pse_mod,
        )

        source = inspect.getsource(pse_mod)
        assert "preserve_research_context" not in source, (
            "preserve_research_context should no longer be referenced in "
            "parallel_search_engine.py — the class-def usage was a no-op."
        )
