from contextlib import asynccontextmanager
from datetime import datetime, timedelta, timezone
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from nanoid import generate as generate_nanoid
from sqlalchemy import func, update
from sqlalchemy.ext.asyncio import AsyncSession

from src import models
from src.crud.document import CreateDocumentsResult
from src.crud.representation import RepresentationManager
from src.schemas.configuration import (
    ResolvedConfiguration,
    ResolvedDreamConfiguration,
    ResolvedPeerCardConfiguration,
    ResolvedReasoningConfiguration,
    ResolvedSummaryConfiguration,
)
from src.utils.representation import (
    DeductiveObservation,
    ExplicitObservation,
    Representation,
)


def _resolved_config(*, dream_enabled: bool = False) -> ResolvedConfiguration:
    """Build a minimal ResolvedConfiguration for tests that only care about dream.enabled."""
    return ResolvedConfiguration(
        reasoning=ResolvedReasoningConfiguration(enabled=False),
        peer_card=ResolvedPeerCardConfiguration(use=False, create=False),
        summary=ResolvedSummaryConfiguration(
            enabled=False,
            messages_per_short_summary=20,
            messages_per_long_summary=60,
        ),
        dream=ResolvedDreamConfiguration(enabled=dream_enabled),
    )


@asynccontextmanager
async def _fake_tracked_db(_name: str):
    yield object()


def _saved_observations(mock_save: AsyncMock):
    call = mock_save.await_args
    assert call is not None, "mock_save was never awaited"
    if "all_observations" in call.kwargs:
        return call.kwargs["all_observations"]
    if len(call.args) > 1:
        return call.args[1]
    raise AssertionError("missing all_observations in await args")


class TestRepresentationManagerSoftDelete:
    """Tests that RepresentationManager query methods exclude soft-deleted documents."""

    async def _setup(
        self,
        db_session: AsyncSession,
        test_workspace: models.Workspace,
        test_peer: models.Peer,
    ) -> tuple[models.Peer, models.Session, models.Collection, RepresentationManager]:
        """Create peers, session, collection, and a RepresentationManager."""
        test_peer2 = models.Peer(
            name=str(generate_nanoid()), workspace_name=test_workspace.name
        )
        db_session.add(test_peer2)
        await db_session.flush()

        test_session = models.Session(
            name=str(generate_nanoid()), workspace_name=test_workspace.name
        )
        db_session.add(test_session)
        await db_session.flush()

        collection = models.Collection(
            workspace_name=test_workspace.name,
            observer=test_peer.name,
            observed=test_peer2.name,
        )
        db_session.add(collection)
        await db_session.flush()

        manager = RepresentationManager(
            test_workspace.name,
            observer=test_peer.name,
            observed=test_peer2.name,
        )

        return test_peer2, test_session, collection, manager

    @pytest.mark.asyncio
    async def test_query_documents_recent_excludes_soft_deleted(
        self,
        db_session: AsyncSession,
        sample_data: tuple[models.Workspace, models.Peer],
    ):
        """Soft-deleted documents must not appear in the recent-documents query."""
        test_workspace, test_peer = sample_data
        test_peer2, test_session, _, manager = await self._setup(
            db_session, test_workspace, test_peer
        )

        # Create two documents
        doc_live = models.Document(
            workspace_name=test_workspace.name,
            observer=test_peer.name,
            observed=test_peer2.name,
            content="Live observation",
            session_name=test_session.name,
        )
        doc_deleted = models.Document(
            workspace_name=test_workspace.name,
            observer=test_peer.name,
            observed=test_peer2.name,
            content="Deleted observation",
            session_name=test_session.name,
        )
        db_session.add_all([doc_live, doc_deleted])
        await db_session.flush()

        # Soft-delete one
        await db_session.execute(
            update(models.Document)
            .where(models.Document.id == doc_deleted.id)
            .values(deleted_at=func.now())
        )
        await db_session.commit()

        results = await manager._query_documents_recent(db_session, top_k=10)  # pyright: ignore[reportPrivateUsage]

        result_ids = [doc.id for doc in results]
        assert doc_live.id in result_ids
        assert doc_deleted.id not in result_ids

    @pytest.mark.asyncio
    async def test_query_documents_most_derived_excludes_soft_deleted(
        self,
        db_session: AsyncSession,
        sample_data: tuple[models.Workspace, models.Peer],
    ):
        """Soft-deleted documents must not appear in the most-derived query."""
        test_workspace, test_peer = sample_data
        test_peer2, test_session, _, manager = await self._setup(
            db_session, test_workspace, test_peer
        )

        # Create two documents with different times_derived
        doc_live = models.Document(
            workspace_name=test_workspace.name,
            observer=test_peer.name,
            observed=test_peer2.name,
            content="Live observation",
            session_name=test_session.name,
            times_derived=5,
        )
        doc_deleted = models.Document(
            workspace_name=test_workspace.name,
            observer=test_peer.name,
            observed=test_peer2.name,
            content="Deleted high-derived observation",
            session_name=test_session.name,
            times_derived=100,
        )
        db_session.add_all([doc_live, doc_deleted])
        await db_session.flush()

        # Soft-delete the high-derived one
        await db_session.execute(
            update(models.Document)
            .where(models.Document.id == doc_deleted.id)
            .values(deleted_at=func.now())
        )
        await db_session.commit()

        results = await manager._query_documents_most_derived(db_session, top_k=10)  # pyright: ignore[reportPrivateUsage]

        result_ids = [doc.id for doc in results]
        assert doc_live.id in result_ids
        assert doc_deleted.id not in result_ids

    @pytest.mark.asyncio
    async def test_query_documents_most_derived_ties_break_by_recency(
        self,
        db_session: AsyncSession,
        sample_data: tuple[models.Workspace, models.Peer],
    ):
        """Regression: when times_derived ties, the manager's most-derived query
        must fall back to recency, not insertion order. Mirrors the equivalent
        test on crud.query_documents_most_derived -- the query is duplicated in
        both modules and must not drift."""
        test_workspace, test_peer = sample_data
        test_peer2, test_session, _, manager = await self._setup(
            db_session, test_workspace, test_peer
        )

        base = datetime(2026, 1, 1, tzinfo=timezone.utc)
        # Three conclusions, all reinforced once, inserted oldest-first.
        for i in range(3):
            db_session.add(
                models.Document(
                    workspace_name=test_workspace.name,
                    observer=test_peer.name,
                    observed=test_peer2.name,
                    content=f"tie {i}",
                    session_name=test_session.name,
                    times_derived=1,
                    created_at=base + timedelta(days=i),
                )
            )
        # A genuinely reinforced conclusion that is also the oldest of all.
        db_session.add(
            models.Document(
                workspace_name=test_workspace.name,
                observer=test_peer.name,
                observed=test_peer2.name,
                content="hot",
                session_name=test_session.name,
                times_derived=5,
                created_at=base - timedelta(days=10),
            )
        )
        await db_session.flush()

        results = await manager._query_documents_most_derived(db_session, top_k=10)  # pyright: ignore[reportPrivateUsage]

        contents = [doc.content for doc in results]
        # Primary sort still wins: the actually-reinforced conclusion leads.
        assert contents[0] == "hot"
        # Ties break toward most-recent, not oldest-inserted.
        assert contents[1:] == ["tie 2", "tie 1", "tie 0"]


class TestRepresentationManagerSessionScoping:
    """Tests that the session allowlist is applied uniformly to every query path.

    Regression: session_name used to be applied only to the
    recent-documents query; the semantic and most-derived paths ignored it,
    so limit_to_session leaked cross-session conclusions.
    """

    async def _setup(
        self,
        db_session: AsyncSession,
        test_workspace: models.Workspace,
        test_peer: models.Peer,
    ) -> tuple[models.Session, models.Session, RepresentationManager]:
        """Create two sessions and documents in each, plus a session-less doc."""
        test_peer2 = models.Peer(
            name=str(generate_nanoid()), workspace_name=test_workspace.name
        )
        db_session.add(test_peer2)
        await db_session.flush()

        session_a = models.Session(
            name=str(generate_nanoid()), workspace_name=test_workspace.name
        )
        session_b = models.Session(
            name=str(generate_nanoid()), workspace_name=test_workspace.name
        )
        db_session.add_all([session_a, session_b])
        await db_session.flush()

        collection = models.Collection(
            workspace_name=test_workspace.name,
            observer=test_peer.name,
            observed=test_peer2.name,
        )
        db_session.add(collection)
        await db_session.flush()

        db_session.add_all(
            [
                models.Document(
                    workspace_name=test_workspace.name,
                    observer=test_peer.name,
                    observed=test_peer2.name,
                    content="in-scope observation",
                    session_name=session_a.name,
                    times_derived=1,
                ),
                models.Document(
                    workspace_name=test_workspace.name,
                    observer=test_peer.name,
                    observed=test_peer2.name,
                    content="out-of-scope observation",
                    session_name=session_b.name,
                    times_derived=100,
                ),
                # Dream-produced documents have no session_name; a session
                # allowlist must exclude them (fail-closed).
                models.Document(
                    workspace_name=test_workspace.name,
                    observer=test_peer.name,
                    observed=test_peer2.name,
                    content="sessionless dream observation",
                    session_name=None,
                    times_derived=50,
                ),
            ]
        )
        await db_session.flush()

        manager = RepresentationManager(
            test_workspace.name,
            observer=test_peer.name,
            observed=test_peer2.name,
        )
        return session_a, session_b, manager

    @pytest.mark.asyncio
    async def test_recent_respects_session_allowlist(
        self,
        db_session: AsyncSession,
        sample_data: tuple[models.Workspace, models.Peer],
    ):
        test_workspace, test_peer = sample_data
        session_a, _, manager = await self._setup(db_session, test_workspace, test_peer)

        results = await manager._query_documents_recent(  # pyright: ignore[reportPrivateUsage]
            db_session, top_k=10, session_allowlist=[session_a.name]
        )

        contents = [doc.content for doc in results]
        assert contents == ["in-scope observation"]

    @pytest.mark.asyncio
    async def test_most_derived_respects_session_allowlist(
        self,
        db_session: AsyncSession,
        sample_data: tuple[models.Workspace, models.Peer],
    ):
        """The out-of-scope doc has far higher times_derived; it must still be excluded."""
        test_workspace, test_peer = sample_data
        session_a, _, manager = await self._setup(db_session, test_workspace, test_peer)

        results = await manager._query_documents_most_derived(  # pyright: ignore[reportPrivateUsage]
            db_session, top_k=10, session_allowlist=[session_a.name]
        )

        contents = [doc.content for doc in results]
        assert contents == ["in-scope observation"]

    @pytest.mark.asyncio
    async def test_semantic_passes_session_allowlist_as_filters(
        self,
        db_session: AsyncSession,
        sample_data: tuple[models.Workspace, models.Peer],
    ):
        """The semantic path must push the allowlist down to query_documents."""
        test_workspace, test_peer = sample_data
        session_a, _, manager = await self._setup(db_session, test_workspace, test_peer)

        with patch(
            "src.crud.query_documents", new=AsyncMock(return_value=[])
        ) as mock_query:
            await manager._query_documents_semantic(  # pyright: ignore[reportPrivateUsage]
                db_session,
                query="anything",
                top_k=5,
                embedding=[0.1],
                session_allowlist=[session_a.name],
            )

        assert mock_query.await_args is not None
        assert mock_query.await_args.kwargs["filters"] == {
            "session_name": {"in": [session_a.name]},
            # Scoped recall serves only levels with a trustworthy session
            # stamp (ALLOWLIST_SAFE_LEVELS).
            "level": {"in": ["explicit"]},
        }

    @pytest.mark.asyncio
    async def test_semantic_passes_no_filters_when_unscoped(
        self,
        db_session: AsyncSession,
        sample_data: tuple[models.Workspace, models.Peer],
    ):
        test_workspace, test_peer = sample_data
        _, _, manager = await self._setup(db_session, test_workspace, test_peer)

        with patch(
            "src.crud.query_documents", new=AsyncMock(return_value=[])
        ) as mock_query:
            await manager._query_documents_semantic(  # pyright: ignore[reportPrivateUsage]
                db_session,
                query="anything",
                top_k=5,
                embedding=[0.1],
            )

        assert mock_query.await_args is not None
        assert mock_query.await_args.kwargs["filters"] is None

    @pytest.mark.asyncio
    async def test_working_representation_scoped_end_to_end(
        self,
        db_session: AsyncSession,
        sample_data: tuple[models.Workspace, models.Peer],
    ):
        """All blended paths active: only in-scope content may appear."""
        test_workspace, test_peer = sample_data
        session_a, _, manager = await self._setup(db_session, test_workspace, test_peer)

        representation = await manager.get_working_representation(
            db=db_session,
            session_allowlist=[session_a.name],
            include_most_derived=True,
        )

        contents = [obs.content for obs in representation.explicit]
        assert "in-scope observation" in contents
        assert "out-of-scope observation" not in contents
        assert "sessionless dream observation" not in contents

    @pytest.mark.asyncio
    async def test_empty_allowlist_fails_closed(
        self,
        db_session: AsyncSession,
        sample_data: tuple[models.Workspace, models.Peer],
    ):
        """An empty allowlist must return an empty representation, not fall
        back to unscoped behavior (downstream stores drop empty IN clauses)."""
        test_workspace, test_peer = sample_data
        _, _, manager = await self._setup(db_session, test_workspace, test_peer)

        representation = await manager.get_working_representation(
            db=db_session,
            session_allowlist=[],
            include_most_derived=True,
        )

        assert representation.explicit == []
        assert representation.deductive == []

    def test_build_filter_conditions_empty_allowlist_fails_closed(self):
        """The filter-builder layer itself must fail closed, independent of the
        early-return guard in _get_working_representation_internal. An empty
        allowlist emits an empty `in` (renders as always-false downstream), not
        an omitted filter."""
        manager = RepresentationManager(
            "workspace", observer="observer", observed="observed"
        )

        # Scoping also narrows to levels whose session stamp is trustworthy
        # (see ALLOWLIST_SAFE_LEVELS).
        assert manager._build_filter_conditions(session_allowlist=[]) == {  # pyright: ignore[reportPrivateUsage]
            "session_name": {"in": []},
            "level": {"in": ["explicit"]},
        }
        # None means unscoped — no session filter and no level narrowing.
        assert manager._build_filter_conditions(session_allowlist=None) == {}  # pyright: ignore[reportPrivateUsage]
        assert manager._build_filter_conditions(session_allowlist=["s1"]) == {  # pyright: ignore[reportPrivateUsage]
            "session_name": {"in": ["s1"]},
            "level": {"in": ["explicit"]},
        }
        # A requested level outside the safe set yields an empty `in`, which
        # matches nothing rather than falling back to unscoped recall.
        assert manager._build_filter_conditions(  # pyright: ignore[reportPrivateUsage]
            level="inductive", session_allowlist=["s1"]
        ) == {
            "session_name": {"in": ["s1"]},
            "level": {"in": []},
        }
        # ...while an unscoped level filter is left exactly as asked.
        assert manager._build_filter_conditions(level="inductive") == {  # pyright: ignore[reportPrivateUsage]
            "level": "inductive"
        }


class TestRepresentationManagerSave:
    @pytest.mark.asyncio
    async def test_save_representation_filters_blank_observations_before_embedding(
        self,
    ):
        manager = RepresentationManager(
            "workspace",
            observer="observer",
            observed="observed",
        )
        representation = Representation(
            explicit=[
                ExplicitObservation(
                    content="   ",
                    created_at=datetime.now(timezone.utc),
                    message_ids=[1],
                    session_name="session",
                ),
                ExplicitObservation(
                    content=" useful observation ",
                    created_at=datetime.now(timezone.utc),
                    message_ids=[1],
                    session_name="session",
                ),
            ]
        )

        with (
            patch("src.crud.representation.tracked_db", _fake_tracked_db),
            patch(
                "src.crud.representation.embedding_client.simple_batch_embed",
                new=AsyncMock(return_value=[[0.1]]),
            ) as mock_embed,
            patch.object(
                manager,
                "_save_representation_internal",
                new=AsyncMock(
                    return_value=CreateDocumentsResult(created_documents=[MagicMock()])
                ),
            ) as mock_save,
        ):
            saved = await manager.save_representation(
                representation,
                message_ids=[1],
                session_name="session",
                message_created_at=datetime.now(timezone.utc),
                message_level_configuration=_resolved_config(),
            )

        assert len(saved.created_documents) == 1
        mock_embed.assert_awaited_once_with(
            ["useful observation"], on_oversize="truncate"
        )
        saved_observations = _saved_observations(mock_save)
        assert len(saved_observations) == 1
        assert saved_observations[0].content == "useful observation"

    @pytest.mark.asyncio
    async def test_save_representation_filters_blank_deductive_observations(self):
        manager = RepresentationManager(
            "workspace",
            observer="observer",
            observed="observed",
        )
        representation = Representation(
            deductive=[
                DeductiveObservation(
                    conclusion="   ",
                    premises=["premise a"],
                    source_ids=["doc-a"],
                    created_at=datetime.now(timezone.utc),
                    message_ids=[1],
                    session_name="session",
                ),
                DeductiveObservation(
                    conclusion=" inferred conclusion ",
                    premises=["premise b"],
                    source_ids=["doc-b"],
                    created_at=datetime.now(timezone.utc),
                    message_ids=[1],
                    session_name="session",
                ),
            ]
        )

        with (
            patch("src.crud.representation.tracked_db", _fake_tracked_db),
            patch(
                "src.crud.representation.embedding_client.simple_batch_embed",
                new=AsyncMock(return_value=[[0.2]]),
            ) as mock_embed,
            patch.object(
                manager,
                "_save_representation_internal",
                new=AsyncMock(
                    return_value=CreateDocumentsResult(created_documents=[MagicMock()])
                ),
            ) as mock_save,
        ):
            saved = await manager.save_representation(
                representation,
                message_ids=[1],
                session_name="session",
                message_created_at=datetime.now(timezone.utc),
                message_level_configuration=_resolved_config(),
            )

        assert len(saved.created_documents) == 1
        mock_embed.assert_awaited_once_with(
            ["inferred conclusion"], on_oversize="truncate"
        )
        saved_observations = _saved_observations(mock_save)
        assert len(saved_observations) == 1
        assert isinstance(saved_observations[0], DeductiveObservation)
        assert saved_observations[0].conclusion == "inferred conclusion"

    @pytest.mark.asyncio
    async def test_save_representation_skips_all_blank_observations(self):
        manager = RepresentationManager(
            "workspace",
            observer="observer",
            observed="observed",
        )
        representation = Representation(
            explicit=[
                ExplicitObservation(
                    content="",
                    created_at=datetime.now(timezone.utc),
                    message_ids=[1],
                    session_name="session",
                ),
                ExplicitObservation(
                    content="\n\t ",
                    created_at=datetime.now(timezone.utc),
                    message_ids=[1],
                    session_name="session",
                ),
            ]
        )

        with (
            patch("src.crud.representation.tracked_db", _fake_tracked_db),
            patch(
                "src.crud.representation.embedding_client.simple_batch_embed",
                new=AsyncMock(),
            ) as mock_embed,
            patch.object(
                manager,
                "_save_representation_internal",
                new=AsyncMock(),
            ) as mock_save,
        ):
            saved = await manager.save_representation(
                representation,
                message_ids=[1],
                session_name="session",
                message_created_at=datetime.now(timezone.utc),
                message_level_configuration=_resolved_config(),
            )

        assert len(saved.created_documents) == 0
        mock_embed.assert_not_awaited()
        mock_save.assert_not_awaited()

    @pytest.mark.asyncio
    async def test_save_representation_embeds_with_truncate_on_oversize(self):
        """One oversize observation must not drop the rest of the batch."""
        manager = RepresentationManager(
            "workspace",
            observer="observer",
            observed="observed",
        )
        representation = Representation(
            explicit=[
                ExplicitObservation(
                    content="short fact",
                    created_at=datetime.now(timezone.utc),
                    message_ids=[1],
                    session_name="session",
                )
            ],
            deductive=[
                DeductiveObservation(
                    conclusion="inferred fact",
                    premises=["premise"],
                    source_ids=["doc-a"],
                    created_at=datetime.now(timezone.utc),
                    message_ids=[1],
                    session_name="session",
                )
            ],
        )

        with (
            patch("src.crud.representation.tracked_db", _fake_tracked_db),
            patch(
                "src.crud.representation.embedding_client.simple_batch_embed",
                new=AsyncMock(return_value=[[0.1], [0.2]]),
            ) as mock_embed,
            patch.object(
                manager,
                "_save_representation_internal",
                new=AsyncMock(
                    return_value=CreateDocumentsResult(created_documents=[MagicMock()])
                ),
            ),
        ):
            await manager.save_representation(
                representation,
                message_ids=[1],
                session_name="session",
                message_created_at=datetime.now(timezone.utc),
                message_level_configuration=_resolved_config(),
            )

        mock_embed.assert_awaited_once_with(
            ["inferred fact", "short fact"], on_oversize="truncate"
        )


class TestVectorQueryTopKFloor:
    """Regression for HONCHO-19Q / HONCHO-4Q4.

    A top_k of 0 reached Turbopuffer, which rejects it with a 400
    ('top_k must be between 1 and 10000'). Two independent paths produced it:
    the working-representation budget split (``total // 3`` rounds to 0 for
    max_conclusions < 3) and the dialectic ``search_memory`` tool, whose
    LLM-supplied top_k has an upper clamp but no floor.
    """

    @pytest.mark.asyncio
    async def test_query_documents_returns_empty_without_querying_on_zero_top_k(self):
        """The choke point every semantic document query routes through."""
        from src.crud.document import query_documents

        with (
            patch(
                "src.crud.document.embedding_client.embed", new=AsyncMock()
            ) as mock_embed,
            patch(
                "src.crud.document.query_external_vector_document_ids",
                new=AsyncMock(),
            ) as mock_vector,
        ):
            for top_k in (0, -1):
                assert (
                    await query_documents(
                        None,
                        "workspace",
                        "query",
                        observer="observer",
                        observed="observed",
                        top_k=top_k,
                    )
                    == []
                )

        mock_embed.assert_not_awaited()
        mock_vector.assert_not_awaited()

    @pytest.mark.asyncio
    async def test_requested_semantic_search_always_gets_budget(
        self,
        db_session: AsyncSession,
        sample_data: tuple[models.Workspace, models.Peer],
    ):
        """max_conclusions < 3 must not allocate 0 to an explicitly requested search."""
        test_workspace, test_peer = sample_data
        manager = RepresentationManager(
            test_workspace.name, observer=test_peer.name, observed=test_peer.name
        )

        for max_observations in (1, 2, 100):
            with patch(
                "src.crud.query_documents", new=AsyncMock(return_value=[])
            ) as mock_query:
                await manager._get_working_representation_internal(  # pyright: ignore[reportPrivateUsage]
                    db_session,
                    include_semantic_query="what do they like?",
                    embedding=[0.1],
                    max_observations=max_observations,
                )

            assert mock_query.await_args is not None
            top_k = mock_query.await_args.kwargs["top_k"]
            assert top_k >= 1, f"max_observations={max_observations} gave top_k={top_k}"
            assert top_k <= max_observations