"""Pre-API preflight compression gate for the conversation turn loop.

Extracted from ``run_conversation``. Runs once per API call after the request pressure
is measured: grow a managed llama.cpp window (last resort), or compress when over
threshold (deferring on noisy estimates, in failure cooldown, or when the review fork's
first request is pending), handle the provider-proven overflow re-check fail-closed,
and emit the deduped blocked/uncompressed overflow warnings. Nothing here imports
``agent.conversation_loop`` at module level (cycle); loop-internal helpers resolve lazily.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from typing import Any, Dict, List, Optional

from agent.context_engine import automatic_compaction_status_message
from agent.conversation_compression import (
    PRE_API_COMPRESSION_STATUS_TEMPLATE,
    compression_blocked_transiently,
    compression_skipped_due_to_lock,
    context_compression_timed_out,
    conversation_history_after_compression,
)
from agent.turn_context import _review_fork_first_request_pending

logger = logging.getLogger("agent.conversation_loop")


@dataclass
class PreflightVerdict:
    """Outcome of ``run_preflight_compression``.

    ``action``: ``"proceed"`` (make the API call), ``"continue"`` (window grown or
    history compacted — the call/budget was refunded, re-enter the turn loop and
    re-measure), ``"break"`` (turn ends: compression timeout or non-actionable
    compaction handoff — ``final_response``/``failed``/``turn_exit_reason`` set) or
    ``"return"`` (typed deferred/exhausted result in ``result``). The remaining fields
    are the loop locals the gate may have rebound."""

    action: str
    result: Optional[Dict[str, Any]]
    messages: List[Dict[str, Any]]
    active_system_prompt: Any
    conversation_history: Any
    api_call_count: int
    compression_attempts: int
    pending_moa_prepared_request: Any
    last_preflight_pressure: Optional[int]
    final_response: Any
    failed: bool
    compression_timeout_exhausted: bool
    turn_exit_reason: Any


def run_preflight_compression(
    agent: Any,
    *,
    compressor: Any,
    request_pressure_tokens: int,
    provider_overflow_preflight: bool,
    preflight_compression_blocked: bool,
    defer_preflight: Any,
    moa_prepared_request: Any,
    pending_moa_prepared_request: Any,
    messages: List[Dict[str, Any]],
    system_message: Any,
    user_message: Any,
    active_system_prompt: Any,
    conversation_history: Any,
    api_call_count: int,
    compression_attempts: int,
    max_compression_attempts: int,
    effective_task_id: Any,
    final_response: Any,
    failed: bool,
    compression_timeout_exhausted: bool,
    turn_exit_reason: Any,
) -> PreflightVerdict:
    """Mirror of the turn-prologue guard chain (defer on noisy estimate → skip in failure
    cooldown → ``should_compress``), #11529. A compression pass that never reaches the
    provider refunds the call/budget in every branch (skip, re-run, timeout) so
    ``api_call_count`` never over-reports; a lock/transient skip refunds the attempt and
    leaves the progress blocker unarmed (#69870, #97488). A forced provider-overflow
    preflight that any gate blocks fails closed (llama.cpp may silently truncate)."""
    from agent.conversation_loop import (
        _COMPRESSION_TIMEOUT_FINAL_RESPONSE,
        _HANDOFF_SKIP_FINAL_RESPONSE,
        _compression_deferred_result,
        _maybe_grow_local_window,
        _provider_overflow_exhausted_result,
        _should_skip_model_call_for_reference_handoff,
    )

    _compressor = compressor
    _provider_overflow_preflight = provider_overflow_preflight
    _preflight_compression_blocked = preflight_compression_blocked
    _defer_preflight = defer_preflight
    _moa_prepared_request = moa_prepared_request
    _last_preflight_pressure = None
    _compression_timeout_exhausted = compression_timeout_exhausted
    _turn_exit_reason = turn_exit_reason

    def _verdict(action: str, result: Optional[Dict[str, Any]] = None) -> PreflightVerdict:
        return PreflightVerdict(
            action=action,
            result=result,
            messages=messages,
            active_system_prompt=active_system_prompt,
            conversation_history=conversation_history,
            api_call_count=api_call_count,
            compression_attempts=compression_attempts,
            pending_moa_prepared_request=pending_moa_prepared_request,
            last_preflight_pressure=_last_preflight_pressure,
            final_response=final_response,
            failed=failed,
            compression_timeout_exhausted=_compression_timeout_exhausted,
            turn_exit_reason=_turn_exit_reason,
        )

    _compression_cooldown = getattr(
        _compressor, "get_active_compression_failure_cooldown", lambda: None
    )()
    if (
        agent.compression_enabled
        and not _review_fork_first_request_pending(agent)
        and len(messages) > 1
        and compression_attempts < max_compression_attempts
        and (
            not _preflight_compression_blocked
            or _provider_overflow_preflight
        )
        and (
            not _defer_preflight(request_pressure_tokens)
            or _provider_overflow_preflight
        )
        and not _compression_cooldown
        and _compressor.should_compress(request_pressure_tokens)
    ):
        # Managed local runtime: grow the context window before compressing (last
        # resort). Only for a llamacpp provider at the supervised base_url.
        _grown_window = _maybe_grow_local_window(
            agent, _compressor, request_pressure_tokens
        )
        if _grown_window:
            # Bigger window granted: recalibrate the compressor and skip compression
            # this pass.
            _compressor.update_model(
                agent.model,
                _grown_window,
                base_url=getattr(agent, "base_url", "") or "",
                api_key=getattr(agent, "api_key", "") or "",
                provider=getattr(agent, "provider", "") or "",
                api_mode=getattr(agent, "api_mode", "") or "",
            )
            agent._buffer_status(
                f"📈 Context window grown to {_grown_window // 1024}K "
                f"(local model; conversation continues uncompressed)"
            )
            # Never reached the provider — refund the call/budget like the
            # compression path does before its continue.
            api_call_count -= 1
            agent._api_call_count = api_call_count
            agent.iteration_budget.refund()
            return _verdict("continue")
        if _moa_prepared_request is not None:
            pending_moa_prepared_request = _moa_prepared_request
        compression_attempts += 1
        # Compression is running: reset the blocked-overflow warning dedup so a
        # later blocked turn warns again (#62625). getattr: test doubles lack it.
        _clear_warn = getattr(agent, "_clear_context_overflow_warn", None)
        if callable(_clear_warn):
            _clear_warn()
        logger.info(
            "Pre-API compression: ~%s request tokens >= %s threshold "
            "(context=%s, attempt=%s/%s)",
            f"{request_pressure_tokens:,}",
            f"{int(getattr(_compressor, 'threshold_tokens', 0) or 0):,}",
            f"{int(getattr(_compressor, 'context_length', 0) or 0):,}"
            if getattr(_compressor, "context_length", 0) else "unknown",
            compression_attempts,
            max_compression_attempts,
        )
        _pre_api_status = automatic_compaction_status_message(
            _compressor,
            phase="pre_api",
            default_message=PRE_API_COMPRESSION_STATUS_TEMPLATE.format(
                tokens=request_pressure_tokens
            ),
            approx_tokens=request_pressure_tokens,
            threshold_tokens=int(
                getattr(_compressor, "threshold_tokens", 0) or 0
            ),
            context_length=int(
                getattr(_compressor, "context_length", 0) or 0
            ),
            model=agent.model,
            attempt=compression_attempts,
            max_attempts=max_compression_attempts,
        )
        if _pre_api_status:
            agent._emit_status(_pre_api_status)
        _last_preflight_pressure = request_pressure_tokens
        _pre_api_input = messages
        messages, active_system_prompt = agent._compress_context(
            messages,
            system_message,
            approx_tokens=request_pressure_tokens,
            task_id=effective_task_id,
        )
        if context_compression_timed_out(agent):
            # Progress-aware timeout (#98722): never reached the provider — refund
            # the call/budget and stop; an overflow retry would only re-compress.
            api_call_count -= 1
            agent._api_call_count = api_call_count
            agent.iteration_budget.refund()
            final_response = _COMPRESSION_TIMEOUT_FINAL_RESPONSE
            failed = True
            _compression_timeout_exhausted = True
            _turn_exit_reason = "context_compression_timeout"
            return _verdict("break")
        if messages is _pre_api_input and (
            compression_skipped_due_to_lock(agent)
            or compression_blocked_transiently(agent)
        ):
            # Temporary DEFER (lock held / cooldown), not evidence about
            # compressibility: refund the attempt, leave the progress blocker
            # unarmed and proceed (#69870, #97488).
            compression_attempts -= 1
            _last_preflight_pressure = None
            if pending_moa_prepared_request is _moa_prepared_request:
                pending_moa_prepared_request = None
        else:
            # Reset retry/empty-response state so the compacted request gets a fresh
            # chance.
            agent._empty_content_retries = 0
            agent._thinking_prefill_retries = 0
            agent._last_content_with_tools = None
            agent._last_content_tools_all_housekeeping = False
            agent._mute_post_response = False
            # Re-baseline the flush cursor: rotation returns None (child flushes
            # whole); in-place returns list(messages) — None would re-append
            # persisted rows. See conversation_history_after_compression().
            conversation_history = conversation_history_after_compression(
                agent, messages, conversation_history
            )
            # Never reaches the provider on skip or re-run — refund the call/budget
            # in BOTH cases, else budget leaks and api_call_count over-reports.
            api_call_count -= 1
            agent._api_call_count = api_call_count
            agent.iteration_budget.refund()
            if _should_skip_model_call_for_reference_handoff(
                messages, user_message
            ):
                # Reference-only handoff must not become the active turn
                # after a completed assistant response (#80622).
                logger.info(
                    "Skipping post-compaction model call: reference-only "
                    "handoff would be the sole active user turn (#80622)"
                )
                if not final_response:
                    final_response = _HANDOFF_SKIP_FINAL_RESPONSE
                _turn_exit_reason = "compaction_handoff_not_actionable"
                return _verdict("break")
            return _verdict("continue")
    elif _provider_overflow_preflight and _compression_cooldown:
        # Provider proved the request cannot fit and the compressor is unavailable:
        # don't resend; let the next user turn retry after cooldown.
        agent._persist_session(messages, conversation_history)
        return _verdict("return", _compression_deferred_result(
            agent,
            messages,
            api_call_count,
            reason="transient_block",
        ))
    elif (
        _provider_overflow_preflight
        and compression_attempts >= max_compression_attempts
    ):
        # All recovery passes consumed and still over threshold: fail closed —
        # llama.cpp may silently truncate an oversized retry.
        return _verdict("return", _provider_overflow_exhausted_result(
            agent,
            messages,
            conversation_history,
            api_call_count,
            request_pressure_tokens,
            max_compression_attempts,
        ))
    elif (
        agent.compression_enabled
        and len(messages) > 1
        and compression_attempts < max_compression_attempts
        and not _defer_preflight(request_pressure_tokens)
        and _compression_cooldown
    ):
        # Summary-LLM cooldown blocks compression: deduped warning only when over
        # threshold (should_compress_info reason is None below it) (#62625).
        _block_reason = None
        try:
            _block_reason = _compressor.should_compress_info(
                request_pressure_tokens
            )[1]
        except Exception:
            _block_reason = None
        if _block_reason:
            agent._warn_context_overflow_blocked(
                _block_reason,
                request_pressure_tokens,
                int(getattr(_compressor, "threshold_tokens", 0) or 0),
            )
    elif not agent.compression_enabled and len(messages) > 1:
        # Uncompressed session guard (#89297): compression is disabled, so warn
        # (deduped) when the request exceeds the context window; the turn-context
        # preflight re-arms the dedup.
        _ctx_len = getattr(
            getattr(agent, "context_compressor", None), "context_length", None
        )
        if (
            isinstance(_ctx_len, int)
            and _ctx_len > 0
            and request_pressure_tokens > _ctx_len
        ):
            _warn_fn = getattr(
                agent, "_warn_uncompressed_context_overflow", None
            )
            if callable(_warn_fn):
                _warn_fn(request_pressure_tokens, _ctx_len)

    if _provider_overflow_preflight:
        # Any other gate blocking the forced preflight (e.g. uncompressible one-
        # message request) must fail closed: the request is proven not to fit.
        return _verdict("return", _provider_overflow_exhausted_result(
            agent,
            messages,
            conversation_history,
            api_call_count,
            request_pressure_tokens,
            max_compression_attempts,
        ))
    return _verdict("proceed")


@dataclass
class PostToolCompressionVerdict:
    """``end_turn`` True → a reference-only compaction handoff would be the sole active
    user turn (#80622): stop without another model call (``final_response`` /
    ``turn_exit_reason`` set)."""

    end_turn: bool
    messages: List[Dict[str, Any]]
    active_system_prompt: Any
    conversation_history: Any
    compression_attempts: int
    final_response: Any
    turn_exit_reason: Any


def compress_after_tool_results(
    agent: Any,
    *,
    messages: List[Dict[str, Any]],
    system_message: Any,
    user_message: Any,
    active_system_prompt: Any,
    conversation_history: Any,
    compression_attempts: int,
    max_compression_attempts: int,
    effective_task_id: Any,
    final_response: Any,
    turn_exit_reason: Any,
) -> PostToolCompressionVerdict:
    """Post-tool-call compression decision. Pressure comes from API-reported
    ``prompt_tokens`` (a tight lower bound; thinking models inflate completion tokens,
    #12026), ``0`` right after compression (no real count yet), else the route-aware
    overhead-inclusive estimate (#14695). Over threshold but blocked → deduped warning
    (#62625) plus the deterministic tool-result-only prune, committed only when the
    engine returns a NEW list (never rebuild ``conversation_history`` for it)."""
    from agent.conversation_loop import (
        _HANDOFF_SKIP_FINAL_RESPONSE,
        _midturn_request_pressure_tokens,
        _should_skip_model_call_for_reference_handoff,
        estimate_request_tokens_rough,
    )

    _turn_exit_reason = turn_exit_reason

    def _verdict(end_turn: bool) -> PostToolCompressionVerdict:
        return PostToolCompressionVerdict(
            end_turn=end_turn,
            messages=messages,
            active_system_prompt=active_system_prompt,
            conversation_history=conversation_history,
            compression_attempts=compression_attempts,
            final_response=final_response,
            turn_exit_reason=_turn_exit_reason,
        )

    # Decide compression from API-reported prompt tokens (tight lower bound;
    # tool results get counted on the next call). If last_prompt_tokens is 0
    # (disconnect / no usage data) fall back to a rough estimate. (#2153)
    _compressor = agent.context_compressor
    if _compressor.last_prompt_tokens > 0:
        # Only prompt_tokens: thinking models inflate completion_tokens with
        # reasoning that uses no context → premature compression. (#12026)
        _real_tokens = _compressor.last_prompt_tokens
    elif _compressor.last_prompt_tokens == -1:
        # Compression just ran, no API prompt count yet: don't treat a rough
        # schema-heavy post-compression estimate as real context pressure.
        _real_tokens = 0
    else:
        # Include tool schemas (20-30K tokens the messages-only estimate
        # misses) and stay route-aware: on a compacted native-Codex session
        # the generic durable-history figure would false-trigger. (#14695)
        _real_tokens = _midturn_request_pressure_tokens(
            agent,
            messages,
            active_system_prompt or "",
            estimate_request_tokens_rough(
                messages, tools=agent.tools or None
            ),
        )

    if (
        agent.compression_enabled
        and compression_attempts < max_compression_attempts
        and _compressor.should_compress(_real_tokens)
    ):
        compression_attempts += 1
        # Compression is running: reset blocked-overflow warning dedup so a
        # future blocked turn can warn again. getattr: test doubles lack it.
        _clear_warn = getattr(agent, "_clear_context_overflow_warn", None)
        if callable(_clear_warn):
            _clear_warn()
        agent._safe_print("  ⟳ compacting context…")
        _post_tool_input = messages
        # Pass overhead-aware _real_tokens, not last_prompt_tokens (0 in
        # the no-usage fallback), so the overflow guard sees the true size.
        messages, active_system_prompt = agent._compress_context(
            messages, system_message,
            approx_tokens=_real_tokens,
            task_id=effective_task_id,
        )
        if (
            messages is _post_tool_input
            and compression_skipped_due_to_lock(agent)
        ):
            # Lock-skip no-op is a temporary defer, not evidence about
            # compressibility: refund so a lock-loser loop doesn't burn the
            # budget toward compression_exhausted. (#69870)
            compression_attempts -= 1
        else:
            conversation_history = conversation_history_after_compression(
                agent, messages, conversation_history
            )
            if _should_skip_model_call_for_reference_handoff(
                messages, user_message
            ):
                logger.info(
                    "Skipping post-tool compaction model call: "
                    "reference-only handoff would be the sole "
                    "active user turn (#80622)"
                )
                if not final_response:
                    final_response = _HANDOFF_SKIP_FINAL_RESPONSE
                _turn_exit_reason = "compaction_handoff_not_actionable"
                return _verdict(True)
    elif agent.compression_enabled:
        # Over threshold but compression blocked (cooldown/anti-thrash):
        # deduped warning so context can't silently overflow. (#62625)
        _block_reason = None
        _info = getattr(_compressor, "should_compress_info", None)
        if _info is not None:
            try:
                _block_reason = _info(_real_tokens)[1]
            except Exception:
                _block_reason = None
        if _block_reason:
            agent._warn_context_overflow_blocked(
                _block_reason,
                _real_tokens,
                int(getattr(_compressor, "threshold_tokens", 0) or 0),
            )
        # Proactive tool-result prune (deterministic, no LLM, keeps tail):
        # no-op unless proactive_prune_tokens is exceeded; commits only past
        # proactive_prune_min_reclaim_tokens so cache breaks stay episodic.
        _prune = getattr(_compressor, "prune_tool_results_only", None)
        if callable(_prune):
            try:
                _pruned_msgs, _pruned_n = _prune(
                    messages, current_tokens=_real_tokens
                )
            except Exception:
                logger.debug(
                    "proactive tool-result prune failed; skipping",
                    exc_info=True,
                )
                _pruned_msgs, _pruned_n = messages, 0
            # Standard no-op caller contract: only commit when the
            # engine returned a NEW list object with a non-zero count.
            if _pruned_n and _pruned_msgs is not messages:
                # Do NOT rebuild conversation_history: rows already carry
                # _DB_PERSISTED_MARKER, and on a stale in-place flag the
                # helper could seed unpersisted rows into history_ids.
                messages = _pruned_msgs
    return _verdict(False)