"""Integration tests for the egress-policy PEPs at the inference call sites.

These do NOT re-test the PDP (``evaluate_llm_endpoint`` /
``evaluate_embeddings``) — that lives in tests/security/test_egress_policy.py
and test_egress_inference_gates_followup.py. Instead they drive the REAL
enforcement points:

    config.llm_config.get_llm(...)
    embeddings.embeddings_config.get_embeddings(...)

and assert the policy gate fires (PolicyDeniedError) BEFORE any model/client
construction happens. Built-in providers are auto-registered in the LLM
registry by ``discover_providers()`` (triggered on importing llm_config), so a
real ``get_llm`` call dispatches through ``get_llm_from_registry`` — we mock
ONLY that final construction step (and the embeddings ``create_embeddings``
classmethods) so what's exercised is the gate + the dispatch decision, never
the network/heavy model build.

Each allow/deny pair is constructed so the deny assertion FAILS if the gate
were removed: the construction mock returns a valid model, so a reverted gate
would let the cloud provider succeed instead of raising — and the
``assert_not_called`` on the construction mock proves the gate short-circuits
upstream of it.
"""

from __future__ import annotations

from unittest.mock import Mock, patch

import pytest
from langchain_core.embeddings import Embeddings
from langchain_core.language_models import BaseChatModel

from local_deep_research.config.llm_config import get_llm
from local_deep_research.embeddings.embeddings_config import get_embeddings
from local_deep_research.embeddings.providers.implementations.ollama import (
    OllamaEmbeddingsProvider,
)
from local_deep_research.embeddings.providers.implementations.openai import (
    OpenAIEmbeddingsProvider,
)
from local_deep_research.embeddings.providers.implementations.sentence_transformers import (
    SentenceTransformersProvider,
)
from local_deep_research.llm.llm_registry import register_llm, unregister_llm
from local_deep_research.security.egress.policy import PolicyDeniedError

# Snapshot fragments that force local-only inference, via the two distinct
# routes the call sites must both honor:
#  - PRIVATE_ONLY scope: context_from_snapshot IMPLIES require_local_*.
#  - the explicit per-knob flag with an otherwise-permissive scope.
# All include a "search.tool": a real run always has a configured primary, and
# the inference PEPs now resolve it via resolve_run_primary_engine (fail closed
# on a missing primary). The explicit scope still drives the decision — the
# primary value here is irrelevant to it.
PRIVATE_ONLY = {"policy.egress_scope": "private_only", "search.tool": "searxng"}
LLM_REQUIRE_LOCAL_FLAG = {
    "policy.egress_scope": "both",
    "llm.require_local_endpoint": True,
    "search.tool": "searxng",
}
EMB_REQUIRE_LOCAL_FLAG = {
    "policy.egress_scope": "both",
    "embeddings.require_local": True,
    "search.tool": "searxng",
}


@pytest.fixture(autouse=True)
def _no_leaked_active_context():
    """Belt-and-suspenders: these PEPs gate inside get_llm/get_embeddings
    without arming the audit hook, but if any test ever does, never let the
    armed context leak into a sibling test's socket path."""
    from local_deep_research.security.egress.audit_hook import (
        clear_active_context,
    )

    try:
        yield
    finally:
        clear_active_context()


# ---------------------------------------------------------------------------
# get_llm — cloud providers blocked under require-local (both routes)
# ---------------------------------------------------------------------------


class TestGetLlmCloudBlocked:
    @pytest.mark.parametrize("provider", ["openai", "anthropic"])
    def test_cloud_provider_denied_under_private_only(self, provider):
        """A cloud LLM under PRIVATE_ONLY must raise PolicyDeniedError before
        the registry dispatch — proven by asserting the construction mock is
        never reached."""
        snapshot = {
            **PRIVATE_ONLY,
            "llm.openai.api_key": "k",  # gitleaks:allow
            "llm.anthropic.api_key": "k",  # gitleaks:allow
        }
        with patch(
            "local_deep_research.config.llm_config.get_llm_from_registry"
        ) as mock_registry:
            mock_registry.return_value = Mock(spec=BaseChatModel)
            with pytest.raises(PolicyDeniedError) as exc:
                get_llm(
                    provider=provider,
                    model_name="some-cloud-model",
                    settings_snapshot=snapshot,
                )
            assert exc.value.decision.reason == "provider_cloud_only"
            # The gate fired upstream of model construction.
            mock_registry.assert_not_called()

    def test_cloud_provider_denied_under_explicit_require_local_flag(self):
        """The explicit llm.require_local_endpoint=true flag (scope BOTH) must
        gate cloud providers the same way the PRIVATE_ONLY scope does."""
        snapshot = {
            **LLM_REQUIRE_LOCAL_FLAG,
            "llm.openai.api_key": "k",
        }  # gitleaks:allow
        with patch(
            "local_deep_research.config.llm_config.get_llm_from_registry"
        ) as mock_registry:
            mock_registry.return_value = Mock(spec=BaseChatModel)
            with pytest.raises(PolicyDeniedError):
                get_llm(
                    provider="openai",
                    model_name="gpt-4o",
                    settings_snapshot=snapshot,
                )
            mock_registry.assert_not_called()

    def test_cloud_provider_allowed_when_no_local_requirement(self):
        """Mirror (allow side): with scope BOTH and the flag off, the SAME
        cloud provider passes the gate and reaches construction. Without this
        pairing the deny tests could pass for an unrelated reason."""
        snapshot = {
            "policy.egress_scope": "both",
            "llm.openai.api_key": "k",  # gitleaks:allow
            "search.tool": "searxng",
        }
        with patch(
            "local_deep_research.config.llm_config.get_llm_from_registry"
        ) as mock_registry:
            mock_registry.return_value = Mock(spec=BaseChatModel)
            result = get_llm(
                provider="openai",
                model_name="gpt-4o",
                settings_snapshot=snapshot,
            )
            assert result is not None
            mock_registry.assert_called_once()


# ---------------------------------------------------------------------------
# get_llm — local providers allowed under require-local
# ---------------------------------------------------------------------------


class TestGetLlmLocalAllowed:
    @pytest.mark.parametrize("provider", ["ollama", "lmstudio"])
    def test_local_default_provider_allowed_under_private_only(self, provider):
        """A localhost-default provider (no URL override) passes the gate under
        PRIVATE_ONLY and reaches construction (provider_local_default)."""
        with patch(
            "local_deep_research.config.llm_config.get_llm_from_registry"
        ) as mock_registry:
            mock_registry.return_value = Mock(spec=BaseChatModel)
            result = get_llm(
                provider=provider,
                model_name="local-model",
                settings_snapshot=PRIVATE_ONLY,
            )
            assert result is not None
            # The real registry dispatch ran (gate did not short-circuit it).
            mock_registry.assert_called_once()

    def test_local_provider_with_remote_url_denied(self):
        """The deny half of the local-provider pair: ollama pointed at a REMOTE
        url is refused even though it is normally a local-default provider —
        the gate classifies the configured endpoint, not the provider name."""
        snapshot = {
            **PRIVATE_ONLY,
            "llm.ollama.url": "https://ollama.example.com",
        }
        with patch(
            "local_deep_research.config.llm_config.get_llm_from_registry"
        ) as mock_registry:
            mock_registry.return_value = Mock(spec=BaseChatModel)
            with patch(
                "local_deep_research.security.egress.policy._classify_host",
                return_value=False,
            ):
                with pytest.raises(PolicyDeniedError) as exc:
                    get_llm(
                        provider="ollama",
                        model_name="llama3",
                        settings_snapshot=snapshot,
                    )
            assert exc.value.decision.reason == "provider_remote"
            mock_registry.assert_not_called()


# ---------------------------------------------------------------------------
# get_llm — user-registered in-process LLM is exempt (allowed)
# ---------------------------------------------------------------------------


class TestGetLlmUserRegistered:
    def test_user_registered_llm_allowed_under_private_only(self):
        """An operator-injected in-process LLM (programmatic API / plugin) has
        no endpoint to classify, so the gate allows it under PRIVATE_ONLY
        (user_registered_llm). The audit hook is the backstop for any stray
        socket it might open."""
        name = "_pep_callsite_inproc_llm"
        register_llm(name, Mock(spec=BaseChatModel))
        try:
            result = get_llm(
                provider=name,
                settings_snapshot=PRIVATE_ONLY,
            )
            assert result is not None
        finally:
            unregister_llm(name)

    def test_shadowing_registration_of_cloud_name_still_blocked(self):
        """Deny half: registering an in-process object under a BUILT-IN cloud
        name ("openai") must NOT smuggle it past the gate — the cloud-block
        fires first on the discovered name. Guards the _is_user_registered_llm
        shadowing discriminator at the real call site."""
        register_llm("openai", Mock(spec=BaseChatModel))
        try:
            with pytest.raises(PolicyDeniedError) as exc:
                get_llm(
                    provider="openai",
                    model_name="gpt-4o",
                    settings_snapshot=PRIVATE_ONLY,
                )
            assert exc.value.decision.reason == "provider_cloud_only"
        finally:
            # Restore the auto-discovered built-in factory so we don't leak a
            # mock into other tests' registry.
            from local_deep_research.llm.providers import discover_providers

            discover_providers(force_refresh=True)


# ---------------------------------------------------------------------------
# get_llm — snapshot-less fail-closed allow-list
# ---------------------------------------------------------------------------


class TestGetLlmSnapshotless:
    def test_snapshotless_cloud_provider_fails_closed(self):
        """No snapshot => the require-local toggle is unreadable, so a non-local
        provider must fail closed (no_snapshot_for_provider) rather than
        silently instantiate a cloud client."""
        with patch(
            "local_deep_research.config.llm_config.get_llm_from_registry"
        ) as mock_registry:
            mock_registry.return_value = Mock(spec=BaseChatModel)
            with pytest.raises(PolicyDeniedError) as exc:
                get_llm(
                    provider="openai",
                    model_name="gpt-4o",
                    settings_snapshot=None,
                )
            assert exc.value.decision.reason == "no_snapshot_for_provider"
            mock_registry.assert_not_called()

    def test_snapshotless_local_default_provider_allowed(self):
        """Mirror: a localhost-default provider is on the snapshot-less
        allow-list and proceeds to construction even with no snapshot."""
        with patch(
            "local_deep_research.config.llm_config.get_llm_from_registry"
        ) as mock_registry:
            mock_registry.return_value = Mock(spec=BaseChatModel)
            result = get_llm(
                provider="ollama",
                model_name="llama3",
                settings_snapshot=None,
            )
            assert result is not None
            mock_registry.assert_called_once()


# ---------------------------------------------------------------------------
# get_embeddings — cloud embedder blocked under require-local
# ---------------------------------------------------------------------------


class TestGetEmbeddingsCloudBlocked:
    def test_openai_denied_under_explicit_require_local_flag(self):
        """A cloud embedder under embeddings.require_local=true must raise
        before create_embeddings (which would ship the corpus to the cloud)."""
        with patch.object(
            OpenAIEmbeddingsProvider, "create_embeddings"
        ) as mock_create:
            with pytest.raises(PolicyDeniedError) as exc:
                get_embeddings(
                    provider="openai",
                    settings_snapshot=EMB_REQUIRE_LOCAL_FLAG,
                )
            assert exc.value.decision.reason == "provider_cloud"
            mock_create.assert_not_called()

    def test_openai_allowed_when_no_local_requirement(self):
        """Mirror: with require-local off, the same cloud embedder reaches
        construction — so the deny test isn't passing for an unrelated reason."""
        with patch.object(
            OpenAIEmbeddingsProvider,
            "create_embeddings",
            return_value=Mock(spec=Embeddings),
        ) as mock_create:
            result = get_embeddings(
                provider="openai",
                settings_snapshot={
                    "policy.egress_scope": "both",
                    "embeddings.require_local": False,
                    "search.tool": "searxng",
                },
            )
            assert result is not None
            mock_create.assert_called_once()

    def test_openai_denied_under_private_only_scope(self):
        """PRIVATE_ONLY implies require_local_embeddings even with the raw flag
        left at its default — a cloud embedder is refused via the scope route."""
        with patch.object(
            OpenAIEmbeddingsProvider, "create_embeddings"
        ) as mock_create:
            with pytest.raises(PolicyDeniedError):
                get_embeddings(
                    provider="openai",
                    settings_snapshot=PRIVATE_ONLY,
                )
            mock_create.assert_not_called()


# ---------------------------------------------------------------------------
# get_embeddings — local embedders allowed under require-local
# ---------------------------------------------------------------------------


class TestGetEmbeddingsLocalAllowed:
    def test_sentence_transformers_allowed_under_require_local(self):
        """In-process sentence_transformers passes the gate and constructs."""
        with patch.object(
            SentenceTransformersProvider,
            "create_embeddings",
            return_value=Mock(spec=Embeddings),
        ) as mock_create:
            result = get_embeddings(
                provider="sentence_transformers",
                settings_snapshot=EMB_REQUIRE_LOCAL_FLAG,
            )
            assert result is not None
            mock_create.assert_called_once()

    def test_ollama_local_default_allowed_under_require_local(self):
        """ollama embeddings (no URL override) falls back to its localhost
        default and is allowed."""
        with patch.object(
            OllamaEmbeddingsProvider,
            "create_embeddings",
            return_value=Mock(spec=Embeddings),
        ) as mock_create:
            result = get_embeddings(
                provider="ollama",
                settings_snapshot=EMB_REQUIRE_LOCAL_FLAG,
            )
            assert result is not None
            mock_create.assert_called_once()

    def test_ollama_remote_url_denied_under_require_local(self):
        """Deny half of the ollama pair: a remote ollama endpoint is refused
        even under the local-default provider name."""
        snapshot = {
            **EMB_REQUIRE_LOCAL_FLAG,
            "embeddings.ollama.url": "https://ollama.example.com",
        }
        with patch.object(
            OllamaEmbeddingsProvider, "create_embeddings"
        ) as mock_create:
            with patch(
                "local_deep_research.security.egress.policy._classify_host",
                return_value=False,
            ):
                with pytest.raises(PolicyDeniedError) as exc:
                    get_embeddings(
                        provider="ollama", settings_snapshot=snapshot
                    )
            assert exc.value.decision.reason == "provider_remote"
            mock_create.assert_not_called()


# ---------------------------------------------------------------------------
# get_embeddings — snapshot-less fail-closed allow-list
# ---------------------------------------------------------------------------


class TestGetEmbeddingsSnapshotless:
    def test_snapshotless_openai_fails_closed(self):
        """No snapshot => cloud embedder fails closed (no_snapshot_for_provider)
        rather than silently embedding the local corpus in the cloud."""
        with patch.object(
            OpenAIEmbeddingsProvider, "create_embeddings"
        ) as mock_create:
            with pytest.raises(PolicyDeniedError) as exc:
                get_embeddings(provider="openai", settings_snapshot=None)
            assert exc.value.decision.reason == "no_snapshot_for_provider"
            mock_create.assert_not_called()

    def test_snapshotless_sentence_transformers_allowed(self):
        """Mirror: an in-process local embedder is on the snapshot-less
        allow-list and constructs even with no snapshot."""
        with patch.object(
            SentenceTransformersProvider,
            "create_embeddings",
            return_value=Mock(spec=Embeddings),
        ) as mock_create:
            result = get_embeddings(
                provider="sentence_transformers", settings_snapshot=None
            )
            assert result is not None
            mock_create.assert_called_once()


# ---------------------------------------------------------------------------
# Inference-path fail-OPEN closed: a snapshot with NO primary must not default
# to the public searxng (-> PUBLIC_ONLY -> require_local off) and admit a cloud
# provider for a run whose actual posture is private. The PEPs resolve the
# primary via resolve_run_primary_engine (no default) and fail closed.
# ---------------------------------------------------------------------------


class TestInferencePepsFailClosedOnMissingPrimary:
    def test_llm_refuses_cloud_provider_when_no_primary(self):
        # A present cloud key but NO search.tool (and no explicit scope =>
        # ADAPTIVE). Pre-fix: searxng default -> PUBLIC_ONLY -> require_local
        # off -> openai ADMITTED. Post-fix: fail closed before construction.
        snapshot = {"llm.openai.api_key": "k"}  # gitleaks:allow
        with patch(
            "local_deep_research.config.llm_config.get_llm_from_registry"
        ) as mock_registry:
            mock_registry.return_value = Mock(spec=BaseChatModel)
            with pytest.raises(PolicyDeniedError) as exc:
                get_llm(
                    provider="openai",
                    model_name="gpt-4o",
                    settings_snapshot=snapshot,
                )
            assert exc.value.decision.reason == "invalid_policy_config"
            mock_registry.assert_not_called()

    def test_embeddings_refuses_cloud_provider_when_no_primary(self):
        with patch.object(
            OpenAIEmbeddingsProvider,
            "create_embeddings",
            return_value=Mock(spec=Embeddings),
        ) as mock_create:
            with pytest.raises(PolicyDeniedError) as exc:
                get_embeddings(
                    provider="openai",
                    settings_snapshot={"embeddings.require_local": False},
                )
            assert exc.value.decision.reason == "invalid_policy_config"
            mock_create.assert_not_called()


# ---------------------------------------------------------------------------
# The PR's HEADLINE behavior (not just the missing-primary guardrail): with a
# PRESENT primary under the default ADAPTIVE scope, the inference PEPs now
# resolve the SAME scope the factory does. A private primary ("library") forces
# PRIVATE_ONLY -> a cloud provider is denied (provider_cloud_only) — exactly the
# exfil scenario from the PR body (search_tool="library"). A public primary
# ("searxng") resolves PUBLIC_ONLY -> the same cloud provider is admitted. The
# "library" key classifies private without any DB/username lookup.
# ---------------------------------------------------------------------------


class TestInferencePepFollowsAdaptivePrimary:
    def test_llm_private_primary_denies_cloud_under_adaptive(self):
        snapshot = {
            "search.tool": "library",  # private primary, scope defaults ADAPTIVE
            "llm.openai.api_key": "k",  # gitleaks:allow
        }
        with patch(
            "local_deep_research.config.llm_config.get_llm_from_registry"
        ) as mock_registry:
            mock_registry.return_value = Mock(spec=BaseChatModel)
            with pytest.raises(PolicyDeniedError) as exc:
                get_llm(
                    provider="openai",
                    model_name="gpt-4o",
                    settings_snapshot=snapshot,
                )
            assert exc.value.decision.reason == "provider_cloud_only"
            mock_registry.assert_not_called()

    def test_llm_public_primary_admits_cloud_under_adaptive(self):
        # Mirror: a public primary under ADAPTIVE => PUBLIC_ONLY => admitted.
        snapshot = {
            "search.tool": "searxng",
            "llm.openai.api_key": "k",  # gitleaks:allow
        }
        with patch(
            "local_deep_research.config.llm_config.get_llm_from_registry"
        ) as mock_registry:
            mock_registry.return_value = Mock(spec=BaseChatModel)
            result = get_llm(
                provider="openai",
                model_name="gpt-4o",
                settings_snapshot=snapshot,
            )
            assert result is not None
            mock_registry.assert_called_once()

    def test_embeddings_private_primary_denies_cloud_under_adaptive(self):
        snapshot = {"search.tool": "library"}  # private primary, ADAPTIVE
        with patch.object(
            OpenAIEmbeddingsProvider,
            "create_embeddings",
            return_value=Mock(spec=Embeddings),
        ) as mock_create:
            with pytest.raises(PolicyDeniedError) as exc:
                get_embeddings(provider="openai", settings_snapshot=snapshot)
            assert exc.value.decision.reason in (
                "provider_cloud_only",
                "provider_cloud",
            )
            mock_create.assert_not_called()
