#!/usr/bin/env python3
"""End-to-end verification for the Tier 1 durability work.

Two tiers of coverage:

**Service-layer scenarios (s1-s10)** — exercise the WAL + reconciler +
idempotency + token-usage primitives without spinning up Flask, Celery
workers, or a live LLM:

* Service-level calls (e.g. ``ConversationService.save_user_question``)
  drive the same code paths the routes use.
* Flask ``test_request_context`` lets us hit ``Resource.post()`` directly
  for HTTP idempotency checks.
* The reconciler is invoked synchronously via ``run_reconciliation()``.
* Stuck-row scenarios are simulated by backdating ``timestamp`` /
  ``attempted_at`` columns so the reconciler treats them as past the
  threshold without having to wait the real 5/15 minutes.

**Live scenarios (s11-s13)** — boot the local mock LLM stub from
``scripts/e2e/mock_llm.py`` as a subprocess and drive real network
streams + real Celery workers:

* s11: real OpenAI-protocol stream through WAL + finalize + token_usage.
* s12: full ingest of a synthesised markdown file (chunker + embeddings
  + faiss). The final ``upload_index`` self-call to Flask is patched
  out — everything before that (chunk progress, source id, vector store
  on disk) runs for real.
* s13: ``kill -9`` a Celery worker subprocess mid-task; verify
  ``acks_late`` + ``visibility_timeout`` make the broker redeliver to a
  fresh worker. Uses isolated Redis DBs (11/12) and a unique queue so
  it can't interfere with production work on the same broker.

Usage::

    .venv/bin/python scripts/qa/durability_e2e.py
    .venv/bin/python scripts/qa/durability_e2e.py --list
    .venv/bin/python scripts/qa/durability_e2e.py --only s1,s5
    .venv/bin/python scripts/qa/durability_e2e.py --only s11,s12,s13
"""

from __future__ import annotations

import argparse
import logging
import os
import sys
import time
import traceback
import uuid
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from pathlib import Path
from unittest.mock import MagicMock, patch


# ---------------------------------------------------------------------------
# Bootstrap: import application/* and connect to the configured Postgres.
# ---------------------------------------------------------------------------

ROOT = Path(__file__).resolve().parent.parent.parent
sys.path.insert(0, str(ROOT))

# Suppress AUTO_MIGRATE / AUTO_CREATE_DB side-effects from app boot — this
# script only reads/writes data, the schema is assumed to be at head.
os.environ.setdefault("AUTO_MIGRATE", "false")
os.environ.setdefault("AUTO_CREATE_DB", "false")

from dotenv import load_dotenv  # noqa: E402

load_dotenv(ROOT / ".env")

from sqlalchemy import create_engine, text  # noqa: E402

POSTGRES_URI = os.environ.get("POSTGRES_URI", "")
if not POSTGRES_URI:
    print("POSTGRES_URI not set; cannot run E2E checks.", file=sys.stderr)
    sys.exit(2)
if POSTGRES_URI.startswith("postgresql://"):
    POSTGRES_URI = POSTGRES_URI.replace("postgresql://", "postgresql+psycopg://", 1)

ENGINE = create_engine(POSTGRES_URI, pool_pre_ping=True)
UNIQUE = f"qa_e2e_{int(time.time())}_{uuid.uuid4().hex[:6]}"


# Quiet noisy library loggers so the script's own output stays readable.
for name in (
    "application", "application.api", "application.storage",
    "application.usage", "application.parser",
    "application.api.user.reconciliation",
):
    logging.getLogger(name).setLevel(logging.ERROR)


# ---------------------------------------------------------------------------
# Output helpers
# ---------------------------------------------------------------------------

GREEN = "\033[32m"
RED = "\033[31m"
YELLOW = "\033[33m"
DIM = "\033[2m"
RESET = "\033[0m"


def header(msg: str) -> None:
    print(f"\n{'=' * 72}\n{msg}\n{'=' * 72}")


def step(msg: str) -> None:
    print(f"  {GREEN}✓{RESET} {msg}")


def info(msg: str) -> None:
    print(f"  {DIM}· {msg}{RESET}")


def fail(msg: str) -> None:
    print(f"  {RED}✗ {msg}{RESET}")


def warn(msg: str) -> None:
    print(f"  {YELLOW}! {msg}{RESET}")


# ---------------------------------------------------------------------------
# Patch helpers — reroute service module-level db_session/db_readonly to
# OUR engine so all writes land on the same Postgres the script is observing.
# ---------------------------------------------------------------------------


@contextmanager
def patch_db_session(*module_paths: str):
    """Reroute ``db_session`` / ``db_readonly`` (whichever each module
    actually imports) at our engine. Modules that import only one of the
    two are handled gracefully — patches are skipped for missing attrs.
    """
    import importlib

    @contextmanager
    def _ses():
        with ENGINE.begin() as conn:
            yield conn

    @contextmanager
    def _ro():
        with ENGINE.connect() as conn:
            yield conn

    patches = []
    for module in module_paths:
        mod = importlib.import_module(module)
        if hasattr(mod, "db_session"):
            patches.append(patch(f"{module}.db_session", _ses))
        if hasattr(mod, "db_readonly"):
            patches.append(patch(f"{module}.db_readonly", _ro))
    for p in patches:
        p.start()
    try:
        yield
    finally:
        for p in patches:
            p.stop()


# ---------------------------------------------------------------------------
# Cleanup helper — run after each scenario regardless of pass/fail.
# ---------------------------------------------------------------------------


def cleanup_artifacts(prefix: str = UNIQUE) -> None:
    with ENGINE.begin() as c:
        c.execute(
            text(
                "DELETE FROM stack_logs "
                "WHERE query LIKE 'reconciler_%' "
                "AND timestamp > now() - interval '15 minutes' "
                "AND (user_id LIKE :p OR query LIKE :p)"
            ),
            {"p": f"%{prefix}%"},
        )
        c.execute(
            text(
                "DELETE FROM tool_call_attempts WHERE call_id LIKE :p"
            ),
            {"p": f"{prefix}-%"},
        )
        c.execute(
            text(
                "DELETE FROM conversation_messages WHERE conversation_id IN ("
                "  SELECT id FROM conversations WHERE user_id LIKE :p"
                ")"
            ),
            {"p": f"{prefix}-%"},
        )
        c.execute(
            text(
                "DELETE FROM pending_tool_state WHERE user_id LIKE :p"
            ),
            {"p": f"{prefix}-%"},
        )
        c.execute(
            text(
                "DELETE FROM conversations WHERE user_id LIKE :p"
            ),
            {"p": f"{prefix}-%"},
        )
        c.execute(
            text(
                "DELETE FROM token_usage WHERE user_id LIKE :p"
            ),
            {"p": f"{prefix}-%"},
        )
        c.execute(
            text(
                "DELETE FROM webhook_dedup WHERE idempotency_key LIKE :p"
            ),
            {"p": f"%:{prefix}-%"},
        )
        c.execute(
            text(
                "DELETE FROM task_dedup WHERE idempotency_key LIKE :p"
            ),
            {"p": f"%:{prefix}-%"},
        )
        c.execute(
            text(
                "DELETE FROM ingest_chunk_progress WHERE source_id IN ("
                "  SELECT id FROM sources WHERE name LIKE :p"
                ")"
            ),
            {"p": f"{prefix}-%"},
        )
        c.execute(
            text(
                "DELETE FROM agents WHERE name LIKE :p"
            ),
            {"p": f"{prefix}-%"},
        )
        c.execute(
            text(
                "DELETE FROM prompts WHERE user_id LIKE :p"
            ),
            {"p": f"{prefix}-%"},
        )
        c.execute(
            text(
                "DELETE FROM users WHERE user_id LIKE :p"
            ),
            {"p": f"{prefix}-%"},
        )


# ---------------------------------------------------------------------------
# Scenarios
# ---------------------------------------------------------------------------


def s1_wal_and_reconciler() -> None:
    """WAL placeholder is reserved before LLM call; reconciler escalates
    a stuck pending row to ``failed`` after 3 attempts and emits a
    structured alert."""

    from application.api.answer.services.conversation_service import (
        ConversationService,
        TERMINATED_RESPONSE_PLACEHOLDER,
    )
    from application.api.user.reconciliation import run_reconciliation

    user = f"{UNIQUE}-s1"
    with ENGINE.begin() as c:
        c.execute(
            text("INSERT INTO users (user_id) VALUES (:u) ON CONFLICT DO NOTHING"),
            {"u": user},
        )

    with patch_db_session("application.api.answer.services.conversation_service"):
        res = ConversationService().save_user_question(
            conversation_id=None,
            question=f"{UNIQUE} s1 stuck WAL",
            decoded_token={"sub": user},
        )
    mid, cid = res["message_id"], res["conversation_id"]
    info(f"reserved message_id={mid} conv={cid}")

    with ENGINE.connect() as c:
        row = c.execute(
            text(
                "SELECT status, response FROM conversation_messages "
                "WHERE id = CAST(:m AS uuid)"
            ),
            {"m": mid},
        ).fetchone()
    assert row[0] == "pending", f"expected pending, got {row[0]!r}"
    assert row[1] == TERMINATED_RESPONSE_PLACEHOLDER, (
        f"placeholder text mismatch: {row[1]!r}"
    )
    step("placeholder reserved with status='pending' + terminated text")

    # Backdate timestamp so the reconciler sweep treats it as stuck.
    with ENGINE.begin() as c:
        c.execute(
            text(
                "UPDATE conversation_messages "
                "SET timestamp = now() - interval '6 minutes' "
                "WHERE id = CAST(:m AS uuid)"
            ),
            {"m": mid},
        )

    # Reconciler MAX_MESSAGE_RECONCILE_ATTEMPTS = 3 → 3 ticks to escalate.
    for i in range(3):
        summary = run_reconciliation()
        info(f"tick {i + 1}: {summary}")

    with ENGINE.connect() as c:
        row = c.execute(
            text(
                "SELECT status, message_metadata->>'reconcile_attempts', "
                "message_metadata->>'error' FROM conversation_messages "
                "WHERE id = CAST(:m AS uuid)"
            ),
            {"m": mid},
        ).fetchone()
    assert row[0] == "failed", f"expected failed, got {row[0]!r}"
    assert int(row[1] or 0) >= 3, f"expected attempts>=3, got {row[1]!r}"
    assert row[2] and "reconciler" in row[2], f"missing reconciler error text: {row[2]!r}"
    step(
        f"reconciler escalated to status='failed' "
        f"after attempts={row[1]} (error: {row[2][:60]}…)"
    )

    with ENGINE.connect() as c:
        alerts = c.execute(
            text(
                "SELECT count(*) FROM stack_logs "
                "WHERE query = 'reconciler_message_failed' "
                "AND timestamp > now() - interval '10 minutes'"
            ),
        ).scalar()
    assert alerts >= 1, f"expected stack_logs alert; got {alerts}"
    step(f"alert written to stack_logs (count={alerts})")


def s2_webhook_idempotency_and_scoping() -> None:
    """Webhook double-POST returns the same task_id; concurrent same-key
    POSTs only enqueue once; two agents sharing a raw key each get their
    own dedup row (cross-tenant scoping)."""

    from flask import Flask

    from application.api.user.agents.webhooks import AgentWebhookListener
    from application.storage.db.repositories.agents import AgentsRepository

    user_a = f"{UNIQUE}-s2a"
    user_b = f"{UNIQUE}-s2b"
    with ENGINE.begin() as c:
        for u in (user_a, user_b):
            c.execute(
                text("INSERT INTO users (user_id) VALUES (:u) ON CONFLICT DO NOTHING"),
                {"u": u},
            )
        agent_a = AgentsRepository(c).create(
            user_a, f"{UNIQUE}-s2-a", "published",
            incoming_webhook_token=f"{UNIQUE}-tok-a",
        )
        agent_b = AgentsRepository(c).create(
            user_b, f"{UNIQUE}-s2-b", "published",
            incoming_webhook_token=f"{UNIQUE}-tok-b",
        )
    info(f"seeded agents a={agent_a['id']} b={agent_b['id']}")

    app = Flask(__name__)

    apply_calls: list[dict] = []

    def _apply_async_side_effect(*_a, **kw):
        apply_calls.append(kw)
        return MagicMock(id=kw.get("task_id") or "auto")

    apply_mock = MagicMock(side_effect=_apply_async_side_effect)

    raw_key = f"{UNIQUE}-key-1"

    # Sequential dedup: same key twice → same task_id, single apply_async call.
    with patch_db_session("application.api.user.agents.webhooks"), patch(
        "application.api.user.agents.webhooks.process_agent_webhook.apply_async",
        apply_mock,
    ):
        for _ in range(2):
            with app.test_request_context(
                f"/api/webhooks/agents/{agent_a['incoming_webhook_token']}",
                method="POST",
                json={"event": "x"},
                headers={"Idempotency-Key": raw_key},
            ):
                listener = AgentWebhookListener()
                resp = listener.post(
                    webhook_token=agent_a["incoming_webhook_token"],
                    agent=agent_a,
                    agent_id_str=str(agent_a["id"]),
                )
                assert resp.status_code == 200
        last_seq = resp.json["task_id"]
    assert apply_mock.call_count == 1, (
        f"sequential same-key: expected 1 apply_async, got {apply_mock.call_count}"
    )
    step(f"sequential same-key dedup: 1 apply_async, task_id={last_seq}")

    # Cross-agent same key — must NOT collide. Two distinct apply_async calls.
    apply_mock.reset_mock()
    apply_calls.clear()
    cross_key = f"{UNIQUE}-cross"
    with patch_db_session("application.api.user.agents.webhooks"), patch(
        "application.api.user.agents.webhooks.process_agent_webhook.apply_async",
        apply_mock,
    ):
        for ag in (agent_a, agent_b):
            with app.test_request_context(
                f"/api/webhooks/agents/{ag['incoming_webhook_token']}",
                method="POST",
                json={"event": "x"},
                headers={"Idempotency-Key": cross_key},
            ):
                listener = AgentWebhookListener()
                resp = listener.post(
                    webhook_token=ag["incoming_webhook_token"],
                    agent=ag,
                    agent_id_str=str(ag["id"]),
                )
                assert resp.status_code == 200
    assert apply_mock.call_count == 2, (
        f"cross-agent same-key should NOT collide; got {apply_mock.call_count}"
    )
    with ENGINE.connect() as c:
        rows = c.execute(
            text(
                "SELECT idempotency_key, agent_id FROM webhook_dedup "
                "WHERE idempotency_key LIKE :p ORDER BY idempotency_key"
            ),
            {"p": f"%:{cross_key}"},
        ).fetchall()
    assert len(rows) == 2, f"expected 2 scoped dedup rows; got {len(rows)}"
    scopes = {str(r[1]) for r in rows}
    assert scopes == {str(agent_a["id"]), str(agent_b["id"])}, (
        f"scopes mismatch: {scopes}"
    )
    step(
        f"cross-agent same key did NOT collide; rows={[r[0] for r in rows]}"
    )


def s3_ingest_deterministic_and_resume() -> None:
    """Deterministic source_id is stable; chunk-progress checkpoint
    resumes from the next un-embedded chunk."""

    from application.parser.embedding_pipeline import _read_resume_index
    from application.storage.db.repositories.ingest_chunk_progress import (
        IngestChunkProgressRepository,
    )
    from application.worker import _derive_source_id

    # Same scoped key → same uuid5; different scope → different uuid5.
    a = _derive_source_id("alice:UPLOAD-1")
    b = _derive_source_id("alice:UPLOAD-1")
    c = _derive_source_id("bob:UPLOAD-1")
    assert a == b, f"deterministic mismatch: {a} vs {b}"
    assert a != c, f"cross-user collision: {a} vs {c}"
    step(f"_derive_source_id: stable for same key, distinct across scopes ({a}, {c})")

    # No key → uuid4 (random per call).
    r1 = _derive_source_id(None)
    r2 = _derive_source_id(None)
    assert r1 != r2, "uuid4 fallback should differ across calls"
    step("uuid4 fallback when no key supplied")

    # Checkpoint primitive: seed progress, observe resume index.
    src_id = uuid.uuid4()
    with patch_db_session("application.parser.embedding_pipeline"):
        with ENGINE.begin() as conn:
            repo = IngestChunkProgressRepository(conn)
            repo.init_progress(str(src_id), total_chunks=10)
            repo.record_chunk(str(src_id), last_index=3, embedded_chunks=4)
        idx = _read_resume_index(str(src_id))
    assert idx == 4, f"expected resume index=4 (next un-embedded); got {idx}"
    step(f"chunk-progress resume index correctly skips already-embedded ({idx}/10)")

    # Cleanup
    with ENGINE.begin() as conn:
        conn.execute(
            text("DELETE FROM ingest_chunk_progress WHERE source_id = CAST(:s AS uuid)"),
            {"s": str(src_id)},
        )


def s4_token_usage_atomic_rollback() -> None:
    """Patch ``update_message_by_id`` to raise after the token_usage
    insert; the surrounding transaction must roll back so no orphan
    token_usage row survives."""

    from application.api.answer.services.conversation_service import (
        ConversationService,
    )
    from application.storage.db.repositories import conversations as conv_mod
    from application.storage.db.repositories.conversations import (
        ConversationsRepository,
    )

    user = f"{UNIQUE}-s4"
    with ENGINE.begin() as c:
        c.execute(
            text("INSERT INTO users (user_id) VALUES (:u) ON CONFLICT DO NOTHING"),
            {"u": user},
        )
        repo = ConversationsRepository(c)
        conv = repo.create(user, f"{UNIQUE}-s4-conv")
        msg = repo.reserve_message(
            str(conv["id"]),
            prompt="q?",
            placeholder_response="...",
            request_id="r1",
            status="streaming",
        )
        mid = str(msg["id"])
    info(f"seeded message_id={mid}")

    original = conv_mod.ConversationsRepository.update_message_by_id

    def explode(self, *_a, **_kw):  # noqa: ANN001
        raise RuntimeError("intentional fault")

    try:
        conv_mod.ConversationsRepository.update_message_by_id = explode
        with patch_db_session(
            "application.api.answer.services.conversation_service",
        ), suppress_logging("application.api.answer.services.conversation_service"):
            try:
                ConversationService().finalize_message(
                    mid,
                    "answer",
                    status="complete",
                    token_usage={"prompt_tokens": 100, "generated_tokens": 50},
                    decoded_token={"sub": user},
                )
            except RuntimeError as exc:
                assert "intentional fault" in str(exc)
            else:
                raise AssertionError("expected RuntimeError to propagate")
    finally:
        conv_mod.ConversationsRepository.update_message_by_id = original

    with ENGINE.connect() as c:
        n = c.execute(
            text("SELECT count(*) FROM token_usage WHERE user_id = :u"),
            {"u": user},
        ).scalar()
    assert n == 0, f"expected 0 token_usage rows after rollback; got {n}"
    step("rollback held: 0 token_usage rows for failed finalize")


@contextmanager
def suppress_logging(*names: str):
    saved = []
    for n in names:
        lg = logging.getLogger(n)
        saved.append((lg, lg.disabled, lg.level))
        lg.disabled = True
    try:
        yield
    finally:
        for lg, was_disabled, was_level in saved:
            lg.disabled = was_disabled
            lg.level = was_level


def s5_reconciler_tool_call_sweeps() -> None:
    """Seed stuck tool_call_attempts rows of each kind; the reconciler
    flips proposed→failed and executed→failed (or compensated when the
    tool exposes a working compensate)."""

    from application.api.user.reconciliation import run_reconciliation

    user = f"{UNIQUE}-s5"
    with ENGINE.begin() as c:
        c.execute(
            text("INSERT INTO users (user_id) VALUES (:u) ON CONFLICT DO NOTHING"),
            {"u": user},
        )

    proposed_id = f"{UNIQUE}-prop-{uuid.uuid4().hex[:6]}"
    executed_id = f"{UNIQUE}-exec-{uuid.uuid4().hex[:6]}"

    with ENGINE.begin() as c:
        c.execute(text("ALTER TABLE tool_call_attempts DISABLE TRIGGER USER"))
        try:
            c.execute(
                text(
                    "INSERT INTO tool_call_attempts "
                    "(call_id, tool_name, action_name, arguments, status, "
                    " attempted_at, updated_at) "
                    "VALUES (:c, 'telegram', 'send_message', '{}'::jsonb, "
                    " 'proposed', now() - interval '6 minutes', "
                    " now() - interval '6 minutes')"
                ),
                {"c": proposed_id},
            )
            c.execute(
                text(
                    "INSERT INTO tool_call_attempts "
                    "(call_id, tool_name, action_name, arguments, status, "
                    " attempted_at, updated_at) "
                    "VALUES (:c, 'telegram', 'send_message', '{}'::jsonb, "
                    " 'executed', now() - interval '20 minutes', "
                    " now() - interval '20 minutes')"
                ),
                {"c": executed_id},
            )
        finally:
            c.execute(text("ALTER TABLE tool_call_attempts ENABLE TRIGGER USER"))
    info(f"seeded proposed={proposed_id} executed={executed_id}")

    summary = run_reconciliation()
    info(f"reconciler summary: {summary}")
    assert summary["tool_calls_failed"] >= 2, (
        f"expected ≥2 tool_calls_failed; got {summary}"
    )

    with ENGINE.connect() as c:
        rows = c.execute(
            text(
                "SELECT call_id, status, error FROM tool_call_attempts "
                "WHERE call_id IN (:p, :e)"
            ),
            {"p": proposed_id, "e": executed_id},
        ).fetchall()
    by_id = {r[0]: r for r in rows}
    assert by_id[proposed_id][1] == "failed"
    assert "proposed" in by_id[proposed_id][2]
    assert by_id[executed_id][1] == "failed"
    assert "executed-not-confirmed" in by_id[executed_id][2]
    step("proposed→failed and executed→failed with descriptive error text")

    with ENGINE.connect() as c:
        n = c.execute(
            text(
                "SELECT count(*) FROM stack_logs "
                "WHERE query IN ('reconciler_tool_call_failed_proposed', "
                "'reconciler_tool_call_failed_executed') "
                "AND timestamp > now() - interval '5 minutes'"
            ),
        ).scalar()
    assert n >= 2, f"expected ≥2 reconciler tool-call alerts; got {n}"
    step(f"alerts written to stack_logs (count={n})")


def s6_resume_janitor() -> None:
    """A pending_tool_state row in ``status='resuming'`` whose
    ``resumed_at`` is older than the 10-min grace window is reverted to
    ``status='pending'`` by the cleanup janitor."""


    user = f"{UNIQUE}-s6"
    with ENGINE.begin() as c:
        c.execute(
            text("INSERT INTO users (user_id) VALUES (:u) ON CONFLICT DO NOTHING"),
            {"u": user},
        )
        conv = c.execute(
            text(
                "INSERT INTO conversations (user_id, name) VALUES (:u, :n) "
                "RETURNING id"
            ),
            {"u": user, "n": f"{UNIQUE}-s6"},
        ).fetchone()
        cid = str(conv[0])
        c.execute(
            text(
                "INSERT INTO pending_tool_state "
                "(conversation_id, user_id, messages, pending_tool_calls, "
                " tools_dict, tool_schemas, agent_config, "
                " created_at, expires_at, status, resumed_at) "
                "VALUES (CAST(:c AS uuid), :u, '[]'::jsonb, '[]'::jsonb, "
                " '{}'::jsonb, '[]'::jsonb, '{}'::jsonb, "
                " now(), now() + interval '30 minutes', "
                " 'resuming', now() - interval '11 minutes')"
            ),
            {"c": cid, "u": user},
        )
    info(f"seeded pending_tool_state in 'resuming' (resumed_at=-11m) for conv={cid}")

    from application.api.user.tasks import cleanup_pending_tool_state

    res = cleanup_pending_tool_state.run()
    info(f"janitor result: {res}")
    assert res["reverted"] >= 1, f"expected reverted>=1; got {res}"

    with ENGINE.connect() as c:
        row = c.execute(
            text(
                "SELECT status, resumed_at FROM pending_tool_state "
                "WHERE conversation_id = CAST(:c AS uuid) AND user_id = :u"
            ),
            {"c": cid, "u": user},
        ).fetchone()
    assert row[0] == "pending", f"expected status=pending after revert; got {row[0]!r}"
    assert row[1] is None, f"expected resumed_at=NULL after revert; got {row[1]!r}"
    step("janitor reverted stale 'resuming' → 'pending' and cleared resumed_at")


def s7_token_usage_attribution() -> None:
    """``finalize_message`` reads from ``agent.llm.token_usage`` (the LLM
    the stream actually mutated); side-channel LLMs marked with
    ``_persist_token_usage_inline`` self-write a row; the canary fires
    when ``status='complete'`` arrives with zero counts."""

    from application.api.answer.services.conversation_service import (
        ConversationService,
    )
    from application.storage.db.repositories.conversations import (
        ConversationsRepository,
    )
    from application.usage import _maybe_persist_inline

    user = f"{UNIQUE}-s7"
    with ENGINE.begin() as c:
        c.execute(
            text("INSERT INTO users (user_id) VALUES (:u) ON CONFLICT DO NOTHING"),
            {"u": user},
        )
        repo = ConversationsRepository(c)
        conv = repo.create(user, f"{UNIQUE}-s7-conv")
        msg = repo.reserve_message(
            str(conv["id"]),
            prompt="q?",
            placeholder_response="...",
            request_id="r1",
            status="streaming",
        )
        mid_real = str(msg["id"])

        msg2 = repo.reserve_message(
            str(conv["id"]),
            prompt="q?",
            placeholder_response="...",
            request_id="r2",
            status="streaming",
        )
        mid_zero = str(msg2["id"])

    # 1. Real counts → row written, source='agent_stream'.
    with patch_db_session("application.api.answer.services.conversation_service"):
        ok = ConversationService().finalize_message(
            mid_real,
            "answer",
            status="complete",
            token_usage={"prompt_tokens": 3000, "generated_tokens": 800},
            decoded_token={"sub": user},
        )
    assert ok, "finalize_message returned False"

    with ENGINE.connect() as c:
        row = c.execute(
            text(
                "SELECT prompt_tokens, generated_tokens, source "
                "FROM token_usage WHERE user_id = :u "
                "ORDER BY timestamp DESC LIMIT 1"
            ),
            {"u": user},
        ).fetchone()
    assert row == (3000, 800, "agent_stream"), f"row mismatch: {row}"
    step(f"primary stream row written: {row}")

    # 2. Zero-count complete fires the canary alert.
    canary_logger = logging.getLogger(
        "application.api.answer.services.conversation_service"
    )
    seen: list[logging.LogRecord] = []

    class _Catch(logging.Handler):
        def emit(self, record):
            seen.append(record)

    handler = _Catch(level=logging.WARNING)
    canary_logger.addHandler(handler)
    prior_disabled = canary_logger.disabled
    prior_level = canary_logger.level
    canary_logger.disabled = False
    canary_logger.setLevel(logging.WARNING)
    try:
        with patch_db_session(
            "application.api.answer.services.conversation_service",
        ):
            ConversationService().finalize_message(
                mid_zero,
                "answer",
                status="complete",
                token_usage={"prompt_tokens": 0, "generated_tokens": 0},
                decoded_token={"sub": user},
            )
    finally:
        canary_logger.removeHandler(handler)
        canary_logger.disabled = prior_disabled
        canary_logger.setLevel(prior_level)

    assert any(
        getattr(r, "alert", None) == "token_usage_zero_on_complete" for r in seen
    ), (
        "expected structured warning with alert=token_usage_zero_on_complete; "
        f"saw alerts={[getattr(r, 'alert', None) for r in seen]}"
    )
    step("canary fired on zero-count 'complete' finalize")

    # 3. Side-channel inline persist via _maybe_persist_inline.
    class _SideLLM:
        decoded_token = {"sub": user}
        user_api_key = None
        agent_id = None
        _persist_token_usage_inline = True
        _token_usage_source = "compression"

    with patch_db_session("application.usage"):
        _maybe_persist_inline(
            _SideLLM(),
            {"prompt_tokens": 250, "generated_tokens": 80},
        )
    with ENGINE.connect() as c:
        side = c.execute(
            text(
                "SELECT prompt_tokens, generated_tokens, source "
                "FROM token_usage WHERE user_id = :u AND source = 'compression'"
            ),
            {"u": user},
        ).fetchone()
    assert side == (250, 80, "compression"), f"side-channel row mismatch: {side}"
    step(f"side-channel inline persist row written: {side}")


def s8_regenerate_replaces() -> None:
    """``save_user_question(index=N)`` truncates messages at and after
    position N before reserving the placeholder, so the new row lands at
    position=N (replacing the old) rather than appending at the end."""

    from application.api.answer.services.conversation_service import (
        ConversationService,
    )
    from application.storage.db.repositories.conversations import (
        ConversationsRepository,
    )

    user = f"{UNIQUE}-s8"
    with ENGINE.begin() as c:
        c.execute(
            text("INSERT INTO users (user_id) VALUES (:u) ON CONFLICT DO NOTHING"),
            {"u": user},
        )
        repo = ConversationsRepository(c)
        conv = repo.create(user, f"{UNIQUE}-s8-conv")
        cid = str(conv["id"])
        for i in range(5):
            repo.append_message(
                cid,
                {
                    "prompt": f"q{i}",
                    "response": f"a{i}",
                    "thought": "",
                    "sources": [],
                    "tool_calls": [],
                    "metadata": {},
                },
            )
        before = [m["position"] for m in repo.get_messages(cid)]
    assert before == [0, 1, 2, 3, 4], f"seed wrong: {before}"
    info(f"seeded conv={cid} positions={before}")

    with patch_db_session("application.api.answer.services.conversation_service"):
        result = ConversationService().save_user_question(
            conversation_id=cid,
            question=f"{UNIQUE}-regen",
            decoded_token={"sub": user},
            index=3,
        )
    info(f"regenerated at index=3 → message_id={result['message_id']}")

    with ENGINE.connect() as c:
        repo = ConversationsRepository(c)
        msgs = repo.get_messages(cid)
    positions = [m["position"] for m in msgs]
    assert positions == [0, 1, 2, 3], f"expected [0,1,2,3]; got {positions}"
    regen = next(m for m in msgs if m["position"] == 3)
    assert regen["prompt"] == f"{UNIQUE}-regen", f"prompt mismatch: {regen['prompt']!r}"
    assert regen["status"] == "pending", f"status mismatch: {regen['status']!r}"
    assert not any(m["response"] == "a3" for m in msgs), "old answer at index=3 not gone"
    assert not any(m["prompt"] == "q4" for m in msgs), "tail not truncated"
    step(f"regenerate replaced position 3 and truncated tail (positions={positions})")


def s9_idempotency_concurrent_claim() -> None:
    """N concurrent ``claim_task`` calls against the same scoped key:
    exactly one writer wins (returns a row), the rest get None.
    Confirms the ON CONFLICT DO NOTHING race semantics for the HTTP
    claim-first pattern."""

    from application.storage.db.repositories.idempotency import (
        IdempotencyRepository,
    )

    key = f"alice:{UNIQUE}-claim-race"
    n = 20

    def attempt(_idx: int) -> bool:
        # Each thread checks out its own engine connection so we don't
        # serialize at the SQLAlchemy Connection layer.
        with ENGINE.begin() as conn:
            row = IdempotencyRepository(conn).claim_task(
                key=key, task_name="ingest", task_id=f"task-{_idx}",
            )
        return row is not None

    with ThreadPoolExecutor(max_workers=n) as ex:
        winners = sum(1 for r in ex.map(attempt, range(n)) if r)

    with ENGINE.connect() as conn:
        rows = conn.execute(
            text("SELECT count(*) FROM task_dedup WHERE idempotency_key = :k"),
            {"k": key},
        ).scalar()
    assert winners == 1, f"expected exactly 1 winner; got {winners}"
    assert rows == 1, f"expected exactly 1 dedup row; got {rows}"
    step(f"{n} concurrent claim_task → 1 winner, 1 row")

    with ENGINE.begin() as conn:
        conn.execute(
            text("DELETE FROM task_dedup WHERE idempotency_key = :k"),
            {"k": key},
        )


# ---------------------------------------------------------------------------
# Live-LLM scenarios: bootstrap helpers
# ---------------------------------------------------------------------------

import shutil  # noqa: E402
import signal  # noqa: E402
import socket  # noqa: E402
import subprocess  # noqa: E402
from typing import Optional  # noqa: E402


MOCK_LLM_HOST = "127.0.0.1"
MOCK_LLM_PORT = 7899
MOCK_LLM_BASE_URL = f"http://{MOCK_LLM_HOST}:{MOCK_LLM_PORT}/v1"


def _port_open(host: str, port: int, timeout: float = 0.4) -> bool:
    try:
        with socket.create_connection((host, port), timeout=timeout):
            return True
    except OSError:
        return False


def _wait_for(predicate, timeout_s: float = 15.0, interval_s: float = 0.2) -> bool:
    deadline = time.monotonic() + timeout_s
    while time.monotonic() < deadline:
        if predicate():
            return True
        time.sleep(interval_s)
    return False


class MockLLMService:
    """Owns a subprocess running ``scripts/e2e/mock_llm.py``. Idempotent —
    if the port is already open, treats that as "already running" and
    skips bootstrap (so a developer can run their own ``up.sh`` and have
    this script reuse it)."""

    def __init__(
        self,
        *,
        chunk_delay_ms: int = 0,
        total_delay_ms: int = 0,
    ) -> None:
        self._proc: Optional[subprocess.Popen] = None
        self._reused_existing = False
        self._chunk_delay_ms = chunk_delay_ms
        self._total_delay_ms = total_delay_ms
        self._log_path = (
            ROOT / ".e2e-tmp" / "logs" / f"mock-llm-{UNIQUE}.log"
        )

    def start(self) -> None:
        if _port_open(MOCK_LLM_HOST, MOCK_LLM_PORT):
            info(f"mock LLM already running at {MOCK_LLM_BASE_URL} — reusing")
            self._reused_existing = True
            return
        self._log_path.parent.mkdir(parents=True, exist_ok=True)
        env = os.environ.copy()
        env["MOCK_LLM_HOST"] = MOCK_LLM_HOST
        env["MOCK_LLM_PORT"] = str(MOCK_LLM_PORT)
        if self._chunk_delay_ms:
            env["MOCK_LLM_FORCE_STREAM_CHUNK_DELAY_MS"] = str(self._chunk_delay_ms)
        if self._total_delay_ms:
            env["MOCK_LLM_FORCE_TOTAL_DELAY_MS"] = str(self._total_delay_ms)
        env["PYTHONUNBUFFERED"] = "1"
        self._proc = subprocess.Popen(  # noqa: S603
            [sys.executable, str(ROOT / "scripts" / "e2e" / "mock_llm.py")],
            stdout=open(self._log_path, "w"),
            stderr=subprocess.STDOUT,
            env=env,
        )
        if not _wait_for(
            lambda: _port_open(MOCK_LLM_HOST, MOCK_LLM_PORT), timeout_s=10
        ):
            raise RuntimeError(
                f"mock LLM did not open port {MOCK_LLM_PORT} within 10s; "
                f"see {self._log_path}"
            )
        info(f"mock LLM started (pid={self._proc.pid}); log: {self._log_path}")

    def stop(self) -> None:
        if self._proc is None or self._reused_existing:
            return
        try:
            self._proc.send_signal(signal.SIGTERM)
            self._proc.wait(timeout=5)
        except subprocess.TimeoutExpired:
            self._proc.kill()
            self._proc.wait(timeout=2)
        finally:
            self._proc = None


def _build_openai_llm():
    """Construct an OpenAI-provider LLM pointing at the mock LLM. Skips
    the BYOM resolution path that LLMCreator usually goes through, since
    we just want a direct openai-client wrapper for the mock."""
    from application.llm.openai import OpenAILLM

    return OpenAILLM(
        api_key="mock-key",
        user_api_key=None,
        base_url=MOCK_LLM_BASE_URL,
        model_id="gpt-4o-mini",
    )


# ---------------------------------------------------------------------------
# s11: live LLM stream end-to-end
# ---------------------------------------------------------------------------


def s11_live_llm_stream_through_wal() -> None:
    """Reserve a WAL placeholder, drive a real OpenAI-protocol stream
    against the mock LLM, then ``finalize_message`` with the resulting
    ``llm.token_usage``. Verifies:

    * The OpenAI client + ``stream_token_usage`` decorator accumulate
      non-zero counts on the LLM instance during a real network stream.
    * ``finalize_message`` writes a ``token_usage`` row with
      ``source='agent_stream'`` and the recorded counts.
    * The conversation message is left at ``status='complete'`` with the
      streamed response captured.
    """

    from application.api.answer.services.conversation_service import (
        ConversationService,
    )

    user = f"{UNIQUE}-s11"
    with ENGINE.begin() as c:
        c.execute(
            text("INSERT INTO users (user_id) VALUES (:u) ON CONFLICT DO NOTHING"),
            {"u": user},
        )

    with patch_db_session("application.api.answer.services.conversation_service"):
        res = ConversationService().save_user_question(
            conversation_id=None,
            question=f"{UNIQUE} say hello",
            decoded_token={"sub": user},
        )
    mid = res["message_id"]
    info(f"reserved message_id={mid}")

    llm = _build_openai_llm()
    chunks = list(
        llm.gen_stream(
            model="gpt-4o-mini",
            messages=[
                {"role": "user", "content": f"{UNIQUE} say hello"},
            ],
            stream=True,
            tools=None,
        )
    )
    response_text = "".join(c for c in chunks if isinstance(c, str))
    assert response_text, f"no streamed content received; chunks={chunks!r}"
    assert llm.token_usage["prompt_tokens"] > 0, (
        f"expected prompt_tokens > 0; got {llm.token_usage}"
    )
    assert llm.token_usage["generated_tokens"] > 0, (
        f"expected generated_tokens > 0; got {llm.token_usage}"
    )
    step(
        f"streamed {len(chunks)} chunks; tokens={llm.token_usage}; "
        f"response[:40]={response_text[:40]!r}"
    )

    with patch_db_session("application.api.answer.services.conversation_service"):
        ok = ConversationService().finalize_message(
            mid,
            response_text,
            status="complete",
            token_usage=llm.token_usage,
            decoded_token={"sub": user},
        )
    assert ok, "finalize_message returned False"

    with ENGINE.connect() as c:
        msg_row = c.execute(
            text(
                "SELECT status, response FROM conversation_messages "
                "WHERE id = CAST(:m AS uuid)"
            ),
            {"m": mid},
        ).fetchone()
    assert msg_row[0] == "complete", f"expected complete; got {msg_row[0]!r}"
    assert msg_row[1] == response_text, "response mismatch"
    step("WAL row finalised at status='complete' with the streamed text")

    with ENGINE.connect() as c:
        usage = c.execute(
            text(
                "SELECT prompt_tokens, generated_tokens, source "
                "FROM token_usage WHERE user_id = :u "
                "ORDER BY timestamp DESC LIMIT 1"
            ),
            {"u": user},
        ).fetchone()
    assert usage is not None, "no token_usage row written"
    assert usage[0] == llm.token_usage["prompt_tokens"], (
        f"prompt mismatch: row={usage[0]} llm={llm.token_usage['prompt_tokens']}"
    )
    assert usage[1] == llm.token_usage["generated_tokens"], (
        f"generated mismatch: row={usage[1]} llm={llm.token_usage['generated_tokens']}"
    )
    assert usage[2] == "agent_stream", f"source mismatch: {usage[2]!r}"
    step(f"token_usage row written: {usage}")


# ---------------------------------------------------------------------------
# s12: full embedding loop with a mock markdown document
# ---------------------------------------------------------------------------


def s12_full_embedding_loop() -> None:
    """Drive ``ingest_worker`` end-to-end against a tiny markdown file:
    chunker runs, the mock LLM answers each ``/v1/embeddings`` call,
    chunks land in a faiss index, and the row in ``ingest_chunk_progress``
    reflects ``embedded_chunks == total_chunks``.

    Note: we do NOT go through Flask/Celery here — the worker entrypoint
    is just a regular Python function; ``self`` is a MagicMock with the
    minimum attrs the worker reads."""

    from application.core.settings import settings

    user = f"{UNIQUE}-s12"
    with ENGINE.begin() as c:
        c.execute(
            text("INSERT INTO users (user_id) VALUES (:u) ON CONFLICT DO NOTHING"),
            {"u": user},
        )

    # ``LocalStorage`` enforces a path-traversal guard against the project
    # root, so the markdown file has to live UNDER ``<root>/<UPLOAD_FOLDER>``.
    # We use the configured upload folder (relative path) and create a
    # uniquely-named subdir we can clean up afterwards.
    upload_folder_rel = settings.UPLOAD_FOLDER  # e.g. "inputs"
    upload_root = ROOT / upload_folder_rel
    upload_root.mkdir(parents=True, exist_ok=True)

    safe_user = user.replace("/", "_")
    job_dir_name = f"qa-md-{UNIQUE}"
    file_dir_abs = upload_root / safe_user / job_dir_name
    file_dir_abs.mkdir(parents=True, exist_ok=True)
    md_path = file_dir_abs / "qa.md"
    # The FAISS path seeds the index with ``docs[0]`` *before* the loop
    # that records per-chunk progress. To exercise both the seed AND at
    # least one progress-recorded chunk, the markdown needs to be long
    # enough to chunk into ≥2 pieces under the default token settings
    # (~1000 tokens / chunk). We use a synthesised body of clearly-
    # separated paragraphs to give the chunker something realistic to
    # split on.
    paragraph = (
        "DocsGPT durability is verified by an E2E script that drives "
        "the chunker, the embedding pipeline, and the vector store on "
        "a real Postgres + mock-LLM stack. This paragraph exists to "
        "give the chunker enough text to split into multiple chunks "
        "so the per-chunk progress recorder fires at least once after "
        "the FAISS seed step. "
    )
    body = (paragraph * 80) + "\n\n"  # ~80 paragraphs of ~50 tokens each
    md_path.write_text(
        "# QA E2E Document\n\n"
        + "## Section A\n\n" + body
        + "## Section B\n\n" + body
        + "## Section C\n\n" + body,
        encoding="utf-8",
    )

    # Path passed downstream is the relative form the upload route would
    # produce (``LocalStorage`` joins this against project root for I/O).
    file_path_rel = f"{upload_folder_rel}/{safe_user}/{job_dir_name}"

    # Point embeddings at the mock LLM via env+settings, but leave
    # UPLOAD_FOLDER alone so LocalStorage's cached base_dir still points
    # at the project root.
    saved = {
        "EMBEDDINGS_BASE_URL": settings.EMBEDDINGS_BASE_URL,
        "EMBEDDINGS_KEY": settings.EMBEDDINGS_KEY,
        "EMBEDDINGS_NAME": settings.EMBEDDINGS_NAME,
        "VECTOR_STORE": settings.VECTOR_STORE,
    }
    settings.EMBEDDINGS_BASE_URL = MOCK_LLM_BASE_URL
    settings.EMBEDDINGS_KEY = "mock-key"
    settings.EMBEDDINGS_NAME = "openai_text-embedding-3-small"
    settings.VECTOR_STORE = "faiss"

    idempotency_key = f"{user}:{UNIQUE}-s12-key"
    from application.worker import _derive_source_id

    expected_source_id = str(_derive_source_id(idempotency_key))

    fake_self = MagicMock()
    fake_self.update_state = MagicMock()
    fake_self.request = MagicMock()
    fake_self.request.id = f"{UNIQUE}-s12-task"

    # The final ``upload_index`` step in ingest_worker self-calls Flask to
    # register the index with the API. We're not running Flask here, so
    # patch it to a no-op — everything we want to assert (chunker,
    # embedding pipeline, ``ingest_chunk_progress`` rows, vector store
    # files written to disk) happens BEFORE that registration call.
    upload_index_calls: list = []

    def _fake_upload_index(full_path, file_data):
        upload_index_calls.append({"full_path": full_path, "file_data": file_data})
        # Sanity-check the faiss artefacts exist before we'd POST them.
        assert os.path.exists(full_path), f"vector store dir missing: {full_path}"

    try:
        with patch_db_session(
            "application.parser.embedding_pipeline",
            "application.worker",
            "application.api.answer.services.conversation_service",
        ), patch("application.worker.upload_index", _fake_upload_index):
            from application.worker import ingest_worker

            resp = ingest_worker(
                fake_self,
                directory=upload_folder_rel,
                formats=[".md"],
                job_name=job_dir_name,
                file_path=file_path_rel,
                filename=job_dir_name,
                user=user,
                file_name_map={"qa.md": "qa.md"},
                idempotency_key=idempotency_key,
            )
        info(f"ingest_worker returned: {str(resp)[:80]}")
    finally:
        for k, v in saved.items():
            setattr(settings, k, v)

    assert len(upload_index_calls) == 1, (
        f"expected 1 upload_index call; got {len(upload_index_calls)}"
    )
    captured = upload_index_calls[0]["file_data"]
    assert captured["id"] == expected_source_id, (
        f"upload_index id mismatch: {captured['id']} vs {expected_source_id}"
    )
    step(
        f"deterministic source_id derived from scoped key: {expected_source_id} "
        f"(would-have-been-POSTed to /api/upload_index)"
    )

    with ENGINE.connect() as c:
        prog = c.execute(
            text(
                "SELECT total_chunks, embedded_chunks, last_index "
                "FROM ingest_chunk_progress WHERE source_id = CAST(:s AS uuid)"
            ),
            {"s": expected_source_id},
        ).fetchone()
    assert prog is not None, "no ingest_chunk_progress row for source"
    assert prog[0] >= 1, f"expected at least 1 chunk; got total={prog[0]}"
    assert prog[1] == prog[0], (
        f"expected embedded == total; got embedded={prog[1]} total={prog[0]}"
    )
    step(
        f"ingest_chunk_progress: embedded={prog[1]}/{prog[0]} "
        f"last_index={prog[2]}"
    )

    # Cleanup
    with ENGINE.begin() as c:
        c.execute(
            text(
                "DELETE FROM ingest_chunk_progress WHERE source_id = CAST(:s AS uuid)"
            ),
            {"s": expected_source_id},
        )
    shutil.rmtree(file_dir_abs.parent, ignore_errors=True)


# ---------------------------------------------------------------------------
# s13: SIGKILL Celery worker mid-task; verify acks_late redelivers
# ---------------------------------------------------------------------------


class CeleryWorkerService:
    """Lightweight Celery worker subprocess for the chaos test. Uses
    isolated Redis DBs (11/12) so it can't grab DocsGPT production work
    by accident, and a unique queue name so two consecutive runs don't
    inherit each other's unacked messages."""

    _log_seq = 0

    def __init__(self, queue: str, label: str = "celery") -> None:
        self._proc: Optional[subprocess.Popen] = None
        self._queue = queue
        CeleryWorkerService._log_seq += 1
        self._log_path = (
            ROOT / ".e2e-tmp" / "logs"
            / f"{label}-{UNIQUE}-{CeleryWorkerService._log_seq}.log"
        )

    @property
    def pid(self) -> Optional[int]:
        return self._proc.pid if self._proc else None

    def start(self, env_overrides: dict) -> None:
        self._log_path.parent.mkdir(parents=True, exist_ok=True)
        env = os.environ.copy()
        env.update(env_overrides)
        env["PYTHONUNBUFFERED"] = "1"
        self._proc = subprocess.Popen(  # noqa: S603
            [
                sys.executable, "-m", "celery",
                "-A", "application.app.celery", "worker",
                "-l", "INFO", "--pool=solo",
                "-Q", self._queue,
                "-n", f"qa-{UNIQUE}@%h",
                "--without-gossip", "--without-mingle", "--without-heartbeat",
            ],
            stdout=open(self._log_path, "w"),
            stderr=subprocess.STDOUT,
            env=env,
            cwd=str(ROOT),
        )

        def _ready() -> bool:
            try:
                txt = self._log_path.read_text()
            except OSError:
                return False
            return f"qa-{UNIQUE}@" in txt and "ready." in txt

        if not _wait_for(_ready, timeout_s=20):
            raise RuntimeError(
                f"celery worker not ready within 20s; see {self._log_path}"
            )
        info(f"celery worker started pid={self._proc.pid} log={self._log_path}")

    def kill(self, sig: int = signal.SIGKILL) -> None:
        if self._proc is None:
            return
        try:
            os.kill(self._proc.pid, sig)
        except ProcessLookupError:
            pass
        self._proc.wait(timeout=5)
        info(f"celery worker pid={self._proc.pid} terminated by signal {sig}")
        self._proc = None

    def stop(self) -> None:
        if self._proc is None:
            return
        try:
            self._proc.send_signal(signal.SIGTERM)
            self._proc.wait(timeout=5)
        except subprocess.TimeoutExpired:
            self._proc.kill()
            self._proc.wait(timeout=2)
        self._proc = None


def s13_sigkill_redelivers_via_acks_late() -> None:
    """Real-process chaos test:

    1. Start a Celery worker in a subprocess (isolated Redis DBs + queue).
    2. Dispatch ``process_agent_webhook`` configured to hit a slow mock
       LLM stream (``MOCK_LLM_FORCE_STREAM_CHUNK_DELAY_MS=400``).
    3. ``kill -9`` the worker mid-task before the LLM finishes.
    4. Restart the worker with the same broker.
    5. Watch the broker re-deliver the same task; the second run completes
       the LLM stream and exits cleanly.
    """

    import redis  # type: ignore[import-not-found]

    queue = f"docsgpt-chaos-{UNIQUE}"
    broker_db = 11
    backend_db = 12
    broker_url = f"redis://127.0.0.1:6379/{broker_db}"
    backend_url = f"redis://127.0.0.1:6379/{backend_db}"

    # Make sure these Redis DBs are clean of leftovers from prior runs.
    r = redis.Redis(host="127.0.0.1", port=6379, db=broker_db)
    r.flushdb()
    redis.Redis(host="127.0.0.1", port=6379, db=backend_db).flushdb()
    # Bust the stream-response cache (DB 0) too — without this, a prior
    # run's cached response satisfies the stream instantly and the mock
    # LLM's chunk delay never gets exercised, so SIGKILL fires after the
    # task already completed.
    redis.Redis(host="127.0.0.1", port=6379, db=0).flushdb()

    # The mock LLM running on the main port (started by the runner) has
    # no chunk delay. s13 needs a slow stream so SIGKILL lands while the
    # worker is still consuming chunks. Stop the existing mock and start
    # a delayed one in its place; the delay env reaches mock_llm.py via
    # ``MOCK_LLM_FORCE_STREAM_CHUNK_DELAY_MS``.
    delayed_mock = MockLLMService(chunk_delay_ms=400)
    # Force-restart even if the existing one is reused: kill via port.
    if _port_open(MOCK_LLM_HOST, MOCK_LLM_PORT):
        # The runner's mock is owned by ``main()``; we need to take it
        # offline temporarily so our delayed copy can bind the port.
        # Send SIGTERM to whoever owns the port via lsof (best-effort).
        try:
            pids = subprocess.check_output(
                ["lsof", "-tiTCP:7899", "-sTCP:LISTEN"], text=True
            ).split()
            for pid in pids:
                os.kill(int(pid), signal.SIGTERM)
            time.sleep(0.3)
        except (subprocess.CalledProcessError, FileNotFoundError, ProcessLookupError):
            pass
    delayed_mock.start()

    # We need an agent (with key + webhook token + prompt) to drive the
    # webhook task. The mock LLM has no idea about retrieval, so the
    # 'classic' agent_type is fine.
    user = f"{UNIQUE}-s13"
    with ENGINE.begin() as c:
        c.execute(
            text("INSERT INTO users (user_id) VALUES (:u) ON CONFLICT DO NOTHING"),
            {"u": user},
        )
        prompt_row = c.execute(
            text(
                "INSERT INTO prompts (user_id, name, content) "
                "VALUES (:u, :n, :c) RETURNING id"
            ),
            {"u": user, "n": f"{UNIQUE}-s13", "c": "you are a tester"},
        ).fetchone()
        prompt_id = str(prompt_row[0])
        agent_id = str(uuid.uuid4())
        agent_key = f"{UNIQUE}-akey-{uuid.uuid4().hex[:8]}"
        c.execute(
            text(
                "INSERT INTO agents (id, user_id, name, status, retriever, "
                "prompt_id, tools, agent_type, key, chunks) "
                "VALUES (CAST(:id AS uuid), :u, :n, 'published', 'classic', "
                "CAST(:p AS uuid), '[]'::jsonb, 'classic', :k, '2')"
            ),
            {
                "id": agent_id, "u": user, "n": f"{UNIQUE}-s13",
                "p": prompt_id, "k": agent_key,
            },
        )
    info(f"seeded agent={agent_id} user={user}")

    env_overrides = {
        "POSTGRES_URI": os.environ["POSTGRES_URI"],
        "CELERY_BROKER_URL": broker_url,
        "CELERY_RESULT_BACKEND": backend_url,
        "LLM_PROVIDER": "openai",
        "LLM_NAME": "gpt-4o-mini",
        "API_KEY": "mock-key",
        "OPENAI_API_KEY": "mock-key",
        "OPENAI_BASE_URL": MOCK_LLM_BASE_URL,
        "EMBEDDINGS_BASE_URL": MOCK_LLM_BASE_URL,
        "EMBEDDINGS_KEY": "mock-key",
        "EMBEDDINGS_NAME": "openai_text-embedding-3-small",
        "VECTOR_STORE": "faiss",
        "MOCK_LLM_FORCE_STREAM_CHUNK_DELAY_MS": "400",
        "AUTO_MIGRATE": "false",
        "AUTO_CREATE_DB": "false",
        # Short visibility window so the broker re-queues the SIGKILLed
        # task within seconds rather than the 1-hour production default.
        # Has to be set on BOTH workers — the unack entry's age is
        # checked by whichever worker scans next.
        "CELERY_VISIBILITY_TIMEOUT": "5",
    }

    worker = CeleryWorkerService(queue=queue)
    worker.start(env_overrides)

    # Dispatch via celery's send_task using the same broker URL the worker
    # listens on. The DocsGPT celery instance has already loaded
    # ``application.celeryconfig`` (broker = production .env DB 0); since
    # Celery's conf re-reads ``CELERY_BROKER_URL`` from the env at publish
    # time, we override the env for the *whole* dispatch + wait block.
    from celery import Celery

    saved_broker = os.environ.get("CELERY_BROKER_URL")
    saved_backend = os.environ.get("CELERY_RESULT_BACKEND")
    os.environ["CELERY_BROKER_URL"] = broker_url
    os.environ["CELERY_RESULT_BACKEND"] = backend_url

    try:
        chaos_app = Celery("chaos", broker=broker_url, backend=backend_url)
        chaos_app.conf.broker_url = broker_url
        chaos_app.conf.result_backend = backend_url
        chaos_app.conf.task_default_queue = queue
        chaos_app.conf.task_default_exchange = queue
        chaos_app.conf.task_default_routing_key = queue
        chaos_app.conf.task_acks_late = True
        chaos_app.conf.task_reject_on_worker_lost = True
        info(
            f"chaos_app broker={chaos_app.conf.broker_url} "
            f"queue={chaos_app.conf.task_default_queue}"
        )

        idempotency_key = f"{agent_id}:{UNIQUE}-s13-key"
        task_id = str(uuid.uuid4())
        info(f"dispatching task_id={task_id} key={idempotency_key}")
        chaos_app.send_task(
            "application.api.user.tasks.process_agent_webhook",
            kwargs={
                "agent_id": agent_id,
                "payload": {"event": f"{UNIQUE}-s13"},
                "idempotency_key": idempotency_key,
            },
            queue=queue,
            task_id=task_id,
        )

        # Sanity: the task should be visible in the broker for at least a
        # brief moment before the worker pulls it. If it never appears,
        # the publish missed the queue entirely.
        queue_seen = _wait_for(
            lambda: r.llen(queue) >= 1, timeout_s=2.0, interval_s=0.05,
        )
        info(
            f"after send_task: queue len={r.llen(queue)} "
            f"(visible-briefly={queue_seen}) "
            f"unacked={r.hlen('unacked') if r.exists('unacked') else 0}"
        )

        # Wait until the worker actually starts running the task. The
        # mock LLM streams 5 chunks at 400ms each (= 2s), giving us a
        # window to SIGKILL while the task is still in flight.
        def _task_received() -> bool:
            try:
                txt = worker._log_path.read_text()  # noqa: SLF001
            except OSError:
                return False
            return task_id in txt

        started = _wait_for(_task_received, timeout_s=20)
        if not started:
            try:
                tail = worker._log_path.read_text()[-2000:]  # noqa: SLF001
                warn(f"worker log tail:\n{tail}")
            except OSError:
                pass
            qlen = r.llen(queue)
            unacked_now = r.hlen("unacked") if r.exists("unacked") else 0
            warn(f"broker state: queue len={qlen} unacked={unacked_now}")
            worker.stop()
            raise AssertionError(
                f"task did not start on worker 1 within 20s — "
                f"see {worker._log_path}"
            )
        info("task picked up by worker 1; killing mid-flight")
        sigkill_pid = worker.pid
        worker.kill(signal.SIGKILL)
        info(f"SIGKILLed worker pid={sigkill_pid}")

        unacked_count = r.hlen("unacked") if r.exists("unacked") else 0
        info(f"broker 'unacked' hash size after SIGKILL: {unacked_count}")

        # Restart the worker — acks_late + reject_on_worker_lost should
        # make the broker redeliver the same task to it.
        worker2 = CeleryWorkerService(queue=queue)
        worker2.start(env_overrides)

        def _task_completed() -> bool:
            try:
                txt = worker2._log_path.read_text()  # noqa: SLF001
            except OSError:
                return False
            return (
                task_id in txt
                and ("succeeded in" in txt or "Task succeeded" in txt)
            )

        completed = _wait_for(_task_completed, timeout_s=60)
        worker2.stop()
        assert completed, (
            f"redelivered task did not complete within 60s; "
            f"see {worker2._log_path}"
        )
        step(f"task {task_id} redelivered after SIGKILL and completed cleanly")
    finally:
        if saved_broker is not None:
            os.environ["CELERY_BROKER_URL"] = saved_broker
        else:
            os.environ.pop("CELERY_BROKER_URL", None)
        if saved_backend is not None:
            os.environ["CELERY_RESULT_BACKEND"] = saved_backend
        else:
            os.environ.pop("CELERY_RESULT_BACKEND", None)

    # Cleanup
    delayed_mock.stop()
    redis.Redis(host="127.0.0.1", port=6379, db=broker_db).flushdb()
    redis.Redis(host="127.0.0.1", port=6379, db=backend_db).flushdb()
    with ENGINE.begin() as c:
        c.execute(text("DELETE FROM agents WHERE id = CAST(:a AS uuid)"), {"a": agent_id})
        c.execute(text("DELETE FROM prompts WHERE id = CAST(:p AS uuid)"), {"p": prompt_id})


def s10_celery_default_queue() -> None:
    """Sanity: the broker queue is project-scoped so a sibling worker on
    the same Redis can't grab DocsGPT tasks."""

    from application import celeryconfig

    assert celeryconfig.task_default_queue == "docsgpt", (
        f"expected default queue 'docsgpt'; got {celeryconfig.task_default_queue!r}"
    )
    assert celeryconfig.task_acks_late is True
    assert celeryconfig.task_reject_on_worker_lost is True
    assert celeryconfig.broker_transport_options.get("visibility_timeout") == 3600, (
        f"visibility_timeout mismatch: "
        f"{celeryconfig.broker_transport_options}"
    )
    step("celeryconfig: docsgpt queue, acks_late, reject_on_lost, visibility=3600")


# ---------------------------------------------------------------------------
# Runner
# ---------------------------------------------------------------------------

SCENARIOS = [
    ("s1", s1_wal_and_reconciler,
     "WAL placeholder + reconciler escalation"),
    ("s2", s2_webhook_idempotency_and_scoping,
     "webhook idempotency + cross-agent scope"),
    ("s3", s3_ingest_deterministic_and_resume,
     "deterministic source_id + chunk-progress resume"),
    ("s4", s4_token_usage_atomic_rollback,
     "token_usage rolls back atomically with message update"),
    ("s5", s5_reconciler_tool_call_sweeps,
     "reconciler escalates stuck tool_call_attempts"),
    ("s6", s6_resume_janitor,
     "pending_tool_state janitor reverts stale 'resuming'"),
    ("s7", s7_token_usage_attribution,
     "token_usage attribution + canary + side-channel inline persist"),
    ("s8", s8_regenerate_replaces,
     "regenerate at index truncates tail and replaces"),
    ("s9", s9_idempotency_concurrent_claim,
     "concurrent claim_task race: exactly one winner"),
    ("s10", s10_celery_default_queue,
     "celeryconfig sanity: docsgpt queue + durability defaults"),
    ("s11", s11_live_llm_stream_through_wal,
     "LIVE: real OpenAI-protocol stream → WAL → finalize → token_usage"),
    ("s12", s12_full_embedding_loop,
     "LIVE: ingest a markdown file end-to-end (chunker + embeddings + faiss)"),
    ("s13", s13_sigkill_redelivers_via_acks_late,
     "LIVE: SIGKILL Celery worker mid-task; verify acks_late redelivers"),
]


# Scenarios that need the mock LLM running. The runner starts/stops it
# automatically when any of these are selected.
_LIVE_SCENARIOS = {"s11", "s12", "s13"}


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--list", action="store_true", help="list scenarios and exit")
    parser.add_argument(
        "--only",
        help="comma-separated scenario keys to run (e.g. s1,s5)",
    )
    args = parser.parse_args()

    if args.list:
        for key, _, desc in SCENARIOS:
            print(f"  {key}: {desc}")
        return 0

    selected = SCENARIOS
    if args.only:
        wanted = {s.strip() for s in args.only.split(",") if s.strip()}
        selected = [s for s in SCENARIOS if s[0] in wanted]
        missing = wanted - {s[0] for s in selected}
        if missing:
            print(f"unknown scenarios: {sorted(missing)}", file=sys.stderr)
            return 2

    print(f"{DIM}prefix: {UNIQUE}{RESET}")
    print(f"{DIM}engine: {POSTGRES_URI[:60]}…{RESET}")

    selected_keys = {k for k, _, _ in selected}
    needs_mock_llm = bool(selected_keys & _LIVE_SCENARIOS)
    mock_llm: Optional[MockLLMService] = None
    if needs_mock_llm:
        mock_llm = MockLLMService()
        try:
            mock_llm.start()
        except Exception as exc:  # noqa: BLE001
            print(f"  {RED}failed to start mock LLM: {exc}{RESET}", file=sys.stderr)
            return 2

    passes = 0
    fails: list[tuple[str, str]] = []

    try:
        for key, fn, desc in selected:
            header(f"[{key}] {desc}")
            try:
                fn()
                passes += 1
            except AssertionError as exc:
                fail(str(exc))
                fails.append((key, str(exc)))
            except Exception as exc:  # noqa: BLE001
                fail(f"{type(exc).__name__}: {exc}")
                traceback.print_exc()
                fails.append((key, f"{type(exc).__name__}: {exc}"))
            finally:
                try:
                    cleanup_artifacts(UNIQUE)
                except Exception as exc:  # noqa: BLE001
                    warn(f"cleanup raised: {exc}")
    finally:
        if mock_llm is not None:
            mock_llm.stop()

    print()
    header("summary")
    print(f"  {GREEN}{passes} passed{RESET}, {RED}{len(fails)} failed{RESET}")
    if fails:
        print()
        for key, why in fails:
            print(f"  {RED}{key}{RESET}: {why}")
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())