"""Context compression gates around and inside a turn."""

from typing import TYPE_CHECKING

import structlog

from navi.config import settings

from ..compressor import should_compress
from ..events import (
    CompressionStarted,
)



log = structlog.get_logger()


class AgentCompressionMixin:
    async def _compression_events_preturn(
        self, session, llm, profile, session_id, pending_tokens: int = 0
    ):
        # real_baseline_estimate, not a character heuristic: the baseline is
        # valid here (the context only grew since the last LLM call), and a
        # character count undercounts code-heavy tool output — firing
        # compression too late is exactly what the baseline machinery was
        # built to fix. The midturn gate below does the same.
        estimated_tokens = self._compressor.real_baseline_estimate(
            session.context, session.context, profile.model
        )
        window = self._window_for(profile)
        # `pending_tokens` prices the user message that has not been appended
        # yet — this gate runs before it is. Without it, a message carrying
        # images could overshoot the window in one step: the context was under
        # the threshold a moment ago and the images were the only thing over.
        if (
            settings.context_compression_enabled
            and len(session.context) > 2
            and should_compress(
                estimated_tokens + pending_tokens,
                window,
                settings.context_compression_threshold,
            )
        ):
            # Confirm there is actually something to compress before announcing
            # it. The token gate can fire for a context that the partition cannot
            # shrink (e.g. a few very large messages in a single turn — handled
            # instead by per-message view truncation in context_builder.build).
            # Emitting CompressionStarted without a following ContextCompressed
            # would show a misleading "compression" status with no work done.
            if not self._compressor.would_compress(
                session.context,
                keep_recent=settings.context_keep_recent,
                threshold=settings.context_compression_threshold,
                max_context_tokens=window,
                # Judge the fallback on the same total that fired the gate.
                current_tokens=estimated_tokens + pending_tokens,
                model=profile.model,
            ):
                return
            yield CompressionStarted(
                # The context being compressed does not include the pending
                # message yet, so report the stored context's own size.
                context_tokens=estimated_tokens,
                max_context_tokens=window,
            )
            event = await self._compressor.compress_and_save_session(
                session=session,
                session_store=self._sessions,
                llm=llm,
                model=profile.model,
                temperature=settings.context_summary_temperature,
                session_id=session_id,
                reason="preturn",
                memory_store=self._memory,
                keep_recent=settings.context_keep_recent,
                max_tokens=settings.context_summary_max_tokens,
                window=window,
            )
            if event:
                yield event

    async def _compression_events_midturn(
        self, session, llm, profile, session_id, iteration, ctx_injections, mem
    ):
        if settings.context_compression_enabled and iteration > 0:
            preflight_ctx = self._ctx_builder.build(
                session.context,
                profile,
                mem,
                extra_system=ctx_injections,
                session_id=session_id,
            )
            estimated_tokens = self._compressor.real_baseline_estimate(
                session.context, preflight_ctx, profile.model
            )
            window = self._window_for(profile)
            if should_compress(
                estimated_tokens,
                window,
                settings.context_compression_threshold,
            ):
                # See _compression_events_preturn: only announce compression when
                # the partition (or token-budget fallback) can actually shrink the
                # stored context. The per-message view truncation in build()
                # already caps single huge messages for the LLM, so a no-op here
                # no longer risks an overflow — but announcing it would still lie
                # to the UI.
                if not self._compressor.would_compress(
                    session.context,
                    keep_recent=settings.context_keep_recent,
                    keep_recent_messages=max(12, settings.context_keep_recent * 2),
                    threshold=settings.context_compression_threshold,
                    max_context_tokens=window,
                    current_tokens=estimated_tokens,
                    model=profile.model,
                ):
                    return
                yield CompressionStarted(
                    context_tokens=estimated_tokens,
                    max_context_tokens=window,
                )
                event = await self._compressor.compress_and_save_session(
                    session=session,
                    session_store=self._sessions,
                    llm=llm,
                    model=profile.model,
                    temperature=settings.context_summary_temperature,
                    session_id=session_id,
                    reason="midturn",
                    memory_store=self._memory,
                    keep_recent=settings.context_keep_recent,
                    max_tokens=settings.context_summary_max_tokens,
                    keep_recent_messages=max(12, settings.context_keep_recent * 2),
                    window=window,
                )
                if event:
                    yield event