"""Short-term memory: transcript token budgeting (build plan Section 6.2).

Long-term (semantic) and episodic memory are backed by Qdrant in Phase 3.
"""

from __future__ import annotations

from legate.llm.base import LLMMessage

#: Default working-context budget (tokens) before the transcript is trimmed.
DEFAULT_CONTEXT_TOKENS = 8_000


def estimate_tokens(text: str) -> int:
    """Cheap token estimate (~4 chars/token) good enough for budgeting."""
    return max(1, len(text) // 4)


def transcript_tokens(messages: list[LLMMessage]) -> int:
    return sum(estimate_tokens(m.content) for m in messages)


def trim_transcript(messages: list[LLMMessage], budget: int) -> list[LLMMessage]:
    """Drop the oldest non-system messages until the transcript fits ``budget``.

    System messages are always retained (instructions only ever come from the
    agent config, never from retrieved content).
    """
    if transcript_tokens(messages) <= budget:
        return list(messages)

    system = [m for m in messages if m.role == "system"]
    rest = [m for m in messages if m.role != "system"]

    base = transcript_tokens(system)
    kept: list[LLMMessage] = []
    running = base
    # Keep the most recent messages, walking backwards.
    for message in reversed(rest):
        cost = estimate_tokens(message.content)
        if running + cost > budget and kept:
            break
        running += cost
        kept.append(message)
    kept.reverse()
    return system + kept
