"""Tests for entity_id/entity_type forwarding (#923) and exclude_source_session (#1220)."""

from __future__ import annotations

import json
import asyncio
from typing import Any, Dict, List, Optional
from unittest.mock import AsyncMock, MagicMock, patch

import pytest


# ---------------------------------------------------------------------------
# Issue #923 — entity_id + entity_type flow through remember pipeline
# ---------------------------------------------------------------------------


class TestEntityIdForwarding:
    """Verify entity_id and entity_type are forwarded from HTTP request through
    to claim storage."""

    def test_enhanced_remember_request_accepts_entity_fields(self):
        """EnhancedRememberRequest model should accept entity_id and entity_type."""
        from cortex.api.models import EnhancedRememberRequest

        req = EnhancedRememberRequest(
            conversation=[{"role": "user", "content": "hello"}],
            entity_id="agent-eva",
            entity_type="agent",
        )
        assert req.entity_id == "agent-eva"
        assert req.entity_type == "agent"

    def test_enhanced_remember_request_entity_fields_optional(self):
        """entity_id and entity_type should be optional (None by default)."""
        from cortex.api.models import EnhancedRememberRequest

        req = EnhancedRememberRequest(
            conversation=[{"role": "user", "content": "hello"}],
        )
        assert req.entity_id is None
        assert req.entity_type is None

    @pytest.mark.asyncio
    async def test_remember_passes_entity_fields_to_run_extraction(self):
        """write_ops.remember() should forward entity_id/entity_type to run_extraction."""
        from cortex.api.plugin.write_ops import WriteOpsMixin

        plugin = MagicMock()
        plugin.config = MagicMock()
        plugin.config.owner_id = "test-owner"
        plugin.config.extraction_enabled = True
        plugin.storage = AsyncMock()
        plugin.extractor = MagicMock()
        plugin.llm = MagicMock()
        plugin.entity_resolver = None
        plugin.cornerstone_guardian = None
        plugin.reconciler = MagicMock()
        plugin.event_emitter = None
        plugin._session_id = "test-session"
        plugin._metrics = None
        plugin._extraction_semaphore = asyncio.Semaphore(1)
        plugin._zero_claim_outcomes = 0
        plugin._budget_check = lambda n: None
        plugin._budget_consume = lambda n: None

        mock_session = MagicMock()
        mock_session.id = "session-pk-1"
        plugin.storage.create_session = AsyncMock(return_value=mock_session)

        captured_kwargs: Dict[str, Any] = {}

        async def mock_run_extraction(**kwargs):
            captured_kwargs.update(kwargs)
            from cortex.capture.wake_pipeline import ExtractionOutcome
            return ExtractionOutcome(
                claims_extracted=0,
                claims_stored=0,
                entities_resolved=0,
                noise_filtered=0,
                errors=[],
                warnings=[],
                ok=True,
                updated_zero_claim_outcomes=0,
            )

        with patch("cortex.api.plugin.write_ops._run_extraction", side_effect=mock_run_extraction):
            await WriteOpsMixin.remember(
                plugin,
                conversation=[{"role": "user", "content": "Test message"}],
                session_id="test-session",
                entity_id="agent-eva",
                entity_type="agent",
            )

        assert captured_kwargs.get("entity_id") == "agent-eva"
        assert captured_kwargs.get("entity_type") == "agent"

    @pytest.mark.asyncio
    async def test_run_extraction_stamps_entity_on_claims(self):
        """run_extraction should stamp scope_entity_id and scope_entity_type on claims."""
        from cortex.capture.wake_pipeline import run_extraction

        mock_claim = MagicMock()
        mock_claim.metadata = {}
        mock_claim.sensitivity = "normal"

        mock_extraction_result = MagicMock()
        mock_extraction_result.claims = [mock_claim]

        mock_extractor = MagicMock()
        mock_extractor.extract = AsyncMock(return_value=mock_extraction_result)

        mock_reconciler = MagicMock()
        mock_reconciliation_report = MagicMock()
        mock_reconciliation_report.results = []
        mock_reconciler.reconcile = AsyncMock(return_value=mock_reconciliation_report)

        outcome = await run_extraction(
            extractor=mock_extractor,
            storage=AsyncMock(),
            llm=MagicMock(embed=AsyncMock(return_value=[0.1] * 128)),
            entity_resolver=None,
            cornerstone_guardian=None,
            reconciler=mock_reconciler,
            event_emitter=None,
            config=MagicMock(
                max_concurrent_extractions=1,
                max_claims_per_extraction=25,
                max_extraction_calls_per_batch=25,
            ),
            extraction_semaphore=asyncio.Semaphore(1),
            conversation=[{"role": "user", "content": "hello"}],
            sid="test-sid",
            db_session_pk="pk-1",
            effective_owner_id="test-owner",
            episode_event_id=None,
            budget_check=lambda n: None,
            budget_consume=lambda n: None,
            errors=[],
            warnings=[],
            zero_claim_outcomes=0,
            entity_id="agent-eva",
            entity_type="agent",
        )

        assert getattr(mock_claim, "scope_entity_id", None) == "agent-eva"
        assert getattr(mock_claim, "scope_entity_type", None) == "agent"

    @pytest.mark.asyncio
    async def test_run_extraction_no_entity_stamp_when_none(self):
        """run_extraction should NOT stamp scope_entity_id when entity_id is None."""
        from cortex.capture.wake_pipeline import run_extraction

        mock_claim = MagicMock()
        mock_claim.metadata = {}
        mock_claim.sensitivity = "normal"
        del mock_claim.scope_entity_id
        del mock_claim.scope_entity_type

        mock_extraction_result = MagicMock()
        mock_extraction_result.claims = [mock_claim]

        mock_extractor = MagicMock()
        mock_extractor.extract = AsyncMock(return_value=mock_extraction_result)

        mock_reconciler = MagicMock()
        mock_reconciliation_report = MagicMock()
        mock_reconciliation_report.results = []
        mock_reconciler.reconcile = AsyncMock(return_value=mock_reconciliation_report)

        await run_extraction(
            extractor=mock_extractor,
            storage=AsyncMock(),
            llm=MagicMock(embed=AsyncMock(return_value=[0.1] * 128)),
            entity_resolver=None,
            cornerstone_guardian=None,
            reconciler=mock_reconciler,
            event_emitter=None,
            config=MagicMock(
                max_concurrent_extractions=1,
                max_claims_per_extraction=25,
                max_extraction_calls_per_batch=25,
            ),
            extraction_semaphore=asyncio.Semaphore(1),
            conversation=[{"role": "user", "content": "hello"}],
            sid="test-sid",
            db_session_pk="pk-1",
            effective_owner_id="test-owner",
            episode_event_id=None,
            budget_check=lambda n: None,
            budget_consume=lambda n: None,
            errors=[],
            warnings=[],
            zero_claim_outcomes=0,
            entity_id=None,
            entity_type=None,
        )

        assert not hasattr(mock_claim, "scope_entity_id")
        assert not hasattr(mock_claim, "scope_entity_type")


# ---------------------------------------------------------------------------
# Issue #1220 — exclude_source_session on retrieve (echo loop prevention)
# ---------------------------------------------------------------------------


class TestExcludeSourceSession:
    """Verify exclude_source_session filters claims correctly."""

    def test_enhanced_retrieve_request_accepts_exclude_source_session(self):
        from cortex.api.models import EnhancedRetrieveRequest

        req = EnhancedRetrieveRequest(
            query="what is my name",
            exclude_source_session="import-session-123",
        )
        assert req.exclude_source_session == "import-session-123"

    def test_enhanced_retrieve_request_exclude_source_session_optional(self):
        from cortex.api.models import EnhancedRetrieveRequest

        req = EnhancedRetrieveRequest(query="what is my name")
        assert req.exclude_source_session is None

    def test_retrieval_constraints_accepts_exclude_source_session(self):
        from cortex.core.retrieval_types import RetrievalConstraints

        constraints = RetrievalConstraints(exclude_source_session="session-abc")
        assert constraints.exclude_source_session == "session-abc"

    def test_retrieval_constraints_exclude_source_session_default_none(self):
        from cortex.core.retrieval_types import RetrievalConstraints

        constraints = RetrievalConstraints()
        assert constraints.exclude_source_session is None

    def test_retrieval_pipeline_filters_excluded_session(self):
        """Items matching exclude_source_session should be filtered out."""
        from cortex.core.retrieval_types import RetrievalConstraints, RetrievedItem

        items = [
            RetrievedItem(
                item_type="claim", item_id="c1", content="claim 1",
                score=0.9, source_session_id="session-to-exclude",
            ),
            RetrievedItem(
                item_type="claim", item_id="c2", content="claim 2",
                score=0.8, source_session_id="other-session",
            ),
            RetrievedItem(
                item_type="claim", item_id="c3", content="claim 3",
                score=0.7, source_session_id=None,
            ),
        ]

        constraints = RetrievalConstraints(exclude_source_session="session-to-exclude")

        _excl = constraints.exclude_source_session
        filtered = [
            it for it in items
            if getattr(it, "source_session_id", None) != _excl
        ]

        assert len(filtered) == 2
        assert {it.item_id for it in filtered} == {"c2", "c3"}

    def test_no_filter_when_exclude_source_session_none(self):
        """When exclude_source_session is None, all items should be returned."""
        from cortex.core.retrieval_types import RetrievalConstraints, RetrievedItem

        items = [
            RetrievedItem(
                item_type="claim", item_id="c1", content="claim 1",
                score=0.9, source_session_id="any-session",
            ),
            RetrievedItem(
                item_type="claim", item_id="c2", content="claim 2",
                score=0.8, source_session_id=None,
            ),
        ]

        constraints = RetrievalConstraints()

        _excl = constraints.exclude_source_session
        if _excl:
            filtered = [
                it for it in items
                if getattr(it, "source_session_id", None) != _excl
            ]
        else:
            filtered = items

        assert len(filtered) == 2

    @pytest.mark.asyncio
    async def test_browse_memories_exclude_source_session_sql(self):
        """SQL filter for exclude_source_session should exclude matching claims."""
        import aiosqlite

        db = await aiosqlite.connect(":memory:")
        await db.execute("PRAGMA journal_mode=WAL")

        await db.execute("""
            CREATE TABLE semantic_claim (
                id TEXT PRIMARY KEY,
                subject TEXT DEFAULT '',
                predicate TEXT DEFAULT '',
                object TEXT DEFAULT '',
                owner_id TEXT DEFAULT 'default',
                source_session TEXT,
                status TEXT DEFAULT 'active',
                is_deleted INTEGER DEFAULT 0,
                valid_to TEXT,
                confidence REAL DEFAULT 0.8,
                category TEXT DEFAULT 'misc',
                access_count INTEGER DEFAULT 0,
                mentions INTEGER DEFAULT 1,
                metadata_json TEXT,
                created_at TEXT DEFAULT (datetime('now')),
                updated_at TEXT DEFAULT (datetime('now'))
            )
        """)

        await db.execute(
            "INSERT INTO semantic_claim (id, subject, predicate, object, owner_id, source_session, metadata_json) VALUES (?, ?, ?, ?, ?, ?, ?)",
            ("c1", "Andrew", "likes", "coffee", "default", "session-import", json.dumps({"source_session_id": "session-import"})),
        )
        await db.execute(
            "INSERT INTO semantic_claim (id, subject, predicate, object, owner_id, source_session, metadata_json) VALUES (?, ?, ?, ?, ?, ?, ?)",
            ("c2", "Eva", "is", "AI", "default", "session-other", None),
        )
        await db.execute(
            "INSERT INTO semantic_claim (id, subject, predicate, object, owner_id, source_session) VALUES (?, ?, ?, ?, ?, ?)",
            ("c3", "Sky", "is", "blue", "default", None),
        )
        await db.commit()

        conditions = ["sc.owner_id = ?"]
        params: list = ["default"]
        conditions.append("(sc.is_deleted = 0 OR sc.is_deleted IS NULL)")
        conditions.append("(sc.valid_to IS NULL OR sc.valid_to > datetime('now'))")

        exclude_source_session = "session-import"
        conditions.append("(sc.source_session IS NULL OR sc.source_session != ?)")
        params.append(exclude_source_session)
        conditions.append(
            "(sc.metadata_json IS NULL OR json_extract(sc.metadata_json, '$.source_session_id') IS NULL"
            " OR json_extract(sc.metadata_json, '$.source_session_id') != ?)"
        )
        params.append(exclude_source_session)

        where = " AND ".join(conditions)
        query = f"SELECT id FROM semantic_claim sc WHERE {where}"
        async with db.execute(query, params) as cur:
            rows = await cur.fetchall()

        ids = {r[0] for r in rows}
        assert "c1" not in ids
        assert "c2" in ids
        assert "c3" in ids

        await db.close()

    @pytest.mark.asyncio
    async def test_browse_no_exclusion_returns_all(self):
        """Without exclude_source_session, all claims should be returned."""
        import aiosqlite

        db = await aiosqlite.connect(":memory:")
        await db.execute("""
            CREATE TABLE semantic_claim (
                id TEXT PRIMARY KEY,
                owner_id TEXT DEFAULT 'default',
                source_session TEXT,
                is_deleted INTEGER DEFAULT 0,
                valid_to TEXT,
                metadata_json TEXT
            )
        """)

        await db.execute(
            "INSERT INTO semantic_claim (id, owner_id, source_session) VALUES (?, ?, ?)",
            ("c1", "default", "session-import"),
        )
        await db.execute(
            "INSERT INTO semantic_claim (id, owner_id, source_session) VALUES (?, ?, ?)",
            ("c2", "default", "session-other"),
        )
        await db.commit()

        conditions = ["sc.owner_id = ?"]
        params: list = ["default"]
        conditions.append("(sc.is_deleted = 0 OR sc.is_deleted IS NULL)")

        where = " AND ".join(conditions)
        query = f"SELECT id FROM semantic_claim sc WHERE {where}"
        async with db.execute(query, params) as cur:
            rows = await cur.fetchall()

        assert len(rows) == 2
        await db.close()

    def test_claims_without_source_session_not_excluded(self):
        """Claims with NULL source_session should never be excluded."""
        from cortex.core.retrieval_types import RetrievedItem

        item_no_session = RetrievedItem(
            item_type="claim", item_id="c-no-session",
            content="I have no source session", score=0.9,
            source_session_id=None,
        )

        exclude = "some-session-id"
        assert getattr(item_no_session, "source_session_id", None) != exclude
