"""
cortex/hippo/engine.py — HippoEngine: main lifecycle manager for LCM conversations.

Implements the core lifecycle hooks:
- initialize() — Setup, run migrations, open DB
- ingest_message() — Store to immutable store, update context items
- get_context() — Assemble active context (summaries + fresh tail)
- should_compact() — Check if compaction threshold is exceeded
- compact() — Trigger compaction pipeline (stub in Phase 1)
- shutdown() — Clean up connections

Issue: #802
"""

from __future__ import annotations

import asyncio
import json
import logging
from typing import Any, Callable, Optional, TypeVar

from cortex.hippo.assembler import ContextAssembler
from cortex.hippo.bootstrap import message_identity, read_leaf_path_messages
from cortex.hippo.budget import CompactionBudget
from cortex.hippo.compaction import estimate_tokens
from cortex.hippo.compaction.engine import CompactionEngine
from cortex.hippo.compaction.hardening import (
    CompactionConcurrencyGuard,
    check_and_skip_empty,
    mark_compaction_end,
    mark_compaction_start,
    maybe_chunk_long_content,
    touch_compaction_heartbeat,
)
from cortex.hippo.compaction.types import SummarizeFn
from cortex.hippo.config import HippoConfig
from cortex.hippo.scheduler import CompactionMode, CompactionScheduler
from cortex.hippo.token_estimation import estimate_session_tokens
from cortex.hippo.db.connection import HippoDatabase
from cortex.hippo.models import (
    AssembleResult,
    CompactionDecision,
    CompactionResult,
    ConversationRecord,
    ConversationStats,
    IngestResult,
)
from cortex.hippo.store.conversation_store import ConversationStore, _estimate_tokens
from cortex.hippo.store.summary_store import SummaryStore

logger = logging.getLogger(__name__)

T = TypeVar("T")


# ── Message part decomposition ───────────────────────────────────────────────


def _estimate_tokens_for_content(content: Any) -> int:
    """Estimate token count for raw content (string or structured blocks).

    For plain strings: uses character count / 4.
    For structured content arrays: JSON-serializes first for a realistic estimate
    that accounts for tool inputs and thinking blocks.
    """
    if isinstance(content, str):
        return _estimate_tokens(content)
    if isinstance(content, (list, dict)):
        try:
            serialized = json.dumps(content)
            return _estimate_tokens(serialized)
        except (TypeError, ValueError):
            return 0
    return 0


def _extract_message_content(content: Any) -> str:
    """Extract plain text from raw content for DB storage.

    Mirrors lossless-claw's extractMessageContent():
    - String: returned as-is
    - List of blocks: only text-type blocks are joined (tool/thinking go to parts)
    - Other: JSON-serialized
    """
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        texts = []
        for block in content:
            if isinstance(block, dict) and block.get("type") == "text":
                text = block.get("text", "")
                if isinstance(text, str) and text:
                    texts.append(text)
        return "\n".join(texts)
    if content is None:
        return ""
    try:
        return json.dumps(content)
    except (TypeError, ValueError):
        return str(content)


def _build_message_parts(role: str, content: Any) -> list[dict]:
    """Decompose message content into message_parts rows.

    Mirrors lossless-claw's buildMessageParts(). Returns a list of dicts
    ready to pass to ConversationStore.create_message_parts().

    Part type mapping:
    - str content          → single "text" part
    - block type "text"    → "text"
    - thinking/reasoning   → "thinking"
    - tool_use variants    → "tool_use" (with tool_call_id, tool_name, tool_input)
    - tool_result variants → "tool_result" (with tool_call_id, tool_output)
    - unknown types        → "agent"
    - empty content        → [] (no rows written)
    """
    if content is None:
        return []

    metadata_base = {"originalRole": role}

    # Plain string → single text part
    if isinstance(content, str):
        if not content:
            return []
        return [{
            "ordinal": 0,
            "part_type": "text",
            "text_content": content,
            "metadata": json.dumps(metadata_base),
        }]

    # Non-list, non-string (e.g. dict)
    if not isinstance(content, list):
        return [{
            "ordinal": 0,
            "part_type": "agent",
            "text_content": None,
            "metadata": json.dumps({**metadata_base, "raw": content}),
        }]

    # Empty array
    if len(content) == 0:
        return []

    parts: list[dict] = []
    for ordinal, block in enumerate(content):
        if not isinstance(block, dict):
            parts.append({
                "ordinal": ordinal,
                "part_type": "agent",
                "text_content": str(block) if block is not None else None,
                "metadata": json.dumps(metadata_base),
            })
            continue

        block_type = block.get("type", "")

        if block_type == "text":
            parts.append({
                "ordinal": ordinal,
                "part_type": "text",
                "text_content": block.get("text"),
                "metadata": json.dumps({**metadata_base, "rawType": block_type}),
            })

        elif block_type in ("thinking", "reasoning"):
            # Thinking blocks may have "thinking", "text", or nested "summary"
            thinking_text = (
                block.get("thinking")
                or block.get("text")
                or block.get("summary")
            )
            # Handle summary as nested structure (e.g. list of {text: ...})
            if not isinstance(thinking_text, str) and thinking_text is not None:
                try:
                    thinking_text = json.dumps(thinking_text)
                except (TypeError, ValueError):
                    thinking_text = str(thinking_text)
            parts.append({
                "ordinal": ordinal,
                # Store as "reasoning" to match lossless-claw convention and
                # what the assembler's _block_from_part() expects.
                "part_type": "reasoning",
                "text_content": thinking_text,
                "metadata": json.dumps({**metadata_base, "rawType": block_type}),
            })

        elif block_type in (
            "tool_use", "toolUse", "tool-use",
            "toolCall", "functionCall", "function_call",
        ):
            tool_call_id = (
                block.get("id")
                or block.get("tool_call_id")
                or block.get("toolCallId")
                or block.get("toolUseId")
                or block.get("tool_use_id")
                or block.get("call_id")
            )
            tool_input = block.get("input")
            if tool_input is None:
                tool_input = block.get("arguments")
            if tool_input is None:
                tool_input = block.get("tool_input")
            if tool_input is not None and not isinstance(tool_input, str):
                try:
                    tool_input = json.dumps(tool_input)
                except (TypeError, ValueError):
                    tool_input = str(tool_input)
            parts.append({
                "ordinal": ordinal,
                "part_type": "tool_use",
                "text_content": None,
                "tool_call_id": str(tool_call_id) if tool_call_id is not None else None,
                "tool_name": (
                    block.get("name")
                    or block.get("toolName")
                    or block.get("tool_name")
                ),
                "tool_input": tool_input,
                "metadata": json.dumps({**metadata_base, "rawType": block_type}),
            })

        elif block_type in (
            "tool_result", "toolResult", "function_call_output",
        ):
            tool_call_id = (
                block.get("tool_use_id")
                or block.get("tool_call_id")
                or block.get("toolCallId")
                or block.get("toolUseId")
                or block.get("id")
            )
            output = (
                block.get("content")
                or block.get("output")
                or block.get("tool_output")
            )
            if output is not None and not isinstance(output, str):
                try:
                    output = json.dumps(output)
                except (TypeError, ValueError):
                    output = str(output)
            parts.append({
                "ordinal": ordinal,
                "part_type": "tool_result",
                "text_content": None,
                "tool_call_id": str(tool_call_id) if tool_call_id is not None else None,
                "tool_output": output,
                "metadata": json.dumps({**metadata_base, "rawType": block_type}),
            })

        else:
            # Unknown block type → "agent" fallback
            parts.append({
                "ordinal": ordinal,
                "part_type": "agent",
                "text_content": block.get("text") or block.get("thinking"),
                "metadata": json.dumps({
                    **metadata_base,
                    "rawType": block_type,
                }),
            })

    return parts

T = TypeVar("T")


class HippoEngine:
    """Main lifecycle manager for LCM conversations.

    This is the primary entry point for all Hippo operations.
    It coordinates the database, stores, and assembler.
    """

    def __init__(
        self,
        config: Optional[HippoConfig] = None,
        summarize_fn: Optional[SummarizeFn] = None,
    ) -> None:
        self.config = config or HippoConfig()
        self._db: Optional[HippoDatabase] = None
        self._conversation_store: Optional[ConversationStore] = None
        self._summary_store: Optional[SummaryStore] = None
        self._assembler: Optional[ContextAssembler] = None
        self._compaction_engine: Optional[CompactionEngine] = None
        self._budget: Optional[CompactionBudget] = None
        self._scheduler: Optional[CompactionScheduler] = None
        self._summarize_fn: Optional[SummarizeFn] = summarize_fn
        self._owner_id: Optional[str] = None
        self._conversation: Optional[ConversationRecord] = None
        self._initialized = False
        self._last_compaction_at: Optional[str] = None
        self._session_locks: dict[int, asyncio.Lock] = {}
        self._background_tasks: set[asyncio.Task] = set()
        # Hardening: per-engine concurrency guard for compaction
        self._compaction_guard = CompactionConcurrencyGuard()

    @property
    def conversation_store(self) -> ConversationStore:
        if self._conversation_store is None:
            raise RuntimeError("HippoEngine not initialized")
        return self._conversation_store

    @property
    def summary_store(self) -> SummaryStore:
        if self._summary_store is None:
            raise RuntimeError("HippoEngine not initialized")
        return self._summary_store

    @property
    def assembler(self) -> ContextAssembler:
        if self._assembler is None:
            raise RuntimeError("HippoEngine not initialized")
        return self._assembler

    async def _with_session_lock(
        self, conversation_id: int, fn: Callable[[], Any]
    ) -> Any:
        """Run *fn()* under a per-conversation asyncio.Lock.

        Creates the lock lazily and removes it once no coroutine is waiting,
        preventing unbounded growth of ``_session_locks``.
        """
        lock = self._session_locks.setdefault(conversation_id, asyncio.Lock())

        async with lock:
            result = await fn()

        return result

    def _spawn_background(self, coro) -> asyncio.Task:
        """Create a background task and prevent GC until it completes."""
        task = asyncio.create_task(coro)
        self._background_tasks.add(task)
        task.add_done_callback(self._background_tasks.discard)
        return task

    async def initialize(
        self,
        db_path: Optional[str] = None,
        owner_id: str = "default",
        conversation_id: Optional[str] = None,
        session_key: Optional[str] = None,
    ) -> ConversationRecord:
        """Initialize the engine: open DB, run migrations, setup stores.

        Parameters
        ----------
        db_path : str, optional
            Override config db_path. Useful for testing with `:memory:`.
        owner_id : str
            Tenant identifier for multi-tenant isolation.
        conversation_id : str, optional
            Session ID to bootstrap/resume. Creates a new conversation if None.

        Returns
        -------
        ConversationRecord
            The active conversation record.
        """
        if db_path:
            self.config = self.config.model_copy(update={"db_path": db_path})

        self._owner_id = owner_id
        self._db = HippoDatabase(self.config)
        await self._db.initialize()

        self._conversation_store = ConversationStore(self._db.conn)
        self._summary_store = SummaryStore(
            self._db.conn,
            cache_max_size=self.config.cache_max_summaries,
            cache_ttl_seconds=self.config.cache_ttl_seconds,
            vec_available=self._db.vec_available,
        )

        # Get or create conversation
        session_id = conversation_id or "default"
        self._conversation = await self._conversation_store.get_or_create_conversation(
            owner_id, session_id, session_key=session_key
        )

        self._assembler = ContextAssembler(
            conversation_store=self._conversation_store,
            summary_store=self._summary_store,
            config=self.config,
        )

        self._budget = CompactionBudget(
            max_summarization_tokens_per_hour=self.config.max_summarization_tokens_per_hour,
            global_max_summarization_tokens_per_hour=self.config.global_max_summarization_tokens_per_hour,
            circuit_breaker_threshold=self.config.circuit_breaker_threshold,
            circuit_breaker_cooldown_seconds=self.config.circuit_breaker_cooldown_seconds,
            min_compaction_reduction_ratio=self.config.min_compaction_reduction_ratio,
            budget_warning_threshold=self.config.budget_warning_threshold,
        )

        # H-13: Adaptive compaction scheduler (#1191)
        if self.config.adaptive_scheduling_enabled:
            self._scheduler = CompactionScheduler(
                idle_threshold_seconds=self.config.compaction_idle_threshold_seconds,
                urgency_bypass=self.config.compaction_urgency_bypass,
                light_threshold=self.config.compaction_light_threshold,
                full_threshold=self.config.compaction_full_threshold,
            )
        else:
            self._scheduler = None

        async def _heartbeat_cb(conversation_id: int) -> None:
            """Heartbeat callback wired into CompactionEngine (B-006)."""
            if self._db and self._db.conn:
                await touch_compaction_heartbeat(self._db.conn, conversation_id)

        self._compaction_engine = CompactionEngine(
            conversation_store=self._conversation_store,
            summary_store=self._summary_store,
            config=self.config,
            budget=self._budget,
            heartbeat_fn=_heartbeat_cb,
        )

        self._initialized = True

        logger.info(
            "HippoEngine initialized: owner=%s, conversation=%d, session=%s",
            owner_id,
            self._conversation.conversation_id,
            session_id,
        )

        return self._conversation

    async def ingest_message(
        self,
        role: str,
        content: Any,
        metadata: Optional[str] = None,
        message_parts: Optional[list[dict]] = None,
    ) -> IngestResult:
        """Ingest a message into the immutable store and update context items.

        Parameters
        ----------
        role : str
            Message role: 'user', 'assistant', 'system', 'tool'.
        content : Any
            Message content: plain string or structured content block array.
        metadata : str, optional
            JSON metadata blob.
        message_parts : list[dict], optional
            Structured message parts (tool calls, results, etc.) to store
            alongside the flattened content. Each dict should contain keys
            matching the ``hippo_message_parts`` schema.

        Returns
        -------
        IngestResult
            Contains message_id, conversation_id, token_count, seq.
        """
        if not self._initialized or not self._conversation:
            raise RuntimeError("HippoEngine not initialized — call initialize() first")

        # Hardening: skip empty messages
        if check_and_skip_empty(role, content):
            # Return a zero-token sentinel result without writing to DB
            conv = self._conversation
            return IngestResult(
                message_id=0,
                conversation_id=conv.conversation_id,
                token_count=0,
                seq=0,
            )

        # Hardening: chunk very long messages (>100K tokens) before ingesting
        was_chunked, content_pieces = maybe_chunk_long_content(role, content)
        if was_chunked:
            # Ingest each chunk as a separate message; return the last result
            result = IngestResult(message_id=0, conversation_id=0, token_count=0, seq=0)
            for piece in content_pieces:
                result = await self.ingest_message(
                    role=role,
                    content=piece,
                    metadata=metadata,
                    message_parts=None,  # parts will be rebuilt per-piece
                )
            return result

        conv = self._conversation

        async def _do_ingest() -> IngestResult:
            # Extract plain-text copy for messages.content column
            plain_content = _extract_message_content(content)
            # Token count on full content shape (not just stripped text)
            token_count = _estimate_tokens_for_content(content)

            msg = await self._conversation_store.create_message(
                conversation_id=conv.conversation_id,
                owner_id=self._owner_id,
                role=role,
                content=plain_content,
                metadata=metadata,
                token_count=token_count,
            )

            # Write message parts (R-036 #1: previously never called)
            # K-005: allow caller-provided parts (e.g. from extract_message_parts)
            parts = message_parts if message_parts is not None else _build_message_parts(role, content)
            if parts:
                await self._conversation_store.create_message_parts(
                    msg.message_id, parts
                )

            # Append to context items
            await self._summary_store.append_context_item(
                conversation_id=conv.conversation_id,
                item_type="message",
                message_id=msg.message_id,
            )

            # H-13: record message for adaptive scheduling (#1191)
            if self._scheduler is not None:
                self._scheduler.record_message(conv.conversation_id)

            logger.debug(
                "Ingested message %d (role=%s, tokens=%d, seq=%d)",
                msg.message_id,
                role,
                msg.token_count,
                msg.seq,
            )

            return IngestResult(
                message_id=msg.message_id,
                conversation_id=conv.conversation_id,
                token_count=msg.token_count,
                seq=msg.seq,
            )

        return await self._with_session_lock(conv.conversation_id, _do_ingest)

    async def get_context(
        self, max_tokens: Optional[int] = None
    ) -> AssembleResult:
        """Assemble the active context for the model.

        Returns an ordered list of context items: cornerstones + summaries + fresh tail.

        Parameters
        ----------
        max_tokens : int, optional
            Token budget override. Defaults to config.max_tokens.
        """
        if not self._initialized or not self._conversation:
            raise RuntimeError("HippoEngine not initialized")

        budget = max_tokens or self.config.max_tokens

        return await self._assembler.assemble(
            conversation_id=self._conversation.conversation_id,
            owner_id=self._owner_id,
            token_budget=budget,
        )

    async def should_compact(
        self, threshold: Optional[float] = None
    ) -> CompactionDecision:
        """Check if compaction should run based on token usage.

        Returns a CompactionDecision indicating whether compaction is needed.
        """
        if not self._initialized or not self._conversation:
            raise RuntimeError("HippoEngine not initialized")

        conv = self._conversation
        thresh = threshold or self.config.context_threshold
        max_tokens = self.config.max_tokens
        threshold_tokens = int(max_tokens * thresh)

        # Calculate current context token count
        total_msg_tokens = await self._conversation_store.get_total_message_tokens(
            conv.conversation_id, self._owner_id
        )
        total_sum_tokens = await self._summary_store.get_total_summary_tokens(
            conv.conversation_id, self._owner_id
        )
        current_tokens = total_msg_tokens + total_sum_tokens

        should = current_tokens > threshold_tokens

        return CompactionDecision(
            should_compact=should,
            current_tokens=current_tokens,
            threshold_tokens=threshold_tokens,
            reason=(
                f"Tokens {current_tokens} > threshold {threshold_tokens}"
                if should
                else f"Tokens {current_tokens} <= threshold {threshold_tokens}"
            ),
        )

    async def compact(
        self,
        mode: str = "leaf",
        token_budget: Optional[int] = None,
        summarize_fn: Optional[SummarizeFn] = None,
        conversation_id: Optional[int] = None,
        max_rounds: Optional[int] = None,
    ) -> CompactionResult:
        """Trigger compaction pipeline.

        Parameters
        ----------
        mode : str
            ``"leaf"`` for incremental leaf pass (default),
            ``"full"`` for a full sweep (leaf + condensation).
        token_budget : int, optional
            Token budget for compaction evaluation. Defaults to config.max_tokens.
        summarize_fn : SummarizeFn, optional
            Override the engine-level summarize function for this call.
        conversation_id : int, optional
            Conversation to compact. Falls back to self._conversation if not provided.

        Returns
        -------
        CompactionResult
            Result from compaction engine (models.CompactionResult).
        """
        if not self._initialized or not self._conversation:
            raise RuntimeError("HippoEngine not initialized")

        # Kill switch: skip compaction when master toggle is off or autocompact disabled
        if not self.config.enabled:
            logger.debug("compact() skipped: Hippo disabled (enabled=False)")
            return CompactionResult()

        # Use provided conversation_id or fall back to the engine's active conversation
        conv = self._conversation
        effective_conversation_id = conversation_id if conversation_id is not None else conv.conversation_id
        budget = token_budget or self.config.max_tokens
        summarize = summarize_fn or self._summarize_fn
        if summarize is None:
            logger.warning(
                "compact() called without summarize_fn — returning no-op result"
            )
            return CompactionResult()

        async def _do_compact() -> CompactionResult:
            import datetime as _dt
            assert self._compaction_engine is not None

            # B-021/B-022: Wire concurrency guard + DB compaction flags
            guard = self._compaction_guard
            async with guard.acquire(effective_conversation_id) as acquired:
                if not acquired:
                    logger.debug(
                        "compact() skipped: already in progress for conversation %d",
                        effective_conversation_id,
                    )
                    return CompactionResult()

                # B-022: Set DB-level compaction flag
                conn = self._db.conn if self._db else None
                if conn is not None:
                    await mark_compaction_start(conn, effective_conversation_id)
                try:
                    if mode == "full":
                        ce_result = await self._compaction_engine.compact_full_sweep(
                            conversation_id=effective_conversation_id,
                            token_budget=budget,
                            summarize=summarize,
                        )
                    else:
                        ce_result = await self._compaction_engine.compact_leaf(
                            conversation_id=effective_conversation_id,
                            token_budget=budget,
                            summarize=summarize,
                        )
                    if ce_result.action_taken:
                        self._last_compaction_at = _dt.datetime.utcnow().isoformat() + "Z"
                    # Convert compaction.types.CompactionResult → models.CompactionResult
                    return CompactionResult(
                        summaries_created=len(ce_result.created_summary_ids),
                        messages_compacted=0,  # calculated by compaction internals
                        tokens_before=ce_result.tokens_before,
                        tokens_after=ce_result.tokens_after,
                        rounds=ce_result.rounds,
                    )
                finally:
                    # B-022: Always clear DB flag
                    if conn is not None:
                        try:
                            await mark_compaction_end(conn, effective_conversation_id)
                        except Exception:
                            logger.warning(
                                "compact(): failed to clear compaction flag for conv %d",
                                effective_conversation_id,
                                exc_info=True,
                            )

        return await self._with_session_lock(effective_conversation_id, _do_compact)

    async def get_stats(self) -> ConversationStats:
        """Get statistics for the current conversation (#1235: extended fields)."""
        if not self._initialized or not self._conversation:
            raise RuntimeError("HippoEngine not initialized")

        conv = self._conversation
        msg_count = await self._conversation_store.get_message_count(
            conv.conversation_id, self._owner_id
        )
        sum_count = await self._summary_store.get_summary_count(
            conv.conversation_id, self._owner_id
        )
        max_depth = await self._summary_store.get_max_depth(
            conv.conversation_id, self._owner_id
        )
        total_msg_tokens = await self._conversation_store.get_total_message_tokens(
            conv.conversation_id, self._owner_id
        )
        total_sum_tokens = await self._summary_store.get_total_summary_tokens(
            conv.conversation_id, self._owner_id
        )
        ctx_count = await self._summary_store.get_context_item_count(
            conv.conversation_id
        )

        depth_dist = await self._summary_store.get_depth_distribution(conv.conversation_id, owner_id=self._owner_id)

        return ConversationStats(
            conversation_id=conv.conversation_id,
            message_count=msg_count,
            summary_count=sum_count,
            max_depth=max_depth,
            total_message_tokens=total_msg_tokens,
            total_summary_tokens=total_sum_tokens,
            context_item_count=ctx_count,
            last_compaction_at=self._last_compaction_at,
            summary_depth_distribution=depth_dist,
        )

    # ------------------------------------------------------------------
    # W3-4: bootstrap() and session reconciliation (#1165)
    # ------------------------------------------------------------------

    async def bootstrap(
        self,
        session_id: str,
        session_file: str,
        session_key: Optional[str] = None,
    ) -> dict:
        """Bootstrap a conversation from a session file.

        - Get or create conversation for session_id
        - If already bootstrapped → return early (no-op)
        - Read messages from session file, deduplicate, ingest
        - Prune HEARTBEAT_OK turns
        - Mark conversation as bootstrapped

        Returns dict with ``bootstrapped: bool`` and ``message_count: int``.
        """
        if not self._initialized:
            raise RuntimeError("HippoEngine not initialized — call initialize() first")

        conv = await self.conversation_store.get_or_create_conversation(
            self._owner_id, session_id, session_key=session_key
        )

        # Already bootstrapped → no-op
        if conv.bootstrapped_at:
            return {"bootstrapped": False, "message_count": 0, "conversation_id": conv.conversation_id}

        # Read messages from file
        file_messages = read_leaf_path_messages(session_file)
        if not file_messages:
            await self.conversation_store.mark_conversation_bootstrapped(
                conv.conversation_id
            )
            return {"bootstrapped": False, "message_count": 0, "conversation_id": conv.conversation_id}

        # Deduplicate against existing DB messages using a counter (multiset)
        # to handle repeated identical messages correctly (#1229).
        # A plain set would collapse all duplicates — e.g. three "ok" messages
        # map to the same identity, so re-bootstrap would skip ALL of them
        # even if only two were previously ingested.
        from collections import Counter

        existing = await self.conversation_store.get_messages(
            conv.conversation_id, self._owner_id
        )
        existing_identity_counts = Counter(
            message_identity(m.role, m.content) for m in existing
        )

        new_messages = []
        for m in file_messages:
            identity = message_identity(m.get("role", ""), _extract_message_content(m.get("content", "")))
            if existing_identity_counts[identity] > 0:
                existing_identity_counts[identity] -= 1  # consume one match
            else:
                new_messages.append(m)

        # Prune heartbeat turns before ingest
        new_messages = self._prune_heartbeat_ok_turns(new_messages)

        # Ingest new messages
        VALID_ROLES = {"user", "assistant", "system", "tool"}
        ROLE_MAP = {"toolResult": "tool", "function": "tool"}
        for msg in new_messages:
            raw_role = msg.get("role", "user")
            role = ROLE_MAP.get(raw_role, raw_role) if raw_role not in VALID_ROLES else raw_role
            if role not in VALID_ROLES:
                role = "assistant"  # align with lossless-claw fallback
            raw_content = msg.get("content", "")
            plain_content = _extract_message_content(raw_content)
            token_count = _estimate_tokens_for_content(raw_content)
            stored = await self.conversation_store.create_message(
                conversation_id=conv.conversation_id,
                owner_id=self._owner_id,
                role=role,
                content=plain_content,
                token_count=token_count,
            )
            # Write message parts (R-036 #1 fix)
            parts = _build_message_parts(role, raw_content)
            if parts:
                await self.conversation_store.create_message_parts(
                    stored.message_id, parts
                )
            await self.summary_store.append_context_item(
                conversation_id=conv.conversation_id,
                item_type="message",
                message_id=stored.message_id,
            )

        # Mark bootstrapped
        await self.conversation_store.mark_conversation_bootstrapped(
            conv.conversation_id
        )

        return {"bootstrapped": True, "message_count": len(new_messages), "conversation_id": conv.conversation_id}

    # ------------------------------------------------------------------
    # W3-5: evaluateLeafTrigger (#1166)
    # ------------------------------------------------------------------

    async def evaluate_leaf_trigger(
        self, messages: list[dict]
    ) -> tuple[bool, int]:
        """Check if live context tokens exceed leaf compaction threshold.

        Returns (should_trigger, token_count).
        """
        token_count = estimate_session_tokens(messages)
        threshold = self.config.context_window_tokens - self.config.leaf_trigger_margin
        return (token_count > threshold, token_count)

    # ------------------------------------------------------------------
    # W2-4: afterTurn dual-trigger (#1155)
    # ------------------------------------------------------------------

    async def after_turn(
        self,
        conversation_id: int,
        messages: list[dict],
        pre_prompt_count: int = 0,
        token_budget: Optional[int] = None,
        is_heartbeat: bool = False,
        session_key: Optional[str] = None,
    ) -> None:
        """Post-turn hook: ingest new messages + optionally trigger compaction.

        This is the primary integration point called by the gateway after each
        model turn completes.  It is best-effort — errors are logged but never
        propagated to the caller.

        Parameters
        ----------
        conversation_id : int
            Active conversation ID.
        messages : list[dict]
            The **full** message list for this turn (system + history + new).
        pre_prompt_count : int
            Index into *messages* where the new (post-prompt) messages begin.
        token_budget : int, optional
            If provided, triggers dual compaction evaluation after ingest.
        is_heartbeat : bool
            When True, prunes HEARTBEAT_OK turn-pairs before ingest.
        """

        async def _do_after_turn() -> None:
            new_messages = messages[pre_prompt_count:]

            # Kill switch: skip everything when disabled
            if not self.config.enabled:
                logger.debug("after_turn skipped: Hippo disabled")
                return

            # W2-3: prune heartbeat-OK turns before ingest (#1154)
            if is_heartbeat:
                new_messages = self._prune_heartbeat_ok_turns(new_messages)

            # --- Ingest new messages (best-effort per message) ---
            # Call store directly to avoid re-entering the session lock
            # (asyncio.Lock is not reentrant).
            # Use the conversation_id passed to after_turn, not self._conversation
            conv_id = conversation_id
            _VALID_ROLES = {"user", "assistant", "system", "tool"}
            _ROLE_MAP = {"toolResult": "tool", "function": "tool"}
            for msg in new_messages:
                try:
                    raw_role = msg.get("role", "user")
                    role = _ROLE_MAP.get(raw_role, raw_role) if raw_role not in _VALID_ROLES else raw_role
                    if role not in _VALID_ROLES:
                        role = "assistant"  # align with lossless-claw fallback
                    raw_content = msg.get("content", "")
                    plain_content = _extract_message_content(raw_content)
                    token_count = _estimate_tokens_for_content(raw_content)
                    stored = await self._conversation_store.create_message(
                        conversation_id=conv_id,
                        owner_id=self._owner_id,
                        role=role,
                        content=plain_content,
                        token_count=token_count,
                    )
                    # Write message parts (R-036 #1 fix)
                    parts = _build_message_parts(role, raw_content)
                    if parts:
                        await self._conversation_store.create_message_parts(
                            stored.message_id, parts
                        )
                    await self._summary_store.append_context_item(
                        conversation_id=conv_id,
                        item_type="message",
                        message_id=stored.message_id,
                    )
                    # H-13: record message for adaptive scheduling (#1191)
                    if self._scheduler is not None:
                        self._scheduler.record_message(conv_id)
                except Exception:
                    logger.warning(
                        "after_turn: failed to ingest message (role=%s), skipping",
                        msg.get("role"),
                        exc_info=True,
                    )

            # --- Compaction trigger ---
            if (
                token_budget is not None
                and self._compaction_engine is not None
                and not self.config.autocompact_disabled
            ):
                # H-13: adaptive scheduling (#1191) — replaces fixed threshold
                if self._scheduler is not None:
                    _, live_tokens = await self.evaluate_leaf_trigger(messages)
                    sched_decision = self._scheduler.evaluate(
                        conv_id, live_tokens, token_budget,
                    )
                    logger.debug(
                        "after_turn: scheduler decision=%s (%s)",
                        sched_decision.mode.value,
                        sched_decision.reason,
                    )

                    if sched_decision.mode == CompactionMode.SKIP:
                        return

                    # Budget gate — scheduler allows, but budget may block
                    if self._budget is not None:
                        budget_decision = self._budget.check_budget(
                            str(conv_id)
                        )
                        if not budget_decision.allowed:
                            logger.warning(
                                "after_turn: compaction skipped — %s",
                                budget_decision.reason,
                            )
                            return

                    summarize = self._summarize_fn
                    if summarize is None:
                        logger.debug(
                            "after_turn: compaction needed but no summarize_fn configured"
                        )
                        return

                    try:
                        if sched_decision.mode == CompactionMode.FORCE:
                            # Safety valve — aggressive compaction
                            async def _force_compact() -> None:
                                async def _do_force() -> None:
                                    assert self._compaction_engine is not None
                                    await self._compaction_engine.compact_until_under(
                                        conversation_id=conv_id,
                                        token_budget=token_budget,
                                        summarize=summarize,
                                    )
                                await self._with_session_lock(
                                    conv_id, _do_force
                                )

                            self._spawn_background(_force_compact())
                            logger.debug(
                                "after_turn: FORCE compaction fired (pressure=%.1f%%)",
                                sched_decision.token_pressure * 100,
                            )

                        elif sched_decision.mode == CompactionMode.FULL:
                            # Full sweep
                            async def _full_compact() -> None:
                                async def _do_full() -> None:
                                    assert self._compaction_engine is not None
                                    ce_result = await self._compaction_engine.compact_full_sweep(
                                        conversation_id=conv_id,
                                        token_budget=token_budget,
                                        summarize=summarize,
                                    )
                                    if self.config.auto_compaction_summary and ce_result and ce_result.action_taken:
                                        sums_created = len(ce_result.created_summary_ids)
                                        if sums_created > 0:
                                            summary_text = (
                                                f"Context was compacted. "
                                                f"{ce_result.rounds} rounds, "
                                                f"{sums_created} summaries created."
                                            )
                                            await self._conversation_store.create_message(
                                                conversation_id=conv_id,
                                                owner_id=self._owner_id,
                                                role="user",
                                                content=summary_text,
                                            )
                                await self._with_session_lock(
                                    conv_id, _do_full
                                )

                            self._spawn_background(_full_compact())
                            logger.debug(
                                "after_turn: FULL compaction fired (urgency=%.2f)",
                                sched_decision.urgency,
                            )

                        elif sched_decision.mode == CompactionMode.LIGHT:
                            # Single leaf pass
                            async def _light_compact() -> None:
                                async def _do_light() -> None:
                                    assert self._compaction_engine is not None
                                    await self._compaction_engine.compact_full_sweep(
                                        conversation_id=conv_id,
                                        token_budget=token_budget,
                                        summarize=summarize,
                                        force=True,
                                    )
                                    # LIGHT = single round, so we rely on compact_leaf
                                    # via compact_full_sweep with force=True; the
                                    # max_rounds config limits total work anyway.
                                await self._with_session_lock(
                                    conv_id, _do_light
                                )

                            self._spawn_background(_light_compact())
                            logger.debug(
                                "after_turn: LIGHT compaction fired (urgency=%.2f)",
                                sched_decision.urgency,
                            )
                    except Exception:
                        logger.warning(
                            "after_turn: adaptive compaction failed",
                            exc_info=True,
                        )
                else:
                    # Legacy fixed-threshold path (adaptive_scheduling_enabled=False)
                    # H-1: budget gate — skip compaction if breaker tripped (#1178)
                    if self._budget is not None:
                        budget_decision = self._budget.check_budget(
                            str(conv_id)
                        )
                        if not budget_decision.allowed:
                            logger.warning(
                                "after_turn: compaction skipped — %s",
                                budget_decision.reason,
                            )
                            return

                    # W3-5: evaluate_leaf_trigger gate — only proceed if live
                    # context exceeds the leaf trigger threshold (#1166)
                    should_fire, live_tokens = await self.evaluate_leaf_trigger(messages)
                    if not should_fire:
                        return

                    summarize = self._summarize_fn

                    # Leaf trigger via CompactionEngine.evaluate()
                    try:
                        decision = await self._compaction_engine.evaluate(
                            conversation_id=conv_id,
                            token_budget=token_budget,
                            observed_token_count=live_tokens,
                        )
                        if decision.should_compact and summarize is not None:
                            async def _compact_and_summarize() -> None:
                                # Acquire session lock to prevent races with
                                # concurrent ingest/after_turn calls (#1183)
                                async def _do_compact_bg() -> None:
                                    assert self._compaction_engine is not None
                                    ce_result = await self._compaction_engine.compact_leaf(
                                        conversation_id=conv_id,
                                        token_budget=token_budget,
                                        summarize=summarize,
                                    )
                                    if self.config.auto_compaction_summary and ce_result:
                                        msgs_compacted = ce_result.messages_compacted
                                        sums_created = len(ce_result.created_summary_ids)
                                        if sums_created > 0:
                                            summary_text = (
                                                f"Context was compacted. "
                                                f"{msgs_compacted} messages were summarized "
                                                f"into {sums_created} summaries."
                                            )
                                            await self._conversation_store.create_message(
                                                conversation_id=conv_id,
                                                owner_id=self._owner_id,
                                                role="user",
                                                content=summary_text,
                                            )

                                await self._with_session_lock(
                                    conv_id, _do_compact_bg
                                )

                            self._spawn_background(_compact_and_summarize())
                            logger.debug(
                                "after_turn: leaf compaction fired (live=%d, threshold=%d)",
                                live_tokens,
                                decision.threshold,
                            )
                        elif decision.should_compact:
                            logger.debug(
                                "after_turn: compaction needed but no summarize_fn configured"
                            )
                    except Exception:
                        logger.warning(
                            "after_turn: leaf compaction check failed",
                            exc_info=True,
                        )

        try:
            await self._with_session_lock(conversation_id, _do_after_turn)
        except Exception:
            logger.warning(
                "after_turn: unexpected error for conversation %d",
                conversation_id,
                exc_info=True,
            )

    # ------------------------------------------------------------------
    # W2-3: Heartbeat pruning helper (#1154)
    # ------------------------------------------------------------------

    @staticmethod
    def _turn_has_tool_calls(messages: list[dict], start: int, end: int) -> bool:
        """Check if any message in [start, end) contains tool call parts."""
        for idx in range(start, min(end, len(messages))):
            msg = messages[idx]
            # Check for tool_call_id on the message itself
            if msg.get("tool_call_id"):
                return True
            # Check for tool_calls list (OpenAI-style)
            if msg.get("tool_calls"):
                return True
            # Check content parts for tool_use / tool_result
            content = msg.get("content")
            if isinstance(content, list):
                for part in content:
                    if isinstance(part, dict) and part.get("type") in (
                        "tool_use", "tool_result", "tool_call",
                    ):
                        return True
            # Check metadata for tool call markers
            metadata = msg.get("metadata") or {}
            if isinstance(metadata, dict) and metadata.get("tool_call_id"):
                return True
        return False

    @staticmethod
    def _is_heartbeat_turn(user_msg: dict, assistant_msg: dict) -> bool:
        """Determine if a user+assistant pair is a heartbeat turn.

        Detection priority:
        1. metadata.heartbeat_ack flag on the assistant message → prune
        2. String-based fallback: assistant content exactly "HEARTBEAT_OK"
           AND assistant has no tool_call parts
        """
        # Priority 1: metadata flag
        metadata = assistant_msg.get("metadata") or {}
        if isinstance(metadata, dict) and metadata.get("heartbeat_ack"):
            return True

        # Priority 2: exact string match (strict)
        assistant_content = assistant_msg.get("content", "")
        # Only match plain string content, not structured content with tool parts
        if not isinstance(assistant_content, str):
            return False
        if assistant_content.strip() != "HEARTBEAT_OK":
            return False

        # Ensure assistant has no tool calls
        if assistant_msg.get("tool_calls"):
            return False

        return True

    def _prune_heartbeat_ok_turns(self, messages: list[dict]) -> list[dict]:
        """Remove heartbeat turn-pairs from a message list.

        A heartbeat turn-pair is:
          1. A user message (any content)
          2. Immediately followed by an assistant heartbeat response

        Heartbeat detection uses metadata.heartbeat_ack flag first,
        then falls back to exact "HEARTBEAT_OK" string match (with
        additional safety: no tool calls allowed in the turn).

        Respects config.prune_heartbeats kill switch.
        Returns a new list with heartbeat pairs removed.
        """
        # Kill switch
        if not self.config.prune_heartbeats:
            return messages

        pruned: list[dict] = []
        i = 0
        while i < len(messages):
            if (
                i + 1 < len(messages)
                and messages[i].get("role") == "user"
                and messages[i + 1].get("role") == "assistant"
            ):
                # Skip if any message in the turn has tool calls
                if self._turn_has_tool_calls(messages, i, i + 2):
                    pruned.append(messages[i])
                    i += 1
                    continue

                if self._is_heartbeat_turn(messages[i], messages[i + 1]):
                    # Prune both messages
                    i += 2
                    continue

            pruned.append(messages[i])
            i += 1
        return pruned

    async def shutdown(self) -> None:
        """Clean up database connections."""
        # Cancel and await any in-flight background tasks before closing the DB,
        # so they don't run against a closed connection.
        if self._background_tasks:
            tasks = list(self._background_tasks)
            for task in tasks:
                task.cancel()
            await asyncio.gather(*tasks, return_exceptions=True)
        if self._db:
            await self._db.close()
            self._db = None
        self._initialized = False
        logger.info("HippoEngine shut down")
