"""Rolling message-history compaction — the context safety net.

We let the research history grow into the model's full context window (131k+ on gpt-oss,
absorbed by oMLX's paged SSD KV cache) and only compact when the prompt actually approaches
that limit. When it does, we fold the older turns into a running text summary and keep the
most recent turns verbatim.
"""

from __future__ import annotations

from langchain_core.messages import HumanMessage, SystemMessage

COMPACT_RATIO = 0.9  # compact once the prompt reaches this fraction of the context window
KEEP_LAST = 20  # verbatim recent turns to retain after compaction
_CHARS_PER_TOKEN = 4  # rough fallback when the server doesn't report token counts


def estimate_tokens(messages: list) -> int:
    chars = sum(
        len(m.content if isinstance(m.content, str) else str(m.content)) for m in messages
    )
    return chars // _CHARS_PER_TOKEN


def should_compact(prompt_tokens: int, context_length: int, ratio: float = COMPACT_RATIO) -> bool:
    return prompt_tokens >= ratio * context_length


def _render(messages: list) -> str:
    out = []
    for m in messages:
        role = m.__class__.__name__.replace("Message", "")
        content = m.content if isinstance(m.content, str) else str(m.content)
        out.append(f"{role}: {content[:800]}")
    return "\n".join(out)


def compact(summarizer_llm, messages: list, running_summary: str, keep_last: int = KEEP_LAST):
    """Return (new_messages, new_summary). Keeps system[0] + a rolling summary + recent turns."""
    if len(messages) <= keep_last + 1:
        return messages, running_summary
    system = messages[0]
    older = messages[1:-keep_last]
    recent = messages[-keep_last:]
    if not older:
        return messages, running_summary

    prompt = (
        "Summarize the research progress so far into concise bullet points. Preserve tickers, "
        "prices, key numbers, findings, and any tentative picks.\n\n"
        f"Previous summary:\n{running_summary or '(none)'}\n\nNew activity:\n{_render(older)}"
    )
    try:
        resp = summarizer_llm.invoke(
            [
                SystemMessage(content="You compress research notes faithfully and briefly."),
                HumanMessage(content=prompt),
            ]
        )
        summary = resp.content if isinstance(resp.content, str) else str(resp.content)
    except Exception:
        summary = running_summary  # on failure, keep the old summary and just trim

    new_messages = [system, SystemMessage(content=f"Research summary so far:\n{summary}")] + recent
    return new_messages, summary