"""
Tests for the immediate message-embedding fast path (src/reconciler/embed_now.py).

These exercise embed_messages_now end-to-end against the test database: it opens
its own tracked_db sessions (patched to the test engine in conftest), so each test
creates committed fixture rows and asserts on the result via the provided session.
"""

from unittest.mock import AsyncMock, patch

import pytest
from fastapi import BackgroundTasks
from nanoid import generate as generate_nanoid
from prometheus_client import REGISTRY
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker

from src import models
from src.config import settings
from src.reconciler.embed_now import (
    embed_messages_now,
    embed_task_gate,
    reset_embed_semaphore,
)
from src.vector_store import VectorStore


async def _create_message_with_pending_chunks(
    db_session: AsyncSession,
    workspace: models.Workspace,
    peer: models.Peer,
    chunk_contents: list[str],
) -> tuple[str, list[int]]:
    """Create a message plus one pending MessageEmbedding row per chunk.

    Returns (message public_id, ordered embedding row ids).
    """
    session = models.Session(name=str(generate_nanoid()), workspace_name=workspace.name)
    db_session.add(session)
    await db_session.commit()

    message_id = str(generate_nanoid())
    message = models.Message(
        public_id=message_id,
        session_name=session.name,
        workspace_name=workspace.name,
        peer_name=peer.name,
        content=" ".join(chunk_contents),
        seq_in_session=1,
    )
    db_session.add(message)
    await db_session.commit()

    rows = [
        models.MessageEmbedding(
            content=chunk,
            message_id=message_id,
            workspace_name=workspace.name,
            session_name=session.name,
            peer_name=peer.name,
            sync_state="pending",
            embedding=None,
        )
        for chunk in chunk_contents
    ]
    db_session.add_all(rows)
    await db_session.commit()
    for row in rows:
        await db_session.refresh(row)
    return message_id, [row.id for row in rows]


@pytest.fixture(autouse=True)
def reset_semaphore_fixture():
    """Rebuild the module semaphore per test so it binds to the active loop,
    and clear the admission gate's in-flight count."""
    reset_embed_semaphore()
    embed_task_gate.in_flight = 0
    yield
    reset_embed_semaphore()
    embed_task_gate.in_flight = 0


@pytest.mark.asyncio
class TestEmbedTaskGate:
    """Admission gate for immediate-embed background tasks
    (EMBEDDING.MAX_PENDING_EMBED_TASKS)."""

    async def test_admits_under_cap_and_releases_slot(self) -> None:
        """Under the cap, the task is scheduled; running it embeds the given ids
        and releases the slot."""
        tasks = BackgroundTasks()
        with (
            patch.object(settings.EMBEDDING, "MAX_PENDING_EMBED_TASKS", 2),
            patch(
                "src.reconciler.embed_now.embed_messages_now", new=AsyncMock()
            ) as mock_embed,
        ):
            assert embed_task_gate.try_schedule(tasks, ["msg_1"]) is True
            assert embed_task_gate.in_flight == 1
            await tasks()
        mock_embed.assert_awaited_once_with(["msg_1"])
        assert embed_task_gate.in_flight == 0

    async def test_rejects_at_cap_without_scheduling(self) -> None:
        """At the cap, nothing is scheduled and False is returned."""
        tasks = BackgroundTasks()
        with patch.object(settings.EMBEDDING, "MAX_PENDING_EMBED_TASKS", 1):
            assert embed_task_gate.try_schedule(tasks, ["msg_1"]) is True
            assert embed_task_gate.try_schedule(tasks, ["msg_2"]) is False
        assert len(tasks.tasks) == 1
        assert embed_task_gate.in_flight == 1

    async def test_slot_released_when_task_raises(self) -> None:
        """A failing task still releases its slot (finally path)."""
        tasks = BackgroundTasks()
        with (
            patch.object(settings.EMBEDDING, "MAX_PENDING_EMBED_TASKS", 1),
            patch(
                "src.reconciler.embed_now.embed_messages_now",
                new=AsyncMock(side_effect=RuntimeError("boom")),
            ),
        ):
            assert embed_task_gate.try_schedule(tasks, ["msg_1"]) is True
            with pytest.raises(RuntimeError):
                await tasks()
        assert embed_task_gate.in_flight == 0

    async def test_gauge_mirrors_in_flight_count(self) -> None:
        """With metrics enabled, the Prometheus gauge tracks the gate's
        in-flight count through schedule and release."""

        def gauge_value() -> float | None:
            return REGISTRY.get_sample_value(
                "embed_now_tasks_in_flight", {"namespace": "test"}
            )

        tasks = BackgroundTasks()
        with (
            patch.object(settings.EMBEDDING, "MAX_PENDING_EMBED_TASKS", 2),
            patch.object(settings.METRICS, "ENABLED", True),
            patch.object(settings.METRICS, "NAMESPACE", "test"),
            patch("src.reconciler.embed_now.embed_messages_now", new=AsyncMock()),
        ):
            assert embed_task_gate.try_schedule(tasks, ["msg_1"]) is True
            assert gauge_value() == 1
            await tasks()
            assert gauge_value() == 0

    async def test_zero_cap_disables_fast_path(self) -> None:
        """MAX_PENDING_EMBED_TASKS=0 rejects every schedule attempt."""
        tasks = BackgroundTasks()
        with patch.object(settings.EMBEDDING, "MAX_PENDING_EMBED_TASKS", 0):
            assert embed_task_gate.try_schedule(tasks, ["msg_1"]) is False
        assert len(tasks.tasks) == 0


@pytest.mark.asyncio
class TestEmbedMessagesNow:
    async def test_pgvector_happy_path_marks_synced(
        self,
        db_session: AsyncSession,
        sample_data: tuple[models.Workspace, models.Peer],
    ) -> None:
        """pgvector-only mode: rows get a vector and flip to synced immediately."""
        workspace, peer = sample_data
        message_id, emb_ids = await _create_message_with_pending_chunks(
            db_session, workspace, peer, ["hello world"]
        )

        await embed_messages_now([message_id])

        for emb_id in emb_ids:
            row = await db_session.get(models.MessageEmbedding, emb_id)
            assert row is not None
            await db_session.refresh(row)
            assert row.sync_state == "synced"
            assert row.embedding is not None
            assert row.sync_attempts == 0

    async def test_no_message_ids_is_noop(self) -> None:
        """Empty input returns without touching the DB or embedding."""
        with patch(
            "src.embedding_client.embedding_client.simple_batch_embed"
        ) as mock_embed:
            await embed_messages_now([])
        mock_embed.assert_not_called()

    async def test_already_synced_rows_not_reclaimed(
        self,
        db_session: AsyncSession,
        sample_data: tuple[models.Workspace, models.Peer],
    ) -> None:
        """A second run finds no pending rows and does not re-embed."""
        workspace, peer = sample_data
        message_id, _ = await _create_message_with_pending_chunks(
            db_session, workspace, peer, ["first content"]
        )
        await embed_messages_now([message_id])

        with patch(
            "src.embedding_client.embedding_client.simple_batch_embed"
        ) as mock_embed:
            await embed_messages_now([message_id])
        mock_embed.assert_not_called()

    async def test_embed_failure_leaves_rows_pending_and_leased(
        self,
        db_session: AsyncSession,
        sample_data: tuple[models.Workspace, models.Peer],
    ) -> None:
        """Embedding failure must leave rows pending + leased, attempts untouched,
        so the reconciler owns retry accounting."""
        workspace, peer = sample_data
        message_id, emb_ids = await _create_message_with_pending_chunks(
            db_session, workspace, peer, ["will fail"]
        )

        with patch(
            "src.embedding_client.embedding_client.simple_batch_embed",
            new=AsyncMock(side_effect=RuntimeError("provider down")),
        ):
            await embed_messages_now([message_id])

        for emb_id in emb_ids:
            row = await db_session.get(models.MessageEmbedding, emb_id)
            assert row is not None
            await db_session.refresh(row)
            assert row.sync_state == "pending"
            assert row.embedding is None
            assert row.sync_attempts == 0  # lease only, no attempt bump
            assert row.last_sync_at is not None  # leased

    async def test_external_store_upserts_with_chunk_positioned_ids(
        self,
        db_session: AsyncSession,
        sample_data: tuple[models.Workspace, models.Peer],
        mock_vector_store: VectorStore,
    ) -> None:
        """External-store mode: upsert each chunk with id {message_id}_{position}
        and mark rows synced."""
        workspace, peer = sample_data
        message_id, emb_ids = await _create_message_with_pending_chunks(
            db_session, workspace, peer, ["chunk a", "chunk b", "chunk c"]
        )

        with patch(
            "src.reconciler.embed_now.get_external_vector_store",
            return_value=mock_vector_store,
        ):
            await embed_messages_now([message_id])

        upsert_mock: AsyncMock = mock_vector_store.upsert_many  # pyright: ignore[reportAssignmentType]
        upsert_mock.assert_awaited()
        upserted_ids = {
            record.id for call in upsert_mock.await_args_list for record in call.args[1]
        }
        assert upserted_ids == {
            f"{message_id}_0",
            f"{message_id}_1",
            f"{message_id}_2",
        }

        for emb_id in emb_ids:
            row = await db_session.get(models.MessageEmbedding, emb_id)
            assert row is not None
            await db_session.refresh(row)
            assert row.sync_state == "synced"

    async def test_locked_chunk_skipped_keeps_positions_stable(
        self,
        db_session: AsyncSession,
        db_engine: AsyncEngine,
        sample_data: tuple[models.Workspace, models.Peer],
        mock_vector_store: VectorStore,
    ) -> None:
        """If a sibling chunk is locked by another txn, SKIP LOCKED skips it but
        chunk positions still come from the full sibling ordering — so the claimed
        chunks keep their {message_id}_0 / _2 ids (not _0 / _1)."""
        workspace, peer = sample_data
        message_id, emb_ids = await _create_message_with_pending_chunks(
            db_session, workspace, peer, ["chunk a", "chunk b", "chunk c"]
        )
        locked_id = emb_ids[1]  # middle chunk -> position 1

        # Hold a row lock on the middle chunk from an independent transaction.
        lock_factory = async_sessionmaker(bind=db_engine, expire_on_commit=False)
        lock_session = lock_factory()
        await lock_session.execute(
            select(models.MessageEmbedding)
            .where(models.MessageEmbedding.id == locked_id)
            .with_for_update()
        )
        try:
            with patch(
                "src.reconciler.embed_now.get_external_vector_store",
                return_value=mock_vector_store,
            ):
                await embed_messages_now([message_id])
        finally:
            await lock_session.rollback()
            await lock_session.close()

        upsert_mock: AsyncMock = mock_vector_store.upsert_many  # pyright: ignore[reportAssignmentType]
        upserted_ids = {
            record.id for call in upsert_mock.await_args_list for record in call.args[1]
        }
        assert upserted_ids == {f"{message_id}_0", f"{message_id}_2"}

        # The locked chunk stays pending; the other two are synced.
        locked_row = await db_session.get(models.MessageEmbedding, locked_id)
        assert locked_row is not None
        await db_session.refresh(locked_row)
        assert locked_row.sync_state == "pending"
        for emb_id in (emb_ids[0], emb_ids[2]):
            row = await db_session.get(models.MessageEmbedding, emb_id)
            assert row is not None
            await db_session.refresh(row)
            assert row.sync_state == "synced"

    async def test_external_store_unavailable_leaves_rows_pending(
        self,
        db_session: AsyncSession,
        sample_data: tuple[models.Workspace, models.Peer],
        mock_vector_store: VectorStore,
    ) -> None:
        """External-store mode: if upsert_many raises VectorStoreError, rows must
        stay pending with no vector and untouched attempts, so the reconciler
        heals them. embed_now never bumps sync_attempts."""
        from src.exceptions import VectorStoreError

        workspace, peer = sample_data
        message_id, emb_ids = await _create_message_with_pending_chunks(
            db_session, workspace, peer, ["chunk a", "chunk b"]
        )

        upsert_mock: AsyncMock = mock_vector_store.upsert_many  # pyright: ignore[reportAssignmentType]
        upsert_mock.side_effect = VectorStoreError("vector store down")

        with patch(
            "src.reconciler.embed_now.get_external_vector_store",
            return_value=mock_vector_store,
        ):
            await embed_messages_now([message_id])

        for emb_id in emb_ids:
            row = await db_session.get(models.MessageEmbedding, emb_id)
            assert row is not None
            await db_session.refresh(row)
            assert row.sync_state == "pending"
            assert row.embedding is None
            assert row.sync_attempts == 0  # embed_now never bumps attempts