"""Context Engine: three-stage lifecycle for context window management.

Stages:
  assemble(messages, budget) → trim history to fit token budget
  maintain(messages)         → post-turn cleanup (compress old tool results, remove stale metadata)
  compact(messages, target)  → adaptive chunked summarization via LLM
"""

from __future__ import annotations

import json
from typing import TYPE_CHECKING, Any

from loguru import logger

from asedclaw.agent.token_utils import estimate_messages_tokens, estimate_tokens

if TYPE_CHECKING:
    from asedclaw.config.schema import ContextBudgetConfig, ToolResultBudgetConfig
    from asedclaw.providers.base import LLMProvider

# Placeholder text when tool output is compacted away
_COMPACTED_MARKER = "[compacted: tool output removed]"


class ContextEngine:
    """Manages context window budget across the agent lifecycle."""

    def __init__(
        self,
        budget_config: ContextBudgetConfig | None = None,
        result_budget_config: ToolResultBudgetConfig | None = None,
    ):
        from asedclaw.config.schema import ContextBudgetConfig, ToolResultBudgetConfig

        self._cfg = budget_config or ContextBudgetConfig()
        self._result_cfg = result_budget_config or ToolResultBudgetConfig()

    # ── Stage 1: Assemble ────────────────────────────────────────────

    def assemble(
        self,
        messages: list[dict[str, Any]],
        budget_tokens: int | None = None,
    ) -> tuple[list[dict[str, Any]], int]:
        """Trim message history to fit within token budget.

        Keeps system prompt + most recent messages that fit.
        Returns (trimmed_messages, estimated_tokens).
        """
        if budget_tokens is None:
            budget_tokens = int(self._cfg.context_window * self._cfg.input_budget_ratio)

        if not messages:
            return messages, 0

        # System prompt always kept
        system_msgs = [m for m in messages if m.get("role") == "system"]
        non_system = [m for m in messages if m.get("role") != "system"]

        system_tokens = estimate_messages_tokens(system_msgs)
        remaining_budget = budget_tokens - system_tokens

        if remaining_budget <= 0:
            logger.warning("System prompt ({} tokens) exceeds budget ({})", system_tokens, budget_tokens)
            return system_msgs + non_system[-2:], estimate_messages_tokens(system_msgs + non_system[-2:])

        # Walk backwards from most recent, accumulating messages
        kept: list[dict[str, Any]] = []
        running_tokens = 0

        for m in reversed(non_system):
            msg_tokens = estimate_messages_tokens([m])
            if running_tokens + msg_tokens > remaining_budget:
                break
            kept.append(m)
            running_tokens += msg_tokens

        kept.reverse()

        # Ensure we always keep at least the last user message
        if not kept and non_system:
            kept = non_system[-1:]
            running_tokens = estimate_messages_tokens(kept)

        result = system_msgs + kept
        total_tokens = system_tokens + running_tokens

        trimmed_count = len(non_system) - len(kept)
        if trimmed_count > 0:
            logger.info(
                "Context assemble: trimmed {} messages, kept {}, ~{} tokens (budget {})",
                trimmed_count, len(kept), total_tokens, budget_tokens,
            )

        return result, total_tokens

    # ── Stage 2: Maintain ────────────────────────────────────────────

    def maintain(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
        """Post-turn maintenance: compress old tool results to save context space.

        Only compresses tool results older than the most recent user message.
        """
        # Find index of the most recent user message
        last_user_idx = -1
        for i in range(len(messages) - 1, -1, -1):
            if messages[i].get("role") == "user":
                last_user_idx = i
                break

        if last_user_idx <= 0:
            return messages

        result = list(messages)
        max_old_tool_chars = int(self._result_cfg.hard_max_tool_result_chars * 0.2)

        for i in range(last_user_idx):
            m = result[i]
            if m.get("role") != "tool":
                continue
            content = m.get("content", "")
            if not isinstance(content, str):
                continue
            if content == _COMPACTED_MARKER:
                continue
            if len(content) > max_old_tool_chars:
                result[i] = {**m, "content": content[:max_old_tool_chars] + "\n... (compacted)"}

        return result

    # ── Stage 3: Preemptive overflow compression ─────────────────────

    def compress_tool_results_if_needed(
        self,
        messages: list[dict[str, Any]],
    ) -> list[dict[str, Any]]:
        """L3: If total tokens exceed preemptive ratio, compress old tool results."""
        total_tokens = estimate_messages_tokens(messages)
        threshold = int(self._cfg.context_window * self._result_cfg.preemptive_ratio)

        if total_tokens <= threshold:
            return messages

        logger.info(
            "Preemptive compression triggered: {} tokens > {} threshold",
            total_tokens, threshold,
        )

        result = list(messages)
        # Compress tool results from oldest to newest until under threshold
        for i, m in enumerate(result):
            if m.get("role") != "tool":
                continue
            content = m.get("content", "")
            if not isinstance(content, str) or content == _COMPACTED_MARKER:
                continue
            result[i] = {**m, "content": _COMPACTED_MARKER}
            total_tokens = estimate_messages_tokens(result)
            if total_tokens <= threshold:
                break

        logger.info("After preemptive compression: ~{} tokens", total_tokens)
        return result

    # ── Tool result budget enforcement (L1 + L2) ─────────────────────

    def enforce_tool_result_budget(
        self,
        result_text: str,
        round_result_tokens: int,
    ) -> tuple[str, int]:
        """Enforce L1 (single result) and L2 (total round) budgets.

        Returns (possibly truncated result, updated round total).
        """
        cfg = self._result_cfg

        # Hard max chars
        if len(result_text) > cfg.hard_max_tool_result_chars:
            result_text = result_text[:cfg.hard_max_tool_result_chars] + "\n... (truncated)"

        # L1: single result token budget
        l1_budget = int(self._cfg.context_window * cfg.single_result_share)
        result_tokens = estimate_tokens(result_text)

        if result_tokens > l1_budget:
            # Estimate chars to keep based on ratio
            keep_ratio = l1_budget / max(result_tokens, 1)
            keep_chars = int(len(result_text) * keep_ratio)
            result_text = result_text[:keep_chars] + "\n... (truncated: exceeded single result budget)"
            result_tokens = estimate_tokens(result_text)
            logger.debug("L1 truncation: {} tokens → {} tokens", result_tokens, l1_budget)

        # L2: total round budget
        l2_budget = int(self._cfg.context_window * cfg.total_result_share)
        if round_result_tokens + result_tokens > l2_budget:
            available = max(100, l2_budget - round_result_tokens)
            keep_ratio = available / max(result_tokens, 1)
            if keep_ratio < 1.0:
                keep_chars = int(len(result_text) * keep_ratio)
                result_text = result_text[:keep_chars] + "\n... (truncated: total result budget exceeded)"
                result_tokens = estimate_tokens(result_text)
                logger.debug("L2 truncation: round total would exceed {} tokens", l2_budget)

        return result_text, round_result_tokens + result_tokens
