"""Tests for owner_id isolation in retrieval paths.

Verifies that BM25/FTS, vector (sqlite-vec + brute-force), and
lightweight_retrieve properly filter results by owner_id — ensuring
multi-tenant isolation (issue #1398).
"""
from __future__ import annotations

import asyncio
from typing import List, Optional
from unittest.mock import MagicMock

import pytest

from cortex.core.retrieval_search import bm25_search, vector_search, cosine_fallback
from cortex.core.retrieval_lightweight import lightweight_retrieve
from cortex.core.retrieval_types import RetrievalConstraints


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _make_config(**overrides):
    cfg = MagicMock()
    cfg.retrieval_temporal_boost_weight = 0
    cfg.retrieval_token_budget = 3000
    cfg.retrieval_rerank_enabled = False
    cfg.retrieval_status_penalties = None
    cfg.owner_id = "default"
    for k, v in overrides.items():
        setattr(cfg, k, v)
    return cfg


# ---------------------------------------------------------------------------
# In-memory storage adapter for isolation testing
# ---------------------------------------------------------------------------

class InMemoryStorageForIsolation:
    """Minimal in-memory storage with owner-scoped claims + FTS + vector."""

    def __init__(self):
        self.claims = {}
        self.embeddings = {}
        self._vec_loaded = False

    def add_claim(self, claim_id: str, owner_id: str, content: str,
                  embedding: Optional[List[float]] = None):
        self.claims[claim_id] = {
            "id": claim_id,
            "owner_id": owner_id,
            "content": content,
            "status": "active",
            "is_deleted": 0,
            "subject": "test",
            "predicate": "test",
            "object_value": content,
            "confidence": 0.9,
            "memory_class": "semantic",
            "salience": "medium",
            "explicitness": "explicit",
            "stability": "stable",
            "category": "test",
            "created_at": "2026-01-01T00:00:00Z",
            "updated_at": "2026-01-01T00:00:00Z",
            "event_date": None,
            "valid_to": None,
            "metadata_json": None,
            "source_session": None,
        }
        if embedding:
            self.embeddings[claim_id] = embedding

    async def fts_search(self, query, table="claims", top_k=50, owner_id="default",
                         entity_id=None, entity_type=None, entity_ids=None):
        from cortex.types import FTSResult
        results = []
        tokens = [t.strip('"').lower() for t in query.split() if t.strip('"')]
        for cid, claim in self.claims.items():
            if claim["owner_id"] != owner_id:
                continue
            if claim.get("is_deleted", 0):
                continue
            content_lower = claim["content"].lower()
            if tokens and any(t in content_lower for t in tokens):
                results.append(FTSResult(
                    item_id=cid, content=claim["content"], rank=-1.0,
                ))
        return results[:top_k]

    async def vector_search(self, query_embedding, item_type=None, top_k=50,
                            owner_id="default"):
        from cortex.types import VectorResult
        results = []
        for cid, emb in self.embeddings.items():
            claim = self.claims.get(cid)
            if not claim:
                continue
            if claim["owner_id"] != owner_id:
                continue
            sim = sum(a * b for a, b in zip(query_embedding, emb))
            results.append(VectorResult(
                item_type="claim", item_id=cid, score=sim,
            ))
        results.sort(key=lambda r: r.score, reverse=True)
        return results[:top_k]

    async def get_claim(self, claim_id):
        data = self.claims.get(claim_id)
        if not data:
            return None
        obj = MagicMock()
        for k, v in data.items():
            setattr(obj, k, v)
        return obj

    async def get_claims_sensitivity_batch(self, claim_ids):
        return {cid: "normal" for cid in claim_ids}

    async def log_retrieval(self, **kwargs):
        pass

    async def get_all_embeddings(self):
        import struct
        rows = []
        for cid, emb in self.embeddings.items():
            blob = struct.pack(f"{len(emb)}f", *emb)
            rows.append({
                "item_type": "claim",
                "item_id": cid,
                "embedding": blob,
                "model_id": "test",
            })
        return rows


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------

OWNER_A = "owner-alpha"
OWNER_B = "owner-beta"
SHARED_CONTENT = "Andrew prefers Python for backend development"


@pytest.fixture
def storage():
    s = InMemoryStorageForIsolation()
    s.add_claim("claim-a1", OWNER_A, SHARED_CONTENT,
                embedding=[0.1, 0.2, 0.3, 0.4])
    s.add_claim("claim-a2", OWNER_A, "Andrew lives in Bangkok",
                embedding=[0.5, 0.6, 0.7, 0.8])
    s.add_claim("claim-b1", OWNER_B, SHARED_CONTENT,
                embedding=[0.1, 0.2, 0.3, 0.4])
    s.add_claim("claim-b2", OWNER_B, "Andrew works at 100Yen",
                embedding=[0.9, 0.1, 0.2, 0.3])
    return s


@pytest.fixture
def embed_fn():
    async def _embed(text):
        return [0.1, 0.2, 0.3, 0.4]
    return _embed


# ---------------------------------------------------------------------------
# BM25 Tests
# ---------------------------------------------------------------------------

class TestBM25OwnerIsolation:

    @pytest.mark.asyncio
    async def test_bm25_returns_only_owner_a_claims(self, storage):
        results = await bm25_search(
            storage, "Python backend", top_k=10, owner_id=OWNER_A,
        )
        ids = {r.item_id for r in results}
        assert "claim-a1" in ids
        assert "claim-b1" not in ids
        assert "claim-b2" not in ids

    @pytest.mark.asyncio
    async def test_bm25_returns_only_owner_b_claims(self, storage):
        results = await bm25_search(
            storage, "Python backend", top_k=10, owner_id=OWNER_B,
        )
        ids = {r.item_id for r in results}
        assert "claim-b1" in ids
        assert "claim-a1" not in ids

    @pytest.mark.asyncio
    async def test_bm25_unknown_owner_returns_empty(self, storage):
        results = await bm25_search(
            storage, "Python", top_k=10, owner_id="owner-unknown",
        )
        assert len(results) == 0


# ---------------------------------------------------------------------------
# Vector Search Tests
# ---------------------------------------------------------------------------

class TestVectorSearchOwnerIsolation:

    @pytest.mark.asyncio
    async def test_vector_returns_only_owner_a(self, storage, embed_fn):
        results = await vector_search(
            storage, "Python backend", top_k=10, owner_id=OWNER_A,
            embed_fn=embed_fn,
        )
        ids = {r.item_id for r in results}
        assert "claim-b1" not in ids
        assert "claim-b2" not in ids
        assert any(r.item_id.startswith("claim-a") for r in results)

    @pytest.mark.asyncio
    async def test_vector_returns_only_owner_b(self, storage, embed_fn):
        results = await vector_search(
            storage, "Python backend", top_k=10, owner_id=OWNER_B,
            embed_fn=embed_fn,
        )
        ids = {r.item_id for r in results}
        assert "claim-a1" not in ids
        assert "claim-a2" not in ids

    @pytest.mark.asyncio
    async def test_vector_unknown_owner_returns_empty(self, storage, embed_fn):
        results = await vector_search(
            storage, "Python", top_k=10, owner_id="owner-unknown",
            embed_fn=embed_fn,
        )
        assert len(results) == 0


# ---------------------------------------------------------------------------
# Cosine Fallback Tests
# ---------------------------------------------------------------------------

class TestCosineFallbackOwnerIsolation:

    @pytest.mark.asyncio
    async def test_cosine_fallback_filters_by_owner(self, storage):
        results = await cosine_fallback(
            storage, [0.1, 0.2, 0.3, 0.4], top_k=10, owner_id=OWNER_A,
        )
        ids = {r.item_id for r in results}
        assert "claim-b1" not in ids
        assert "claim-b2" not in ids

    @pytest.mark.asyncio
    async def test_cosine_fallback_unknown_owner_returns_empty(self, storage):
        results = await cosine_fallback(
            storage, [0.1, 0.2, 0.3, 0.4], top_k=10, owner_id="owner-unknown",
        )
        assert len(results) == 0

    @pytest.mark.asyncio
    async def test_cosine_fallback_no_owner_returns_all(self, storage):
        """Backward compat: no owner_id → no filtering."""
        results = await cosine_fallback(
            storage, [0.1, 0.2, 0.3, 0.4], top_k=10,
        )
        # Default owner_id="default" won't match either owner, so empty
        # But if we add claims with owner_id="default", they should appear
        storage.add_claim("claim-default", "default", "default claim",
                          embedding=[0.1, 0.2, 0.3, 0.4])
        results = await cosine_fallback(
            storage, [0.1, 0.2, 0.3, 0.4], top_k=10,
        )
        ids = {r.item_id for r in results}
        assert "claim-default" in ids


# ---------------------------------------------------------------------------
# Lightweight Retrieve (end-to-end)
# ---------------------------------------------------------------------------

class TestLightweightRetrieveOwnerIsolation:

    @pytest.mark.asyncio
    async def test_owner_a_only(self, storage, embed_fn):
        items, _info = await lightweight_retrieve(
            storage=storage, config=_make_config(),
            query="Python backend", constraints=RetrievalConstraints(top_k=10),
            owner_id=OWNER_A, rrf_k=60, candidate_multiplier=3,
            embed_fn=embed_fn,
        )
        ids = {it.item_id for it in items}
        assert "claim-b1" not in ids
        assert "claim-b2" not in ids

    @pytest.mark.asyncio
    async def test_owner_b_only(self, storage, embed_fn):
        items, _info = await lightweight_retrieve(
            storage=storage, config=_make_config(),
            query="Python backend", constraints=RetrievalConstraints(top_k=10),
            owner_id=OWNER_B, rrf_k=60, candidate_multiplier=3,
            embed_fn=embed_fn,
        )
        ids = {it.item_id for it in items}
        assert "claim-a1" not in ids
        assert "claim-a2" not in ids

    @pytest.mark.asyncio
    async def test_different_owners_no_overlap(self, storage, embed_fn):
        cfg = _make_config()
        constraints = RetrievalConstraints(top_k=10)

        items_a, _info = await lightweight_retrieve(
            storage=storage, config=cfg, query="Python",
            constraints=constraints, owner_id=OWNER_A,
            rrf_k=60, candidate_multiplier=3, embed_fn=embed_fn,
        )
        items_b, _info = await lightweight_retrieve(
            storage=storage, config=cfg, query="Python",
            constraints=constraints, owner_id=OWNER_B,
            rrf_k=60, candidate_multiplier=3, embed_fn=embed_fn,
        )

        ids_a = {it.item_id for it in items_a}
        ids_b = {it.item_id for it in items_b}
        assert not ids_a.intersection(ids_b), \
            "Owner A and B results must have zero overlap"
