"""Session message compression and rehydration for efficient context loading.

This module implements message storage and compression to keep conversation history
within context windows while preserving full content via REM LOOKUP.

Message Types and Storage Strategy
===================================

All messages are stored UNCOMPRESSED in the database for full audit/analysis.
Compression happens only on RELOAD when reconstructing context for the LLM.

Message Types:
- `user`: User messages - stored and reloaded as-is
- `tool`: Tool call messages (e.g., register_metadata) - stored and reloaded as-is
         NEVER compressed - contains important structured metadata
- `assistant`: Assistant text responses - stored uncompressed, but MAY BE
              compressed on reload if long (>400 chars) with REM LOOKUP hints

Example Session Flow:
```
Turn 1 (stored uncompressed):
  - user: "I have a headache"
  - tool: register_metadata({confidence: 0.3, collected_fields: {...}})
  - assistant: "I'm sorry to hear that. How long has this been going on?"

Turn 2 (stored uncompressed):
  - user: "About 3 days, really bad"
  - tool: register_metadata({confidence: 0.6, collected_fields: {...}})
  - assistant: "Got it - 3 days. On a scale of 1-10..."

On reload (for LLM context):
  - user messages: returned as-is
  - tool messages: returned as-is (never compressed)
  - assistant messages: compressed if long, with REM LOOKUP hint for full retrieval
```

REM LOOKUP Pattern:
- Long assistant messages get truncated with hint: "... [REM LOOKUP session-{id}-msg-{idx}] ..."
- Agent can retrieve full content on-demand using the LOOKUP key
- Keeps context window efficient while preserving data integrity

Token Threshold Safety:
- filter_within_token_threshold() ensures loaded messages fit within context window
- Estimates tokens as chars/4 (reasonable approximation)
- When truncating, finds latest moment for session and inserts it temporally
- Moment summary provides lookback to older memories that were truncated

Key Design Decisions:
1. Store everything uncompressed - full audit trail in database
2. Compress only on reload - optimize for LLM context window
3. Never compress tool messages - structured metadata must stay intact
4. REM LOOKUP enables on-demand retrieval of full assistant responses
5. Token threshold safety prevents context overflow with moment-based fallback
"""

import json
from typing import Any

from loguru import logger

# Max length for entity keys (kv_store.entity_key is varchar(255))
MAX_ENTITY_KEY_LENGTH = 255


def truncate_key(key: str, max_length: int = MAX_ENTITY_KEY_LENGTH) -> str:
    """Truncate a key to max length, preserving useful suffix if possible."""
    if len(key) <= max_length:
        return key
    # Keep first part and add hash suffix for uniqueness
    import hashlib
    hash_suffix = hashlib.md5(key.encode()).hexdigest()[:8]
    truncated = key[:max_length - 9] + "-" + hash_suffix
    logger.warning(f"Truncated key from {len(key)} to {len(truncated)} chars: {key[:50]}...")
    return truncated

from rem.models.entities import Message, Session
from rem.services.postgres import PostgresService, Repository
from rem.settings import settings


class MessageCompressor:
    """Compress and decompress session messages with REM lookup keys."""

    def __init__(self, truncate_length: int = 200):
        """
        Initialize message compressor.

        Args:
            truncate_length: Number of characters to keep from start/end (default: 200)
        """
        self.truncate_length = truncate_length
        self.min_length_for_compression = truncate_length * 2

    def compress_message(
        self, message: dict[str, Any], entity_key: str | None = None
    ) -> dict[str, Any]:
        """
        Compress a message by truncating long content and adding REM lookup key.

        Args:
            message: Message dict with role and content
            entity_key: Optional REM lookup key for full message recovery

        Returns:
            Compressed message dict
        """
        content = message.get("content") or ""

        # Don't compress short messages or system messages
        if (
            len(content) <= self.min_length_for_compression
            or message.get("role") == "system"
        ):
            return message.copy()

        # Compress long messages
        n = self.truncate_length
        start = content[:n]
        end = content[-n:]

        # Create compressed content with REM lookup hint
        if entity_key:
            compressed_content = f"{start}\n\n... [Message truncated - REM LOOKUP {entity_key} to recover full content] ...\n\n{end}"
        else:
            compressed_content = f"{start}\n\n... [Message truncated - {len(content) - 2*n} characters omitted] ...\n\n{end}"

        compressed_message = message.copy()
        compressed_message["content"] = compressed_content
        compressed_message["_compressed"] = True
        compressed_message["_original_length"] = len(content)
        if entity_key:
            compressed_message["_entity_key"] = entity_key

        logger.debug(
            f"Compressed message from {len(content)} to {len(compressed_content)} chars (key={entity_key})"
        )

        return compressed_message

    def decompress_message(
        self, message: dict[str, Any], full_content: str
    ) -> dict[str, Any]:
        """
        Decompress a message by restoring full content.

        Args:
            message: Compressed message dict
            full_content: Full content to restore

        Returns:
            Decompressed message dict
        """
        decompressed = message.copy()
        decompressed["content"] = full_content
        decompressed.pop("_compressed", None)
        decompressed.pop("_original_length", None)
        decompressed.pop("_entity_key", None)

        return decompressed

    def is_compressed(self, message: dict[str, Any]) -> bool:
        """Check if a message is compressed."""
        return message.get("_compressed", False)

    def get_entity_key(self, message: dict[str, Any]) -> str | None:
        """Get REM lookup key from compressed message."""
        return message.get("_entity_key")


class SessionMessageStore:
    """Store and retrieve session messages with compression."""

    def __init__(
        self,
        user_id: str,
        compressor: MessageCompressor | None = None,
    ):
        """
        Initialize session message store.

        Args:
            user_id: User identifier for data isolation
            compressor: Optional message compressor (creates default if None)
        """
        self.user_id = user_id
        self.compressor = compressor or MessageCompressor()
        self.repo = Repository(Message)
        self._session_repo = Repository(Session, table_name="sessions")

    async def _ensure_session_exists(
        self,
        session_id: str,
        user_id: str | None = None,
    ) -> None:
        """
        Ensure session exists, creating it if necessary.

        Args:
            session_id: Session UUID from X-Session-Id header
            user_id: Optional user identifier
        """
        try:
            # Check if session already exists by UUID
            existing = await self._session_repo.get_by_id(session_id)
            if existing:
                return  # Session already exists

            # Create new session with the provided UUID as id
            session = Session(
                id=session_id,  # Use the provided UUID as session id
                name=session_id,  # Default name to UUID, can be updated later
                user_id=user_id or self.user_id,
                tenant_id=self.user_id,  # tenant_id set to user_id for scoping
            )
            await self._session_repo.upsert(session)
            logger.info(f"Created session {session_id} for user {user_id or self.user_id}")

        except Exception as e:
            # Log but don't fail - session creation is best-effort
            logger.warning(f"Failed to ensure session exists: {e}")

    async def store_message(
        self,
        session_id: str,
        message: dict[str, Any],
        message_index: int,
        user_id: str | None = None,
    ) -> str:
        """
        Store a long assistant message as a Message entity for REM lookup.

        Args:
            session_id: Parent session identifier
            message: Message dict to store
            message_index: Index of message in conversation
            user_id: Optional user identifier

        Returns:
            Entity key for REM lookup (message ID)
        """
        if not settings.postgres.enabled:
            logger.debug("Postgres disabled, skipping message storage")
            return f"msg-{message_index}"

        # Create entity key for REM LOOKUP: session-{session_id}-msg-{index}
        # Truncate to avoid exceeding kv_store.entity_key varchar(255) limit
        entity_key = truncate_key(f"session-{session_id}-msg-{message_index}")

        # Create Message entity for assistant response
        # Use pre-generated id from message dict if available (for frontend feedback)
        msg = Message(
            id=message.get("id"),  # Use pre-generated ID if provided
            content=message.get("content") or "",
            message_type=message.get("role", "assistant"),
            session_id=session_id,
            tenant_id=self.user_id,  # Set tenant_id to user_id (application scoped to user)
            user_id=user_id or self.user_id,
            trace_id=message.get("trace_id"),
            span_id=message.get("span_id"),
            metadata={
                "message_index": message_index,
                "entity_key": entity_key,  # Store entity key for LOOKUP
                "timestamp": message.get("timestamp"),
            },
        )

        # Store in database
        await self.repo.upsert(msg)

        logger.debug(f"Stored assistant response: {entity_key} (id={msg.id})")
        return entity_key

    async def retrieve_message(self, entity_key: str) -> str | None:
        """
        Retrieve full message content by REM lookup key.

        Uses LOOKUP query pattern: finds message by entity_key in metadata.

        Args:
            entity_key: REM lookup key (session-{id}-msg-{index})

        Returns:
            Full message content or None if not found
        """
        if not settings.postgres.enabled:
            logger.debug("Postgres disabled, cannot retrieve message")
            return None

        try:
            # LOOKUP pattern: find message by entity_key in metadata
            query = """
                SELECT * FROM messages
                WHERE metadata->>'entity_key' = $1
                  AND user_id = $2
                  AND deleted_at IS NULL
                LIMIT 1
            """

            if not self.repo.db:
                logger.warning("Database not available for message lookup")
                return None

            row = await self.repo.db.fetchrow(query, entity_key, self.user_id)

            if row:
                msg = Message.model_validate(dict(row))
                logger.debug(f"Retrieved message via LOOKUP: {entity_key}")
                return msg.content

            logger.warning(f"Message not found via LOOKUP: {entity_key}")
            return None

        except Exception as e:
            logger.error(f"Failed to retrieve message {entity_key}: {e}")
            return None

    async def store_session_messages(
        self,
        session_id: str,
        messages: list[dict[str, Any]],
        user_id: str | None = None,
        compress: bool = True,
    ) -> list[dict[str, Any]]:
        """
        Store all session messages and return compressed versions.

        Ensures session exists before storing messages.

        Args:
            session_id: Session UUID
            messages: List of messages to store
            user_id: Optional user identifier
            compress: Whether to compress messages (default: True)

        Returns:
            List of compressed messages with REM lookup keys
        """
        if not settings.postgres.enabled:
            logger.debug("Postgres disabled, returning messages uncompressed")
            return messages

        # Ensure session exists before storing messages
        await self._ensure_session_exists(session_id, user_id)

        compressed_messages = []

        for idx, message in enumerate(messages):
            content = message.get("content") or ""

            # Only store and compress long assistant responses
            if (
                message.get("role") == "assistant"
                and len(content) > self.compressor.min_length_for_compression
            ):
                # Store full message as separate Message entity
                entity_key = await self.store_message(
                    session_id, message, idx, user_id
                )

                if compress:
                    compressed_msg = self.compressor.compress_message(
                        message, entity_key
                    )
                    compressed_messages.append(compressed_msg)
                else:
                    msg_copy = message.copy()
                    msg_copy["_entity_key"] = entity_key
                    compressed_messages.append(msg_copy)
            else:
                # Short assistant messages, user messages, tool messages, and system messages stored as-is
                # Store ALL messages in database for full audit trail
                # Build metadata dict with standard fields
                msg_metadata = {
                    "message_index": idx,
                    "timestamp": message.get("timestamp"),
                }

                # For tool messages, include tool call details in metadata
                # Note: tool_arguments is stored only when provided (parent tool calls)
                # For child tool calls (e.g., register_metadata), args are in content as JSON
                if message.get("role") == "tool":
                    if message.get("tool_call_id"):
                        msg_metadata["tool_call_id"] = message.get("tool_call_id")
                    if message.get("tool_name"):
                        msg_metadata["tool_name"] = message.get("tool_name")
                    if message.get("tool_arguments"):
                        msg_metadata["tool_arguments"] = message.get("tool_arguments")

                msg = Message(
                    id=message.get("id"),  # Use pre-generated ID if provided
                    content=content,
                    message_type=message.get("role", "user"),
                    session_id=session_id,
                    tenant_id=self.user_id,  # Set tenant_id to user_id (application scoped to user)
                    user_id=user_id or self.user_id,
                    trace_id=message.get("trace_id"),
                    span_id=message.get("span_id"),
                    metadata=msg_metadata,
                )
                await self.repo.upsert(msg)
                compressed_messages.append(message.copy())

        return compressed_messages

    async def load_session_messages(
        self,
        session_id: str,
        user_id: str | None = None,
        compress_on_load: bool = True,
        max_messages: int | None = None,
    ) -> tuple[list[dict[str, Any]], bool]:
        """
        Load session messages from database, optionally compressing long assistant messages.

        Compression on Load:
        - Tool messages (role: "tool") are NEVER compressed - they contain structured metadata
        - User messages are returned as-is
        - Assistant messages MAY be compressed if long (>400 chars) with REM LOOKUP hints

        CTE Query Pattern (when max_messages is set):
        - Uses CTE to select last N messages (DESC order)
        - Returns them in conversation order (ASC order)
        - Efficiently limits context while preserving chronological flow

        Args:
            session_id: Session identifier
            user_id: Optional user identifier for filtering
            compress_on_load: Whether to compress long assistant messages (default: True)
            max_messages: Optional limit on messages to load (uses CTE for efficiency)

        Returns:
            Tuple of:
            - List of session messages in chronological order
            - Boolean indicating if a session_partition event was found in loaded messages
        """
        if not settings.postgres.enabled:
            logger.debug("Postgres disabled, returning empty message list")
            return [], False

        try:
            # Use CTE query when max_messages is specified for efficient limiting
            if max_messages is not None:
                return await self._load_with_cte(
                    session_id, user_id, compress_on_load, max_messages
                )

            # Standard query (load all messages)
            # Note: tenant_id column in messages table maps to user_id (user-scoped partitioning)
            filters = {"session_id": session_id, "tenant_id": self.user_id}
            if user_id:
                filters["user_id"] = user_id

            messages = await self.repo.find(filters, order_by="created_at ASC")

            # Convert Message entities to dict format
            message_dicts = []
            has_partition_event = False

            for idx, msg in enumerate(messages):
                role = msg.message_type or "assistant"
                msg_dict = {
                    "role": role,
                    "content": msg.content,
                    "timestamp": msg.created_at.isoformat() if msg.created_at else None,
                }

                # For tool messages, reconstruct tool call metadata
                # Note: tool_arguments may be in metadata (parent calls) or parsed from
                # content (child calls like register_metadata) by pydantic_messages.py
                if role == "tool" and msg.metadata:
                    if msg.metadata.get("tool_call_id"):
                        msg_dict["tool_call_id"] = msg.metadata["tool_call_id"]
                    if msg.metadata.get("tool_name"):
                        msg_dict["tool_name"] = msg.metadata["tool_name"]
                        # Check for partition event
                        if msg.metadata["tool_name"] == "session_partition":
                            has_partition_event = True
                    if msg.metadata.get("tool_arguments"):
                        msg_dict["tool_arguments"] = msg.metadata["tool_arguments"]

                # Compress long ASSISTANT messages on load (never tool messages)
                if (
                    compress_on_load
                    and role == "assistant"
                    and len(msg.content) > self.compressor.min_length_for_compression
                ):
                    # Generate entity key for REM LOOKUP
                    entity_key = truncate_key(f"session-{session_id}-msg-{idx}")
                    msg_dict = self.compressor.compress_message(msg_dict, entity_key)

                message_dicts.append(msg_dict)

            logger.debug(
                f"Loaded {len(message_dicts)} messages for session {session_id} "
                f"(compress_on_load={compress_on_load}, has_partition={has_partition_event})"
            )
            return message_dicts, has_partition_event

        except Exception as e:
            logger.error(f"Failed to load session messages: {e}")
            return [], False

    async def _load_with_cte(
        self,
        session_id: str,
        user_id: str | None,
        compress_on_load: bool,
        max_messages: int,
    ) -> tuple[list[dict[str, Any]], bool]:
        """
        Load last N messages using CTE query, returned in conversation order.

        SQL Pattern:
        WITH recent_messages AS (
            SELECT * FROM messages
            WHERE session_id = $1 AND user_id = $2
            ORDER BY created_at DESC
            LIMIT $3
        )
        SELECT * FROM recent_messages ORDER BY created_at ASC;

        This efficiently fetches only the most recent messages while
        returning them in chronological order for context building.
        """
        from rem.services.postgres import get_postgres_service

        postgres = get_postgres_service()
        if not postgres:
            return [], False

        await postgres.connect()
        try:
            effective_user_id = user_id or self.user_id

            # CTE query: get last N messages, return in chronological order
            query = """
                WITH recent_messages AS (
                    SELECT *
                    FROM messages
                    WHERE session_id = $1
                      AND user_id = $2
                      AND deleted_at IS NULL
                    ORDER BY created_at DESC
                    LIMIT $3
                )
                SELECT * FROM recent_messages
                ORDER BY created_at ASC
            """

            rows = await postgres.fetch(query, session_id, effective_user_id, max_messages)

            # Convert rows to message dicts
            message_dicts = []
            has_partition_event = False

            for idx, row in enumerate(rows):
                role = row["message_type"] or "assistant"
                content = row["content"] or ""
                # Handle metadata - might be a JSON string or dict
                metadata_raw = row["metadata"] or {}
                if isinstance(metadata_raw, str):
                    try:
                        metadata = json.loads(metadata_raw)
                    except json.JSONDecodeError:
                        metadata = {}
                else:
                    metadata = metadata_raw

                msg_dict = {
                    "role": role,
                    "content": content,
                    "timestamp": row["created_at"].isoformat() if row["created_at"] else None,
                }

                # For tool messages, reconstruct tool call metadata
                if role == "tool" and metadata:
                    if metadata.get("tool_call_id"):
                        msg_dict["tool_call_id"] = metadata["tool_call_id"]
                    if metadata.get("tool_name"):
                        msg_dict["tool_name"] = metadata["tool_name"]
                        # Check for partition event
                        if metadata["tool_name"] == "session_partition":
                            has_partition_event = True
                    if metadata.get("tool_arguments"):
                        msg_dict["tool_arguments"] = metadata["tool_arguments"]

                # Also check tool_name column directly
                if row.get("tool_name") == "session_partition":
                    has_partition_event = True

                # Compress long ASSISTANT messages on load (never tool messages)
                if (
                    compress_on_load
                    and role == "assistant"
                    and len(content) > self.compressor.min_length_for_compression
                ):
                    entity_key = truncate_key(f"session-{session_id}-msg-{idx}")
                    msg_dict = self.compressor.compress_message(msg_dict, entity_key)

                message_dicts.append(msg_dict)

            logger.debug(
                f"Loaded {len(message_dicts)} messages via CTE for session {session_id} "
                f"(max={max_messages}, has_partition={has_partition_event})"
            )
            return message_dicts, has_partition_event

        except Exception as e:
            logger.error(f"CTE load failed for session {session_id}: {e}")
            return [], False
        # NOTE: Do NOT disconnect() here - pool reuse improves TTFT

    async def retrieve_full_message(self, session_id: str, message_index: int) -> str | None:
        """
        Retrieve full message content by session and message index (for REM LOOKUP).

        This is used when an agent needs to recover full content from a compressed
        message that has a REM LOOKUP hint.

        Args:
            session_id: Session identifier
            message_index: Index of message in session (from REM LOOKUP key)

        Returns:
            Full message content or None if not found
        """
        entity_key = truncate_key(f"session-{session_id}-msg-{message_index}")
        return await self.retrieve_message(entity_key)


async def calculate_session_tokens(
    session_id: str,
    user_id: str,
    last_moment_idx: int = 0,
    model: str | None = None,
) -> int:
    """
    Calculate total tokens for unprocessed messages in a session using tiktoken.

    This provides accurate token counting for moment builder trigger decisions,
    replacing the defunct session.total_tokens field which was never updated.

    Args:
        session_id: Session identifier
        user_id: User identifier for tenant filtering
        last_moment_idx: Index of last message included in a moment (0 if no moments)
        model: Optional model name for tiktoken encoding

    Returns:
        Total tokens for messages after last_moment_idx
    """
    from rem.services.postgres import get_postgres_service
    from rem.settings import settings
    from rem.utils.agentic_chunking import estimate_tokens

    if not settings.postgres.enabled:
        return 0

    postgres = get_postgres_service()
    if not postgres:
        return 0

    await postgres.connect()
    try:
        # Query message content for unprocessed messages
        # Using row_number to identify messages after last_moment_idx
        query = """
            WITH numbered AS (
                SELECT
                    content,
                    ROW_NUMBER() OVER (ORDER BY created_at) as rn
                FROM messages
                WHERE session_id = $1
                  AND user_id = $2
                  AND deleted_at IS NULL
            )
            SELECT content
            FROM numbered
            WHERE rn > $3
        """
        rows = await postgres.fetch(query, (session_id, user_id, last_moment_idx))

        total_tokens = 0
        for row in rows:
            content = row.get("content") or ""
            # Add ~4 tokens overhead per message for role/structure
            total_tokens += estimate_tokens(content, model) + 4

        logger.debug(
            f"Calculated {total_tokens} tokens for session {session_id} "
            f"({len(rows)} messages after idx {last_moment_idx})"
        )
        return total_tokens

    except Exception as e:
        logger.error(f"Token calculation failed for session {session_id}: {e}")
        return 0
    # NOTE: Do NOT disconnect() here - pool reuse improves TTFT


async def _get_session_token_sum_sql(session_id: str, user_id: str) -> int | None:
    """
    Get total token count for session using SQL SUM (fast path).

    Returns None if token_count column has NULLs or query fails.
    """
    from rem.services.postgres import get_postgres_service

    postgres = get_postgres_service()
    if not postgres:
        return None

    try:
        await postgres.connect()
        query = """
            SELECT
                SUM(COALESCE(token_count, 0)) as total_tokens,
                COUNT(*) FILTER (WHERE token_count IS NULL) as null_count
            FROM messages
            WHERE session_id = $1
              AND user_id = $2
              AND deleted_at IS NULL
        """
        row = await postgres.fetchrow(query, session_id, user_id)
        if row and row["null_count"] == 0:
            # All messages have token_count populated
            return row["total_tokens"] or 0
        return None  # Some messages missing token_count, fall back to Python
    except Exception as e:
        logger.debug(f"SQL token sum failed: {e}")
        return None
    # NOTE: Do NOT disconnect() here - pool reuse improves TTFT


async def filter_within_token_threshold(
    messages: list[dict[str, Any]],
    session_id: str,
    user_id: str,
    token_threshold: int,
    model: str | None = None,
) -> tuple[list[dict[str, Any]], int]:
    """
    Filter messages to fit within token threshold, inserting moment summary if truncated.

    Uses SQL-based fast path when token_count is populated in DB, otherwise falls back
    to tiktoken estimation. When messages are truncated, finds the latest moment for
    the session and inserts it to maintain lookback capability to older memories.

    Args:
        messages: List of message dicts (chronological order, oldest first)
        session_id: Session ID for moment lookup
        user_id: User ID for moment lookup
        token_threshold: Maximum tokens allowed (from settings.moment_builder.token_threshold)
        model: Optional model name for tiktoken encoding

    Returns:
        Tuple of (filtered messages, total tokens used)

    Example:
        >>> messages = [msg1, msg2, msg3, ...]  # 150K tokens
        >>> filtered, tokens = await filter_within_token_threshold(
        ...     messages, session_id, user_id, token_threshold=90000
        ... )
        >>> # Returns most recent messages + moment summary for truncated portion
    """
    from rem.utils.agentic_chunking import estimate_tokens

    if not messages:
        return [], 0

    # OPTIMIZATION: Try SQL-based fast path first
    # If all messages have token_count populated, use SUM from DB
    sql_total = await _get_session_token_sum_sql(session_id, user_id)
    if sql_total is not None and sql_total <= token_threshold:
        # Fast path: SQL confirms we're under threshold, skip O(n) iteration
        logger.debug(f"SQL fast path: {sql_total} tokens <= {token_threshold} threshold")
        return messages, sql_total

    # Fall back to Python-based token counting
    # Calculate tokens for each message (content + role overhead ~4 tokens)
    message_tokens = []
    for msg in messages:
        content = msg.get("content") or ""
        tokens = estimate_tokens(content, model) + 4  # overhead for role/structure
        message_tokens.append(tokens)

    total_tokens = sum(message_tokens)

    # If within threshold, return as-is
    if total_tokens <= token_threshold:
        logger.debug(f"Messages within token threshold: {total_tokens}/{token_threshold}")
        return messages, total_tokens

    # Need to truncate - work backwards from most recent to find cutoff
    # Reserve 10% of threshold for moment summary
    available_tokens = int(token_threshold * 0.9)
    accumulated = 0
    cutoff_idx = len(messages)

    for i in range(len(messages) - 1, -1, -1):
        if accumulated + message_tokens[i] > available_tokens:
            cutoff_idx = i + 1  # Keep messages from i+1 onwards
            break
        accumulated += message_tokens[i]

    # If we'd keep all messages, just return them
    if cutoff_idx == 0:
        return messages, total_tokens

    # Truncate older messages
    kept_messages = messages[cutoff_idx:]
    truncated_count = cutoff_idx

    logger.info(
        f"Token threshold exceeded ({total_tokens}/{token_threshold}), "
        f"truncating {truncated_count} oldest messages, keeping {len(kept_messages)}"
    )

    # Check if a moment system message is already in the kept messages
    # (e.g., from a previous load that already included it)
    has_moment = any(
        msg.get("_moment_id") or "[CONVERSATION HISTORY SUMMARY]" in (msg.get("content") or "")
        for msg in kept_messages
        if msg.get("role") == "system"
    )

    if has_moment:
        logger.debug("Moment summary already present in messages, skipping moment lookup")
        return kept_messages, accumulated

    # Try to find moment summary for the truncated portion
    moment_message = await _get_moment_for_truncated_messages(
        session_id, user_id, messages[:cutoff_idx]
    )

    if moment_message:
        # Insert moment at the beginning (it summarizes the truncated older messages)
        kept_messages = [moment_message] + kept_messages
        moment_tokens = estimate_tokens(moment_message.get("content", ""), model)
        accumulated += moment_tokens
        logger.info(
            f"Inserted moment summary ({moment_tokens} tokens) for {truncated_count} truncated messages"
        )

    return kept_messages, accumulated


async def _get_moment_for_truncated_messages(
    session_id: str,
    user_id: str,
    truncated_messages: list[dict[str, Any]],
) -> dict[str, Any] | None:
    """
    Find the latest moment that covers the truncated messages.

    Queries the moments table for session-compression moments linked to this session,
    ordered by ends_timestamp DESC, and returns the most recent one as a system message.

    Args:
        session_id: Session ID to find moments for
        user_id: User ID for tenant filtering
        truncated_messages: The messages being truncated (for timestamp reference)

    Returns:
        A system message dict with the moment summary, or None if no moment found
    """
    from rem.services.postgres import get_postgres_service
    from rem.settings import settings

    if not settings.postgres.enabled:
        return None

    postgres = get_postgres_service()
    if not postgres:
        return None

    await postgres.connect()
    try:
        # Query for the latest session-compression moment for this session
        query = """
            SELECT id, name, summary, starts_timestamp, ends_timestamp, previous_moment_keys
            FROM moments
            WHERE source_session_id = $1
              AND user_id = $2
              AND category = 'session-compression'
              AND deleted_at IS NULL
            ORDER BY ends_timestamp DESC
            LIMIT 1
        """

        row = await postgres.fetchrow(query, session_id, user_id)

        if not row:
            logger.debug(f"No moment found for session {session_id}")
            return None

        summary = row["summary"]
        moment_name = row["name"]
        ends_ts = row["ends_timestamp"]
        previous_keys = row["previous_moment_keys"] or []

        # Build a system message with the moment summary
        # Include hint about previous moments for deeper lookback
        content = f"[CONVERSATION HISTORY SUMMARY]\n{summary}"

        if previous_keys:
            content += f"\n\n[For older history, see previous moments: {', '.join(previous_keys[:3])}]"

        logger.info(
            f"Found moment '{moment_name}' (ends: {ends_ts}) for truncated session context"
        )

        return {
            "role": "system",
            "content": content,
            "_moment_id": str(row["id"]),
            "_moment_name": moment_name,
        }

    except Exception as e:
        logger.error(f"Failed to get moment for truncated messages: {e}")
        return None
    # NOTE: Do NOT disconnect() here - pool reuse improves TTFT
