"""
Tests for the distributed worker system.

Tests cover:
- BrokerTaskBackend task submission and storage
- WorkerPoller task claiming with FOR UPDATE SKIP LOCKED
- Concurrent workers claiming different tasks (no duplicates)
- Task completion and failure handling
- Retry mechanism
- Worker decommissioning
"""

import asyncio
import json
import uuid

import pytest
import pytest_asyncio

from hindsight_api.engine.task_backend import BrokerTaskBackend, SyncTaskBackend


async def _ensure_bank(pool, bank_id: str) -> None:
    """Upsert a minimal bank row so FK on async_operations passes."""
    await pool.execute(
        "INSERT INTO banks (bank_id, name) VALUES ($1, $2) ON CONFLICT DO NOTHING",
        bank_id,
        bank_id,
    )


# Use loadgroup to ensure these tests run in the same worker
# since they share database state
pytestmark = pytest.mark.xdist_group("worker_tests")


@pytest_asyncio.fixture
async def pool(pg0_db_url):
    """Create a dedicated connection pool for worker tests."""
    import asyncpg

    from hindsight_api.pg0 import resolve_database_url

    # Resolve pg0:// URL to postgresql:// URL if needed
    resolved_url = await resolve_database_url(pg0_db_url)

    pool = await asyncpg.create_pool(
        resolved_url,
        min_size=2,
        max_size=10,
        command_timeout=30,
    )
    yield pool
    await pool.close()


@pytest_asyncio.fixture
async def clean_operations(pool):
    """Clean up async_operations table before and after tests."""
    # Clean before test - covers both 'test-worker-' and 'test_worker_recovery' patterns
    await pool.execute(
        "DELETE FROM async_operations WHERE bank_id LIKE 'test-worker-%' OR bank_id LIKE 'test_worker_%'"
    )
    yield
    # Clean after test
    await pool.execute(
        "DELETE FROM async_operations WHERE bank_id LIKE 'test-worker-%' OR bank_id LIKE 'test_worker_%'"
    )


class TestBrokerTaskBackend:
    """Tests for BrokerTaskBackend task storage."""

    @pytest.mark.asyncio
    async def test_submit_task_updates_existing_operation(self, pool, clean_operations):
        """Test that submit_task updates task_payload for existing operations."""
        # Create an operation record first
        operation_id = uuid.uuid4()
        bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
        await _ensure_bank(pool, bank_id)

        await pool.execute(
            """
            INSERT INTO async_operations (operation_id, bank_id, operation_type, status)
            VALUES ($1, $2, 'test_operation', 'pending')
            """,
            operation_id,
            bank_id,
        )

        # Submit task with same operation_id
        backend = BrokerTaskBackend(pool_getter=lambda: pool)
        await backend.initialize()

        task_dict = {
            "operation_id": str(operation_id),
            "type": "test_task",
            "bank_id": bank_id,
            "data": {"key": "value"},
        }
        await backend.submit_task(task_dict)

        # Verify task_payload was stored
        row = await pool.fetchrow(
            "SELECT task_payload, status FROM async_operations WHERE operation_id = $1",
            operation_id,
        )
        assert row is not None
        assert row["status"] == "pending"
        payload = json.loads(row["task_payload"])
        assert payload["type"] == "test_task"
        assert payload["data"] == {"key": "value"}

    @pytest.mark.asyncio
    async def test_submit_task_creates_new_operation(self, pool, clean_operations):
        """Test that submit_task creates new operation when no operation_id provided."""
        backend = BrokerTaskBackend(pool_getter=lambda: pool)
        await backend.initialize()

        bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
        await _ensure_bank(pool, bank_id)
        task_dict = {
            "type": "access_count_update",
            "bank_id": bank_id,
            "node_ids": ["node1", "node2"],
        }
        await backend.submit_task(task_dict)

        # Verify new operation was created
        row = await pool.fetchrow(
            "SELECT operation_type, status, task_payload FROM async_operations WHERE bank_id = $1",
            bank_id,
        )
        assert row is not None
        assert row["operation_type"] == "access_count_update"
        assert row["status"] == "pending"
        payload = json.loads(row["task_payload"])
        assert payload["node_ids"] == ["node1", "node2"]


class TestWorkerPoller:
    """Tests for WorkerPoller task claiming and execution."""

    @pytest.mark.asyncio
    async def test_claim_batch_claims_pending_tasks(self, pool, clean_operations):
        """Test that claim_batch claims pending tasks with task_payload."""
        from hindsight_api.worker import WorkerPoller

        # Create some pending tasks
        bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
        await _ensure_bank(pool, bank_id)
        for i in range(3):
            op_id = uuid.uuid4()
            payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id})
            await pool.execute(
                """
                INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
                VALUES ($1, $2, 'test', 'pending', $3::jsonb)
                """,
                op_id,
                bank_id,
                payload,
            )

        # Create poller and claim tasks
        executed_tasks = []

        async def mock_executor(task_dict):
            executed_tasks.append(task_dict)

        poller = WorkerPoller(
            pool=pool,
            worker_id="test-worker-1",
            executor=mock_executor,
        )

        claimed = await poller.claim_batch()
        assert len(claimed) == 3

        # ClaimedTask objects have operation_id, task_dict, schema attributes
        for task in claimed:
            assert task.operation_id is not None
            assert task.task_dict is not None

        # Verify tasks are marked as processing with worker_id
        rows = await pool.fetch(
            "SELECT status, worker_id FROM async_operations WHERE bank_id = $1",
            bank_id,
        )
        for row in rows:
            assert row["status"] == "processing"
            assert row["worker_id"] == "test-worker-1"

    @pytest.mark.asyncio
    async def test_claim_batch_respects_max_slots(self, pool, clean_operations):
        """Test that claim_batch respects the max_slots limit."""
        from hindsight_api.worker import WorkerPoller

        # Create 10 pending tasks
        bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
        await _ensure_bank(pool, bank_id)
        for i in range(10):
            op_id = uuid.uuid4()
            payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id})
            await pool.execute(
                """
                INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
                VALUES ($1, $2, 'test', 'pending', $3::jsonb)
                """,
                op_id,
                bank_id,
                payload,
            )

        poller = WorkerPoller(
            pool=pool,
            worker_id="test-worker-1",
            executor=lambda x: None,
            max_slots=3,  # Limit to 3 concurrent tasks
        )

        claimed = await poller.claim_batch()
        assert len(claimed) == 3

    @pytest.mark.asyncio
    async def test_execute_task_executor_marks_completed(self, pool, clean_operations):
        """Test that executor's status marking is preserved by the poller.

        The executor (MemoryEngine.execute_task) handles marking operations as completed/failed.
        The poller should NOT override those status updates.
        """
        from hindsight_api.worker import WorkerPoller
        from hindsight_api.worker.poller import ClaimedTask

        # Create a pending task
        bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
        op_id = uuid.uuid4()
        payload = json.dumps({"type": "test_task", "operation_id": str(op_id), "bank_id": bank_id})
        await _ensure_bank(pool, bank_id)
        await pool.execute(
            """
            INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id)
            VALUES ($1, $2, 'test', 'processing', $3::jsonb, 'test-worker-1')
            """,
            op_id,
            bank_id,
            payload,
        )

        executed = []

        async def mock_executor(task_dict):
            """Executor that marks its own status as completed (like MemoryEngine.execute_task)."""
            executed.append(task_dict)
            await pool.execute(
                """
                UPDATE async_operations
                SET status = 'completed', completed_at = now(), updated_at = now()
                WHERE operation_id = $1
                """,
                op_id,
            )

        poller = WorkerPoller(
            pool=pool,
            worker_id="test-worker-1",
            executor=mock_executor,
        )

        # Execute the task (fire-and-forget)
        task_dict = json.loads(payload)
        claimed_task = ClaimedTask(operation_id=str(op_id), task_dict=task_dict, schema=None)
        await poller.execute_task(claimed_task)

        # Wait for background task to complete
        completed = await poller.wait_for_active_tasks(timeout=5.0)
        assert completed, "Task did not complete within timeout"
        assert len(executed) == 1

        # Verify task is marked as completed (by executor, not overridden by poller)
        row = await pool.fetchrow(
            "SELECT status, completed_at FROM async_operations WHERE operation_id = $1",
            op_id,
        )
        assert row["status"] == "completed"
        assert row["completed_at"] is not None

    @pytest.mark.asyncio
    async def test_executor_exception_triggers_retry(self, pool, clean_operations):
        """Test that exceptions from the executor trigger _retry_or_fail (not a crash).

        When the executor re-raises an exception (as MemoryEngine.execute_task does for
        retryable task failures), the poller calls _retry_or_fail, which resets the task
        back to 'pending' and increments retry_count so it can be reclaimed.

        This is the fix for the consolidation deadlock: previously submit_task was called
        with only a task_payload update, leaving status='processing' forever.
        """
        from hindsight_api.worker import WorkerPoller
        from hindsight_api.worker.poller import ClaimedTask

        bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
        op_id = uuid.uuid4()
        payload = json.dumps({"type": "consolidation", "operation_id": str(op_id), "bank_id": bank_id})
        await _ensure_bank(pool, bank_id)
        await pool.execute(
            """
            INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id, claimed_at)
            VALUES ($1, $2, 'consolidation', 'processing', $3::jsonb, 'test-worker-1', now())
            """,
            op_id,
            bank_id,
            payload,
        )

        from datetime import datetime, timezone

        from hindsight_api.worker.exceptions import RetryTaskAt

        async def failing_executor(task_dict):
            raise RetryTaskAt(retry_at=datetime.now(timezone.utc), message="TimeoutError during recall")

        poller = WorkerPoller(
            pool=pool,
            worker_id="test-worker-1",
            executor=failing_executor,
        )

        task_dict = json.loads(payload)
        claimed_task = ClaimedTask(operation_id=str(op_id), task_dict=task_dict, schema=None)
        await poller.execute_task(claimed_task)

        completed = await poller.wait_for_active_tasks(timeout=5.0)
        assert completed, "Task did not complete within timeout"

        # Task must be reset to 'pending' with worker_id/claimed_at cleared — not left as
        # 'processing', which would cause a permanent deadlock via the NOT EXISTS guard.
        row = await pool.fetchrow(
            "SELECT status, worker_id, claimed_at, retry_count FROM async_operations WHERE operation_id = $1",
            op_id,
        )
        assert row["status"] == "pending", (
            f"REGRESSION: Task status is '{row['status']}' instead of 'pending'. "
            "A task stuck in 'processing' after a retry causes a consolidation deadlock."
        )
        assert row["worker_id"] is None, "worker_id must be cleared on retry"
        assert row["claimed_at"] is None, "claimed_at must be cleared on retry"
        assert row["retry_count"] == 1

    @pytest.mark.asyncio
    async def test_executor_exception_marks_failed_immediately(self, pool, clean_operations):
        """Test that a plain exception (not RetryTaskAt) permanently marks a task as 'failed'.

        With the task-owned retry model, plain exceptions are non-retryable — the poller
        marks them as 'failed' immediately. Tasks that want to be retried must raise RetryTaskAt.
        """
        from hindsight_api.worker import WorkerPoller
        from hindsight_api.worker.poller import ClaimedTask

        bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
        op_id = uuid.uuid4()
        payload = json.dumps({"type": "consolidation", "operation_id": str(op_id), "bank_id": bank_id})
        await _ensure_bank(pool, bank_id)
        await pool.execute(
            """
            INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id, claimed_at, retry_count)
            VALUES ($1, $2, 'consolidation', 'processing', $3::jsonb, 'test-worker-1', now(), 0)
            """,
            op_id,
            bank_id,
            payload,
        )

        async def failing_executor(task_dict):
            raise ValueError("Non-retryable error")

        poller = WorkerPoller(
            pool=pool,
            worker_id="test-worker-1",
            executor=failing_executor,
        )

        task_dict = json.loads(payload)
        claimed_task = ClaimedTask(operation_id=str(op_id), task_dict=task_dict, schema=None)
        await poller.execute_task(claimed_task)

        completed = await poller.wait_for_active_tasks(timeout=5.0)
        assert completed, "Task did not complete within timeout"

        row = await pool.fetchrow(
            "SELECT status, error_message, retry_count FROM async_operations WHERE operation_id = $1",
            op_id,
        )
        assert row["status"] == "failed", f"Expected 'failed' for plain exception, got '{row['status']}'"
        assert row["error_message"] is not None
        assert row["retry_count"] == 0  # not incremented; plain exception = immediate fail

    @pytest.mark.asyncio
    async def test_executor_failed_status_not_overridden(self, pool, clean_operations):
        """REGRESSION TEST: Verify poller does NOT overwrite executor's 'failed' status to 'completed'.

        This test covers the non-retryable failure path (e.g., file_convert_retain):
        1. Executor catches an internal error, marks the operation as 'failed' in the DB
        2. Executor returns normally (does NOT re-raise) — so no exception reaches the poller
        3. The poller must NOT overwrite the 'failed' status to 'completed'

        Retryable failures re-raise instead (see test_executor_exception_triggers_retry).
        """
        from hindsight_api.worker import WorkerPoller
        from hindsight_api.worker.poller import ClaimedTask

        bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
        op_id = uuid.uuid4()
        payload = json.dumps({"type": "test_task", "operation_id": str(op_id), "bank_id": bank_id})
        await _ensure_bank(pool, bank_id)
        await pool.execute(
            """
            INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id)
            VALUES ($1, $2, 'test', 'processing', $3::jsonb, 'test-worker-1')
            """,
            op_id,
            bank_id,
            payload,
        )

        async def executor_that_marks_failed(task_dict):
            """Simulates MemoryEngine.execute_task behavior on internal error.

            The executor catches the error, marks the operation as 'failed',
            and returns normally (does NOT re-raise the exception).
            """
            # Simulate internal failure handling (like MemoryEngine._mark_operation_failed)
            await pool.execute(
                """
                UPDATE async_operations
                SET status = 'failed', error_message = $2, completed_at = now(), updated_at = now()
                WHERE operation_id = $1
                """,
                op_id,
                "Simulated conversion error: file format not supported",
            )
            # Returns normally - this is the key: executor does NOT re-raise

        poller = WorkerPoller(
            pool=pool,
            worker_id="test-worker-1",
            executor=executor_that_marks_failed,
        )

        task_dict = json.loads(payload)
        claimed_task = ClaimedTask(operation_id=str(op_id), task_dict=task_dict, schema=None)
        await poller.execute_task(claimed_task)

        completed = await poller.wait_for_active_tasks(timeout=5.0)
        assert completed, "Task did not complete within timeout"

        # THE KEY ASSERTION: Status must be 'failed', NOT 'completed'
        row = await pool.fetchrow(
            "SELECT status, error_message FROM async_operations WHERE operation_id = $1",
            op_id,
        )
        assert row["status"] == "failed", (
            f"REGRESSION: Poller overwrote executor's 'failed' status to '{row['status']}'. "
            "The poller must not override status set by the executor."
        )
        assert "Simulated conversion error" in row["error_message"]

    @pytest.mark.asyncio
    async def test_claim_batch_skips_consolidation_when_same_bank_processing(self, pool, clean_operations):
        """Test that pending consolidation is skipped if same bank has one processing."""
        from hindsight_api.worker import WorkerPoller

        bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
        await _ensure_bank(pool, bank_id)

        # Create a processing consolidation for bank
        processing_op_id = uuid.uuid4()
        await pool.execute(
            """
            INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id)
            VALUES ($1, $2, 'consolidation', 'processing', $3::jsonb, 'other-worker')
            """,
            processing_op_id,
            bank_id,
            json.dumps({"type": "consolidation", "bank_id": bank_id}),
        )

        # Create a pending consolidation for same bank (should be skipped)
        pending_op_id = uuid.uuid4()
        await pool.execute(
            """
            INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
            VALUES ($1, $2, 'consolidation', 'pending', $3::jsonb)
            """,
            pending_op_id,
            bank_id,
            json.dumps({"type": "consolidation", "bank_id": bank_id}),
        )

        # Create a pending consolidation for different bank (should be claimed)
        other_bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
        await _ensure_bank(pool, other_bank_id)
        other_op_id = uuid.uuid4()
        await pool.execute(
            """
            INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
            VALUES ($1, $2, 'consolidation', 'pending', $3::jsonb)
            """,
            other_op_id,
            other_bank_id,
            json.dumps({"type": "consolidation", "bank_id": other_bank_id}),
        )

        poller = WorkerPoller(
            pool=pool,
            worker_id="test-worker-1",
            executor=lambda x: None,
        )

        claimed = await poller.claim_batch()

        # Should only claim the consolidation for the other bank
        assert len(claimed) == 1
        assert claimed[0].operation_id == str(other_op_id)
        assert claimed[0].task_dict["bank_id"] == other_bank_id

        # Verify the pending consolidation for first bank is still pending
        row = await pool.fetchrow(
            "SELECT status, worker_id FROM async_operations WHERE operation_id = $1",
            pending_op_id,
        )
        assert row["status"] == "pending"
        assert row["worker_id"] is None

    @pytest.mark.asyncio
    async def test_claim_batch_allows_non_consolidation_when_consolidation_processing(self, pool, clean_operations):
        """Test that non-consolidation tasks are still claimed even if consolidation is processing."""
        from hindsight_api.worker import WorkerPoller

        bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
        await _ensure_bank(pool, bank_id)

        # Create a processing consolidation for bank
        await pool.execute(
            """
            INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id)
            VALUES ($1, $2, 'consolidation', 'processing', $3::jsonb, 'other-worker')
            """,
            uuid.uuid4(),
            bank_id,
            json.dumps({"type": "consolidation", "bank_id": bank_id}),
        )

        # Create a pending retain task for same bank (should be claimed)
        retain_op_id = uuid.uuid4()
        await pool.execute(
            """
            INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
            VALUES ($1, $2, 'retain', 'pending', $3::jsonb)
            """,
            retain_op_id,
            bank_id,
            json.dumps({"type": "batch_retain", "bank_id": bank_id}),
        )

        poller = WorkerPoller(
            pool=pool,
            worker_id="test-worker-1",
            executor=lambda x: None,
        )

        claimed = await poller.claim_batch()

        # Should claim the retain task (non-consolidation tasks are unaffected)
        assert len(claimed) == 1
        assert claimed[0].operation_id == str(retain_op_id)


class TestWorkerRecovery:
    """Tests for worker task recovery on startup."""

    @pytest.mark.asyncio
    async def test_recover_own_tasks_resets_processing_to_pending(self, pool, clean_operations):
        """Test that recover_own_tasks resets processing tasks back to pending."""
        from hindsight_api.worker import WorkerPoller

        # Create tasks that were being processed by this worker (simulating a crash)
        bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
        await _ensure_bank(pool, bank_id)
        worker_id = "crashed-worker"
        task_ids = []

        for i in range(3):
            op_id = uuid.uuid4()
            task_ids.append(op_id)
            payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id})
            await pool.execute(
                """
                INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id, claimed_at)
                VALUES ($1, $2, 'test', 'processing', $3::jsonb, $4, now())
                """,
                op_id,
                bank_id,
                payload,
                worker_id,
            )

        # Create poller with same worker_id and call recover
        poller = WorkerPoller(
            pool=pool,
            worker_id=worker_id,
            executor=lambda x: None,
        )

        recovered_count = await poller.recover_own_tasks()
        assert recovered_count == 3

        # Verify all tasks are back to pending with no worker assigned
        rows = await pool.fetch(
            "SELECT status, worker_id, claimed_at FROM async_operations WHERE bank_id = $1",
            bank_id,
        )
        for row in rows:
            assert row["status"] == "pending"
            assert row["worker_id"] is None
            assert row["claimed_at"] is None

    @pytest.mark.asyncio
    async def test_recover_own_tasks_does_not_affect_other_workers(self, pool, clean_operations):
        """Test that recover_own_tasks only affects tasks from the same worker_id."""
        from hindsight_api.worker import WorkerPoller

        bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
        await _ensure_bank(pool, bank_id)

        # Create tasks for worker-1 (the one that will recover)
        for i in range(2):
            op_id = uuid.uuid4()
            payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id})
            await pool.execute(
                """
                INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id)
                VALUES ($1, $2, 'test', 'processing', $3::jsonb, 'worker-1')
                """,
                op_id,
                bank_id,
                payload,
            )

        # Create tasks for worker-2 (should not be affected)
        for i in range(2):
            op_id = uuid.uuid4()
            payload = json.dumps({"type": "test_task", "index": i + 10, "bank_id": bank_id})
            await pool.execute(
                """
                INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id)
                VALUES ($1, $2, 'test', 'processing', $3::jsonb, 'worker-2')
                """,
                op_id,
                bank_id,
                payload,
            )

        # Worker-1 recovers its tasks
        poller = WorkerPoller(
            pool=pool,
            worker_id="worker-1",
            executor=lambda x: None,
        )

        recovered_count = await poller.recover_own_tasks()
        assert recovered_count == 2

        # Verify worker-1 tasks are released
        worker1_rows = await pool.fetch(
            "SELECT status, worker_id FROM async_operations WHERE bank_id = $1 AND worker_id IS NULL",
            bank_id,
        )
        assert len(worker1_rows) == 2

        # Verify worker-2 tasks are unaffected
        worker2_rows = await pool.fetch(
            "SELECT status, worker_id FROM async_operations WHERE bank_id = $1 AND worker_id = 'worker-2'",
            bank_id,
        )
        assert len(worker2_rows) == 2
        for row in worker2_rows:
            assert row["status"] == "processing"

    @pytest.mark.asyncio
    async def test_recover_own_tasks_returns_zero_when_no_stale_tasks(self, pool, clean_operations):
        """Test that recover_own_tasks returns 0 when there are no stale tasks."""
        from hindsight_api.worker import WorkerPoller

        poller = WorkerPoller(
            pool=pool,
            worker_id="fresh-worker",
            executor=lambda x: None,
        )

        recovered_count = await poller.recover_own_tasks()
        assert recovered_count == 0


class TestConcurrentWorkers:
    """Tests for concurrent worker task claiming (FOR UPDATE SKIP LOCKED)."""

    @pytest.mark.asyncio
    async def test_concurrent_workers_claim_different_tasks(self, pool, clean_operations):
        """Test that multiple workers claim different tasks (no duplicates)."""
        from hindsight_api.worker import WorkerPoller

        # Create 10 pending tasks
        bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
        await _ensure_bank(pool, bank_id)
        task_ids = []
        for i in range(10):
            op_id = uuid.uuid4()
            task_ids.append(op_id)
            payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id, "operation_id": str(op_id)})
            await pool.execute(
                """
                INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
                VALUES ($1, $2, 'test', 'pending', $3::jsonb)
                """,
                op_id,
                bank_id,
                payload,
            )

        # Create 3 workers that will claim tasks concurrently
        workers_claimed: dict[str, list[str]] = {"worker-1": [], "worker-2": [], "worker-3": []}

        async def claim_for_worker(worker_id: str):
            poller = WorkerPoller(
                pool=pool,
                worker_id=worker_id,
                executor=lambda x: None,
            )
            claimed = await poller.claim_batch()
            workers_claimed[worker_id] = [task.operation_id for task in claimed]

        # Run all workers concurrently
        await asyncio.gather(
            claim_for_worker("worker-1"),
            claim_for_worker("worker-2"),
            claim_for_worker("worker-3"),
        )

        # Verify no duplicates - each task claimed by exactly one worker
        all_claimed = workers_claimed["worker-1"] + workers_claimed["worker-2"] + workers_claimed["worker-3"]
        assert len(all_claimed) == len(set(all_claimed)), "Duplicate task claimed by multiple workers!"

        # Verify total claimed equals available tasks (10)
        assert len(all_claimed) == 10, f"Expected 10 tasks claimed, got {len(all_claimed)}"

        # Verify each task is assigned to exactly one worker in DB
        rows = await pool.fetch(
            "SELECT operation_id, worker_id FROM async_operations WHERE bank_id = $1",
            bank_id,
        )
        worker_assignments = {str(row["operation_id"]): row["worker_id"] for row in rows}

        # With FOR UPDATE SKIP LOCKED, it's a race condition which workers get tasks.
        # The important invariant is no duplicates and all tasks claimed, which we verified above.
        # Just verify that at least 1 worker got tasks and all tasks have a worker assigned.
        assert len(set(worker_assignments.values())) >= 1, "At least one worker should have claimed tasks"
        assert all(w is not None for w in worker_assignments.values()), "All tasks should have a worker assigned"

    @pytest.mark.asyncio
    async def test_workers_do_not_claim_already_processing_tasks(self, pool, clean_operations):
        """Test that workers skip tasks already being processed by another worker."""
        from hindsight_api.worker import WorkerPoller

        # Create tasks - some pending, some already processing
        bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
        await _ensure_bank(pool, bank_id)

        # Create 3 pending tasks
        for i in range(3):
            op_id = uuid.uuid4()
            payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id})
            await pool.execute(
                """
                INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
                VALUES ($1, $2, 'test', 'pending', $3::jsonb)
                """,
                op_id,
                bank_id,
                payload,
            )

        # Create 2 already-processing tasks owned by another worker
        for i in range(2):
            op_id = uuid.uuid4()
            payload = json.dumps({"type": "test_task", "index": i + 10, "bank_id": bank_id})
            await pool.execute(
                """
                INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id)
                VALUES ($1, $2, 'test', 'processing', $3::jsonb, 'other-worker')
                """,
                op_id,
                bank_id,
                payload,
            )

        # New worker should only claim the 3 pending tasks
        poller = WorkerPoller(
            pool=pool,
            worker_id="new-worker",
            executor=lambda x: None,
        )

        claimed = await poller.claim_batch()
        assert len(claimed) == 3, "Worker should only claim pending tasks"

        # Verify other worker's tasks are still owned by them
        row = await pool.fetchrow(
            "SELECT COUNT(*) as count FROM async_operations WHERE bank_id = $1 AND worker_id = 'other-worker'",
            bank_id,
        )
        assert row["count"] == 2


class TestWorkerDecommission:
    """Tests for worker decommissioning functionality."""

    @pytest.mark.asyncio
    async def test_decommission_releases_worker_tasks(self, pool, clean_operations):
        """Test that decommissioning a worker releases all its processing tasks."""
        # Create tasks being processed by a worker
        bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
        await _ensure_bank(pool, bank_id)
        worker_id = "worker-to-decommission"

        for i in range(5):
            op_id = uuid.uuid4()
            payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id})
            await pool.execute(
                """
                INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id, claimed_at)
                VALUES ($1, $2, 'test', 'processing', $3::jsonb, $4, now())
                """,
                op_id,
                bank_id,
                payload,
                worker_id,
            )

        # Run decommission
        result = await pool.fetch(
            """
            UPDATE async_operations
            SET status = 'pending', worker_id = NULL, claimed_at = NULL, updated_at = now()
            WHERE worker_id = $1 AND status = 'processing'
            RETURNING operation_id
            """,
            worker_id,
        )

        assert len(result) == 5

        # Verify all tasks are back to pending
        rows = await pool.fetch(
            "SELECT status, worker_id, claimed_at FROM async_operations WHERE bank_id = $1",
            bank_id,
        )
        for row in rows:
            assert row["status"] == "pending"
            assert row["worker_id"] is None
            assert row["claimed_at"] is None

    @pytest.mark.asyncio
    async def test_decommission_does_not_affect_other_workers(self, pool, clean_operations):
        """Test that decommissioning one worker doesn't affect another worker's tasks."""
        bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
        await _ensure_bank(pool, bank_id)

        # Create tasks for worker-1
        for i in range(3):
            op_id = uuid.uuid4()
            payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id})
            await pool.execute(
                """
                INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id)
                VALUES ($1, $2, 'test', 'processing', $3::jsonb, 'worker-1')
                """,
                op_id,
                bank_id,
                payload,
            )

        # Create tasks for worker-2
        for i in range(3):
            op_id = uuid.uuid4()
            payload = json.dumps({"type": "test_task", "index": i + 10, "bank_id": bank_id})
            await pool.execute(
                """
                INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload, worker_id)
                VALUES ($1, $2, 'test', 'processing', $3::jsonb, 'worker-2')
                """,
                op_id,
                bank_id,
                payload,
            )

        # Decommission worker-1 only
        await pool.execute(
            """
            UPDATE async_operations
            SET status = 'pending', worker_id = NULL, claimed_at = NULL
            WHERE worker_id = 'worker-1' AND status = 'processing'
            """
        )

        # Verify worker-1 tasks are released
        worker1_rows = await pool.fetch(
            "SELECT status, worker_id FROM async_operations WHERE bank_id = $1 AND worker_id IS NULL",
            bank_id,
        )
        assert len(worker1_rows) == 3

        # Verify worker-2 tasks are unaffected
        worker2_rows = await pool.fetch(
            "SELECT status, worker_id FROM async_operations WHERE bank_id = $1 AND worker_id = 'worker-2'",
            bank_id,
        )
        assert len(worker2_rows) == 3
        for row in worker2_rows:
            assert row["status"] == "processing"


class TestSyncTaskBackend:
    """Tests for SyncTaskBackend (used in tests and embedded mode)."""

    @pytest.mark.asyncio
    async def test_sync_backend_executes_immediately(self):
        """Test that SyncTaskBackend executes tasks immediately."""
        executed = []

        async def mock_executor(task_dict):
            executed.append(task_dict)

        backend = SyncTaskBackend()
        backend.set_executor(mock_executor)
        await backend.initialize()

        task_dict = {"type": "test", "data": "value"}
        await backend.submit_task(task_dict)

        assert len(executed) == 1
        assert executed[0] == task_dict

    @pytest.mark.asyncio
    async def test_sync_backend_handles_errors(self):
        """Test that SyncTaskBackend handles executor errors gracefully."""

        async def failing_executor(task_dict):
            raise ValueError("Test error")

        backend = SyncTaskBackend()
        backend.set_executor(failing_executor)
        await backend.initialize()

        # Should not raise, error is logged
        await backend.submit_task({"type": "test"})


class TestDynamicTenantDiscovery:
    """Tests for dynamic tenant discovery via TenantExtension."""

    @pytest.mark.asyncio
    async def test_poller_discovers_tenants_dynamically(self, pool, clean_operations):
        """Test that poller calls list_tenants() on each poll cycle."""
        from hindsight_api.extensions.tenant import Tenant, TenantExtension
        from hindsight_api.worker import WorkerPoller

        # Create a mock tenant extension that tracks calls
        class MockTenantExtension(TenantExtension):
            def __init__(self):
                self.list_tenants_calls = 0
                self.tenants_to_return: list[Tenant] = [Tenant(schema="public")]

            async def authenticate(self, context):
                raise NotImplementedError("Not used in this test")

            async def list_tenants(self) -> list[Tenant]:
                self.list_tenants_calls += 1
                return self.tenants_to_return

        mock_extension = MockTenantExtension()

        # Create pending tasks in public schema
        bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
        await _ensure_bank(pool, bank_id)
        for i in range(2):
            op_id = uuid.uuid4()
            payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id})
            await pool.execute(
                """
                INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
                VALUES ($1, $2, 'test', 'pending', $3::jsonb)
                """,
                op_id,
                bank_id,
                payload,
            )

        poller = WorkerPoller(
            pool=pool,
            worker_id="test-worker-1",
            executor=lambda x: None,
            tenant_extension=mock_extension,
        )

        # First claim_batch should call list_tenants
        claimed1 = await poller.claim_batch()
        assert mock_extension.list_tenants_calls == 1
        assert len(claimed1) == 2

        # Add more tasks
        for i in range(2):
            op_id = uuid.uuid4()
            payload = json.dumps({"type": "test_task", "index": i + 10, "bank_id": bank_id})
            await pool.execute(
                """
                INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
                VALUES ($1, $2, 'test', 'pending', $3::jsonb)
                """,
                op_id,
                bank_id,
                payload,
            )

        # Second claim_batch should call list_tenants again
        claimed2 = await poller.claim_batch()
        assert mock_extension.list_tenants_calls == 2
        assert len(claimed2) == 2

    @pytest.mark.asyncio
    async def test_poller_picks_up_new_tenants_without_restart(self, pool, clean_operations):
        """Test that new tenants are discovered on subsequent poll cycles."""
        from hindsight_api.extensions.tenant import Tenant, TenantExtension
        from hindsight_api.worker import WorkerPoller

        class DynamicTenantExtension(TenantExtension):
            def __init__(self):
                # Start with just public
                self.tenants: list[Tenant] = [Tenant(schema="public")]
                self.list_tenants_calls = 0

            async def authenticate(self, context):
                raise NotImplementedError("Not used in this test")

            async def list_tenants(self) -> list[Tenant]:
                self.list_tenants_calls += 1
                return self.tenants

        dynamic_extension = DynamicTenantExtension()

        # Create a task in public schema
        bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
        await _ensure_bank(pool, bank_id)
        op_id = uuid.uuid4()
        payload = json.dumps({"type": "test_task", "bank_id": bank_id})
        await pool.execute(
            """
            INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
            VALUES ($1, $2, 'test', 'pending', $3::jsonb)
            """,
            op_id,
            bank_id,
            payload,
        )

        poller = WorkerPoller(
            pool=pool,
            worker_id="test-worker-1",
            executor=lambda x: None,
            tenant_extension=dynamic_extension,
        )

        # First poll - only public schema
        claimed1 = await poller.claim_batch()
        assert len(claimed1) == 1
        assert claimed1[0].schema is None  # public is represented as None
        assert dynamic_extension.list_tenants_calls == 1

        # Simulate tenant list changing (but we won't add a non-existent schema)
        # In real world, the schema would be created before list_tenants returns it
        # Here we just verify that list_tenants is called again

        # Add another task to public
        op_id2 = uuid.uuid4()
        payload2 = json.dumps({"type": "test_task", "bank_id": bank_id})
        await pool.execute(
            """
            INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
            VALUES ($1, $2, 'test', 'pending', $3::jsonb)
            """,
            op_id2,
            bank_id,
            payload2,
        )

        # Second poll - list_tenants should be called again
        claimed2 = await poller.claim_batch()
        assert len(claimed2) == 1
        assert dynamic_extension.list_tenants_calls == 2  # Called again on second poll

        # Third poll with no tasks - still calls list_tenants
        claimed3 = await poller.claim_batch()
        assert len(claimed3) == 0
        assert dynamic_extension.list_tenants_calls == 3  # Called again even with no tasks

    @pytest.mark.asyncio
    async def test_poller_without_tenant_extension_uses_public(self, pool, clean_operations):
        """Test that poller uses public schema when no tenant extension is configured."""
        from hindsight_api.worker import WorkerPoller

        # Create pending tasks
        bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
        await _ensure_bank(pool, bank_id)
        for i in range(3):
            op_id = uuid.uuid4()
            payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id})
            await pool.execute(
                """
                INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
                VALUES ($1, $2, 'test', 'pending', $3::jsonb)
                """,
                op_id,
                bank_id,
                payload,
            )

        # No tenant_extension provided
        poller = WorkerPoller(
            pool=pool,
            worker_id="test-worker-1",
            executor=lambda x: None,
        )

        claimed = await poller.claim_batch()
        assert len(claimed) == 3

        # All tasks should have schema=None (public)
        for task in claimed:
            assert task.schema is None

    @pytest.mark.asyncio
    async def test_poller_with_custom_schema(self, pool):
        """Test that poller uses custom schema when schema parameter is provided."""
        from hindsight_api.worker import WorkerPoller

        # Create a custom schema for testing
        test_schema = "test_custom_schema"

        try:
            # Create schema and copy table structure
            await pool.execute(f'CREATE SCHEMA IF NOT EXISTS "{test_schema}"')
            await pool.execute(
                f"""
                CREATE TABLE "{test_schema}".async_operations (
                    LIKE public.async_operations INCLUDING ALL
                )
                """
            )

            # Create pending tasks in the custom schema
            bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
            task_ids = []
            for i in range(3):
                op_id = uuid.uuid4()
                task_ids.append(str(op_id))
                payload = json.dumps({"type": "test_task", "index": i, "bank_id": bank_id})
                await pool.execute(
                    f"""
                    INSERT INTO "{test_schema}".async_operations (operation_id, bank_id, operation_type, status, task_payload)
                    VALUES ($1, $2, 'test', 'pending', $3::jsonb)
                    """,
                    op_id,
                    bank_id,
                    payload,
                )

            # Create poller with custom schema
            poller = WorkerPoller(
                pool=pool,
                worker_id="test-worker-custom-schema",
                executor=lambda x: None,
                schema=test_schema,
            )

            # Claim tasks
            claimed = await poller.claim_batch()
            assert len(claimed) == 3, f"Expected 3 tasks, got {len(claimed)}"

            # All tasks should have schema=test_schema
            claimed_ids = []
            for task in claimed:
                assert task.schema == test_schema, f"Expected schema '{test_schema}', got '{task.schema}'"
                claimed_ids.append(task.operation_id)

            # Verify claimed tasks match what we inserted
            assert set(claimed_ids) == set(task_ids)

            # Verify tasks are marked as processing in the custom schema
            rows = await pool.fetch(
                f"""
                SELECT operation_id, status, worker_id
                FROM "{test_schema}".async_operations
                WHERE operation_id = ANY($1)
                """,
                [uuid.UUID(tid) for tid in task_ids],
            )
            assert len(rows) == 3
            for row in rows:
                assert row["status"] == "processing"
                assert row["worker_id"] == "test-worker-custom-schema"

        finally:
            # Clean up: drop the custom schema
            await pool.execute(f'DROP SCHEMA IF EXISTS "{test_schema}" CASCADE')


async def test_worker_fire_and_forget_nonblocking(pool, clean_operations):
    """
    Test that worker continues polling while tasks run (fire-and-forget pattern).

    This test verifies the FIX: With the old blocking behavior, the worker would
    wait for all tasks in a batch to complete before claiming more. This test
    would FAIL with the old code because tasks 3-4 wouldn't be claimed until
    tasks 1-2 complete. With fire-and-forget, tasks 3-4 are claimed immediately.
    """
    from hindsight_api.worker.poller import WorkerPoller

    task_started = {}  # operation_id -> Event (set when task starts)
    task_canfinish = {}  # operation_id -> Event (wait before finishing)

    async def blocking_executor(task_dict: dict):
        op_id = task_dict["operation_id"]
        # Signal that this task has started
        started = asyncio.Event()
        task_started[op_id] = started
        started.set()

        # Block until we're told to finish
        finish = asyncio.Event()
        task_canfinish[op_id] = finish
        await finish.wait()

    poller = WorkerPoller(
        pool=pool,
        worker_id="test-worker",
        executor=blocking_executor,
        poll_interval_ms=50,  # Fast polling
        max_slots=10,
        consolidation_max_slots=2,
    )

    bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
    await _ensure_bank(pool, bank_id)

    # Submit initial 2 tasks
    task_ids = []
    for i in range(2):
        op_id = uuid.uuid4()
        task_ids.append(str(op_id))
        payload = json.dumps(
            {"type": "test", "operation_type": "retain", "operation_id": str(op_id), "bank_id": bank_id}
        )
        await pool.execute(
            """
            INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
            VALUES ($1, $2, 'retain', 'pending', $3::jsonb)
            """,
            op_id,
            bank_id,
            payload,
        )

    poll_task = asyncio.create_task(poller.run())

    try:
        # Wait for first 2 tasks to start executing (but not finish)
        for i in range(100):  # Try for up to 1 second
            if len(task_started) >= 2:
                break
            await asyncio.sleep(0.01)
        assert len(task_started) == 2, f"Expected 2 tasks started, got {len(task_started)}"

        # Verify tasks are in_flight
        async with poller._in_flight_lock:
            assert poller._in_flight_count == 2

        # NOW submit 2 more tasks WHILE the first 2 are still running
        for i in range(2):
            op_id = uuid.uuid4()
            task_ids.append(str(op_id))
            payload = json.dumps(
                {"type": "test", "operation_type": "retain", "operation_id": str(op_id), "bank_id": bank_id}
            )
            await pool.execute(
                """
                INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
                VALUES ($1, $2, 'retain', 'pending', $3::jsonb)
                """,
                op_id,
                bank_id,
                payload,
            )

        # KEY ASSERTION: Worker should claim tasks 3-4 WITHOUT waiting for 1-2 to finish
        # This would FAIL with the old blocking behavior
        for i in range(100):  # Try for up to 1 second
            if len(task_started) >= 4:
                break
            await asyncio.sleep(0.01)

        assert len(task_started) == 4, (
            f"Fire-and-forget FAILED: Expected 4 tasks started, got {len(task_started)}. "
            "This means the worker blocked waiting for the first batch to complete."
        )

        # Verify all 4 tasks are in-flight
        async with poller._in_flight_lock:
            assert poller._in_flight_count == 4

        # Clean up: allow all tasks to finish
        for event in task_canfinish.values():
            event.set()

    finally:
        # Ensure cleanup
        for event in task_canfinish.values():
            event.set()
        await poller.shutdown_graceful(timeout=2.0)
        try:
            await asyncio.wait_for(poll_task, timeout=1.0)
        except asyncio.CancelledError:
            pass


async def test_worker_slot_limits_enforced(pool, clean_operations):
    """Test that worker respects max_slots and won't exceed the limit."""
    from hindsight_api.worker.poller import WorkerPoller

    tasks_started = set()
    task_events = {}

    async def controlled_executor(task_dict: dict):
        op_id = task_dict["operation_id"]
        tasks_started.add(op_id)
        event = asyncio.Event()
        task_events[op_id] = event
        await event.wait()

    poller = WorkerPoller(
        pool=pool,
        worker_id="test-worker",
        executor=controlled_executor,
        poll_interval_ms=50,
        max_slots=3,  # Only allow 3 concurrent tasks
        consolidation_max_slots=1,
    )

    # Submit 10 tasks
    bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
    await _ensure_bank(pool, bank_id)
    for i in range(10):
        op_id = uuid.uuid4()
        payload = json.dumps(
            {"type": "test", "operation_type": "retain", "operation_id": str(op_id), "bank_id": bank_id}
        )
        await pool.execute(
            """
            INSERT INTO async_operations (operation_id, bank_id, operation_type, status, task_payload)
            VALUES ($1, $2, 'retain', 'pending', $3::jsonb)
            """,
            op_id,
            bank_id,
            payload,
        )

    poll_task = asyncio.create_task(poller.run())

    try:
        # Wait for slots to fill
        for i in range(100):
            if len(tasks_started) >= 3:
                break
            await asyncio.sleep(0.01)

        # Should have claimed exactly 3 tasks (slot limit)
        assert len(tasks_started) == 3

        # Wait to ensure no additional tasks are claimed
        for i in range(30):
            await asyncio.sleep(0.01)
        assert len(tasks_started) == 3, "Worker exceeded slot limit!"

        # Release tasks one by one and verify remaining are claimed
        completed = 0
        while completed < 10 and len(tasks_started) < 10:
            # Release the next batch
            events_to_release = list(task_events.values())[completed : completed + 3]
            for event in events_to_release:
                event.set()
            completed += len(events_to_release)

            # Wait for new tasks to be claimed
            for i in range(100):
                if len(tasks_started) >= min(completed + 3, 10):
                    break
                await asyncio.sleep(0.01)

        assert len(tasks_started) == 10

    finally:
        for event in task_events.values():
            event.set()
        await poller.shutdown_graceful(timeout=2.0)
        try:
            await asyncio.wait_for(poll_task, timeout=1.0)
        except asyncio.CancelledError:
            pass


class TestMarkFailedParentPropagation:
    """Tests for _mark_failed parent propagation in WorkerPoller.

    When a child retain operation fails via an unhandled exception, the memory
    engine's transaction is rolled back entirely — including any call to
    _maybe_update_parent_operation inside the engine. The poller's fallback
    _mark_failed must detect this and finalise the parent batch_retain itself.
    """

    async def _insert_op(
        self,
        pool,
        *,
        op_id: "uuid.UUID",
        bank_id: str,
        operation_type: str,
        status: str,
        result_metadata: dict | None = None,
    ) -> None:
        meta_json = json.dumps(result_metadata if result_metadata is not None else {})
        await pool.execute(
            """
            INSERT INTO async_operations
                (operation_id, bank_id, operation_type, status, result_metadata)
            VALUES ($1, $2, $3, $4, $5::jsonb)
            """,
            op_id,
            bank_id,
            operation_type,
            status,
            meta_json,
        )

    @pytest.mark.asyncio
    async def test_mark_failed_finalises_parent_when_last_sibling_fails(self, pool, clean_operations):
        """When the last pending child fails, parent batch_retain is marked failed."""
        from hindsight_api.worker import WorkerPoller

        bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
        await _ensure_bank(pool, bank_id)

        parent_id = uuid.uuid4()
        child1_id = uuid.uuid4()
        child2_id = uuid.uuid4()

        # Parent batch_retain still pending
        await self._insert_op(pool, op_id=parent_id, bank_id=bank_id, operation_type="batch_retain", status="pending")

        # child1 already completed
        await self._insert_op(
            pool,
            op_id=child1_id,
            bank_id=bank_id,
            operation_type="retain",
            status="completed",
            result_metadata={"parent_operation_id": str(parent_id)},
        )

        # child2 still processing — this is the one that will fail
        await self._insert_op(
            pool,
            op_id=child2_id,
            bank_id=bank_id,
            operation_type="retain",
            status="processing",
            result_metadata={"parent_operation_id": str(parent_id)},
        )

        poller = WorkerPoller(pool=pool, worker_id="test-worker-1", executor=lambda x: None)
        await poller._mark_failed(str(child2_id), "DB constraint violation", schema=None)

        # child2 must be failed
        child2_row = await pool.fetchrow(
            "SELECT status, error_message FROM async_operations WHERE operation_id = $1", child2_id
        )
        assert child2_row["status"] == "failed"
        assert "DB constraint violation" in child2_row["error_message"]

        # parent must now be failed (all siblings done, at least one failed)
        parent_row = await pool.fetchrow("SELECT status FROM async_operations WHERE operation_id = $1", parent_id)
        assert parent_row["status"] == "failed", (
            f"Parent should be 'failed' when last sibling fails, got '{parent_row['status']}'"
        )

    @pytest.mark.asyncio
    async def test_mark_failed_finalises_parent_when_last_sibling_is_sole_child(self, pool, clean_operations):
        """When the only child fails, parent batch_retain becomes failed."""
        from hindsight_api.worker import WorkerPoller

        bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
        await _ensure_bank(pool, bank_id)

        parent_id = uuid.uuid4()
        child_id = uuid.uuid4()

        await self._insert_op(pool, op_id=parent_id, bank_id=bank_id, operation_type="batch_retain", status="pending")
        await self._insert_op(
            pool,
            op_id=child_id,
            bank_id=bank_id,
            operation_type="retain",
            status="processing",
            result_metadata={"parent_operation_id": str(parent_id)},
        )

        poller = WorkerPoller(pool=pool, worker_id="test-worker-1", executor=lambda x: None)
        await poller._mark_failed(str(child_id), "unexpected error", schema=None)

        parent_row = await pool.fetchrow("SELECT status FROM async_operations WHERE operation_id = $1", parent_id)
        assert parent_row["status"] == "failed"

    @pytest.mark.asyncio
    async def test_mark_failed_does_not_finalise_parent_when_siblings_still_pending(self, pool, clean_operations):
        """Parent is NOT updated while other siblings are still processing/pending."""
        from hindsight_api.worker import WorkerPoller

        bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
        await _ensure_bank(pool, bank_id)

        parent_id = uuid.uuid4()
        child1_id = uuid.uuid4()
        child2_id = uuid.uuid4()

        await self._insert_op(pool, op_id=parent_id, bank_id=bank_id, operation_type="batch_retain", status="pending")

        # child1 is the one failing
        await self._insert_op(
            pool,
            op_id=child1_id,
            bank_id=bank_id,
            operation_type="retain",
            status="processing",
            result_metadata={"parent_operation_id": str(parent_id)},
        )
        # child2 is still pending — not done yet
        await self._insert_op(
            pool,
            op_id=child2_id,
            bank_id=bank_id,
            operation_type="retain",
            status="pending",
            result_metadata={"parent_operation_id": str(parent_id)},
        )

        poller = WorkerPoller(pool=pool, worker_id="test-worker-1", executor=lambda x: None)
        await poller._mark_failed(str(child1_id), "early failure", schema=None)

        # child1 is failed
        child1_row = await pool.fetchrow("SELECT status FROM async_operations WHERE operation_id = $1", child1_id)
        assert child1_row["status"] == "failed"

        # parent must still be pending (child2 not done)
        parent_row = await pool.fetchrow("SELECT status FROM async_operations WHERE operation_id = $1", parent_id)
        assert parent_row["status"] == "pending", (
            f"Parent should remain 'pending' while siblings are outstanding, got '{parent_row['status']}'"
        )

    @pytest.mark.asyncio
    async def test_mark_failed_no_parent_is_safe(self, pool, clean_operations):
        """Operations without a parent (no result_metadata parent_operation_id) fail cleanly."""
        from hindsight_api.worker import WorkerPoller

        bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
        await _ensure_bank(pool, bank_id)

        op_id = uuid.uuid4()
        await self._insert_op(pool, op_id=op_id, bank_id=bank_id, operation_type="retain", status="processing")

        poller = WorkerPoller(pool=pool, worker_id="test-worker-1", executor=lambda x: None)
        # Must not raise
        await poller._mark_failed(str(op_id), "standalone failure", schema=None)

        row = await pool.fetchrow("SELECT status FROM async_operations WHERE operation_id = $1", op_id)
        assert row["status"] == "failed"

    @pytest.mark.asyncio
    async def test_unhandled_exception_via_execute_task_propagates_to_parent(self, pool, clean_operations):
        """End-to-end: executor raises a plain exception, poller calls _mark_failed,
        which then resolves the parent batch_retain to failed."""
        from hindsight_api.worker import WorkerPoller
        from hindsight_api.worker.poller import ClaimedTask

        bank_id = f"test-worker-{uuid.uuid4().hex[:8]}"
        await _ensure_bank(pool, bank_id)

        parent_id = uuid.uuid4()
        child_id = uuid.uuid4()

        await self._insert_op(pool, op_id=parent_id, bank_id=bank_id, operation_type="batch_retain", status="pending")
        await self._insert_op(
            pool,
            op_id=child_id,
            bank_id=bank_id,
            operation_type="retain",
            status="processing",
            result_metadata={"parent_operation_id": str(parent_id)},
        )

        async def crashing_executor(task_dict):
            raise RuntimeError("Simulated DB constraint violation — transaction rolled back")

        poller = WorkerPoller(pool=pool, worker_id="test-worker-1", executor=crashing_executor)

        task_dict = {"type": "retain", "operation_id": str(child_id), "bank_id": bank_id}
        claimed_task = ClaimedTask(operation_id=str(child_id), task_dict=task_dict, schema=None)
        await poller.execute_task(claimed_task)

        completed = await poller.wait_for_active_tasks(timeout=5.0)
        assert completed, "Task did not complete within timeout"

        child_row = await pool.fetchrow("SELECT status FROM async_operations WHERE operation_id = $1", child_id)
        assert child_row["status"] == "failed"

        parent_row = await pool.fetchrow("SELECT status FROM async_operations WHERE operation_id = $1", parent_id)
        assert parent_row["status"] == "failed", (
            f"Parent batch_retain should be 'failed' after child fails via unhandled exception, "
            f"got '{parent_row['status']}'"
        )