"""History compression for conversation context management."""

from __future__ import annotations

from donkit.llm import GenerateRequest, LLMModelAbstract, Message
from loguru import logger


class HistoryCompressor:
    """Compresses conversation history when it exceeds a token threshold.

    Strategies (applied in order):
    1. LLM-based summary of old turns, keeping the last turn intact.
    2. Tool-call compression within the current turn (summarise old tool results).
    3. Emergency truncation of individual oversized messages.

    All thresholds are configurable via constructor arguments.
    """

    def __init__(
        self,
        token_threshold: int = 150_000,
        keep_recent_turns: int = 1,
        keep_recent_tool_pairs: int = 3,
        tool_result_summary_chars: int = 500,
        emergency_msg_max_chars: int = 4_000,
        summary_prompt: str = (
            "Summarize this conversation concisely.\n"
            "Preserve ALL key information: file paths, project names, "
            "configurations, decisions, errors.\n"
            "Format as bullet points. Be brief but complete."
        ),
        fallback_notice: str = (
            "[CONVERSATION HISTORY TRUNCATED]\n"
            "Previous conversation context was too large to summarize. "
            "Key context may have been lost. The most recent interaction is preserved below.\n"
            "[END TRUNCATION NOTICE]"
        ),
    ) -> None:
        self.token_threshold = token_threshold
        self.keep_recent_turns = keep_recent_turns
        self.keep_recent_tool_pairs = keep_recent_tool_pairs
        self.tool_result_summary_chars = tool_result_summary_chars
        self.emergency_msg_max_chars = emergency_msg_max_chars
        self.summary_prompt = summary_prompt
        self.fallback_notice = fallback_notice
        self._tiktoken_encoding = None
        self._tiktoken_loaded = False

    # -- Token estimation ----------------------------------------------------

    def _get_tiktoken_encoding(self):
        if self._tiktoken_loaded:
            return self._tiktoken_encoding
        self._tiktoken_loaded = True
        try:
            import tiktoken

            self._tiktoken_encoding = tiktoken.get_encoding("cl100k_base")
        except ImportError:
            logger.debug("tiktoken not available, using fallback token estimation")
        except Exception as e:
            logger.warning(f"Failed to load tiktoken encoding: {e}")
        return self._tiktoken_encoding

    def _count_tokens(self, text: str) -> int:
        if not text:
            return 0
        encoding = self._get_tiktoken_encoding()
        if encoding:
            try:
                return len(encoding.encode(text))
            except Exception:
                pass
        return len(text) // 4

    def estimate_token_count(self, messages: list[Message]) -> int:
        total = 0
        for msg in messages:
            total += 4  # per-message overhead
            if isinstance(msg.content, str):
                total += self._count_tokens(msg.content)
            elif isinstance(msg.content, list):
                for part in msg.content:
                    if hasattr(part, "content") and isinstance(part.content, str):
                        total += self._count_tokens(part.content)
            if msg.tool_calls:
                for tc in msg.tool_calls:
                    total += self._count_tokens(tc.function.name)
                    total += self._count_tokens(tc.function.arguments)
            if msg.tool_call_id:
                total += self._count_tokens(msg.tool_call_id)
        return total

    # -- Turn splitting ------------------------------------------------------

    @staticmethod
    def _find_recent_complete_turns(
        messages: list[Message], num_turns: int
    ) -> list[Message]:
        if not messages:
            return []
        user_indices = [i for i, m in enumerate(messages) if m.role == "user"]
        if not user_indices:
            return messages
        turns_to_keep = min(num_turns, len(user_indices))
        return messages[user_indices[-turns_to_keep] :]

    # -- Tool-call compression -----------------------------------------------

    def _compress_tool_calls_in_turn(
        self, msgs_to_keep: list[Message]
    ) -> list[Message]:
        if not msgs_to_keep:
            return msgs_to_keep

        leading: list[Message] = []
        tool_sequence: list[Message] = []
        found_tool_call = False

        for msg in msgs_to_keep:
            if (
                not found_tool_call
                and msg.role in ("user", "assistant")
                and not msg.tool_calls
            ):
                leading.append(msg)
            else:
                found_tool_call = True
                tool_sequence.append(msg)

        if not tool_sequence:
            return msgs_to_keep

        pairs: list[list[Message]] = []
        current_pair: list[Message] = []
        for msg in tool_sequence:
            if msg.role == "assistant" and msg.tool_calls and current_pair:
                pairs.append(current_pair)
                current_pair = [msg]
            else:
                current_pair.append(msg)
        if current_pair:
            pairs.append(current_pair)

        if len(pairs) <= self.keep_recent_tool_pairs:
            return msgs_to_keep

        old_pairs = pairs[: -self.keep_recent_tool_pairs]
        recent_pairs = pairs[-self.keep_recent_tool_pairs :]

        summaries: list[str] = []
        for pair in old_pairs:
            for msg in pair:
                if msg.role == "assistant" and msg.tool_calls:
                    for tc in msg.tool_calls:
                        summaries.append(f"- Called {tc.function.name}")
                elif msg.role == "tool":
                    preview = (msg.content or "")[:200]
                    if len(msg.content or "") > 200:
                        preview += "..."
                    tool_name = msg.name or "tool"
                    summaries.append(f"  {tool_name} result: {preview}")

        summary_text = (
            "[COMPRESSED TOOL HISTORY]\n"
            + "\n".join(summaries)
            + "\n[END COMPRESSED TOOL HISTORY]"
        )

        recent_msgs: list[Message] = []
        for pair in recent_pairs:
            recent_msgs.extend(pair)

        return leading + [Message(role="assistant", content=summary_text)] + recent_msgs

    # -- Shrink / truncate helpers -------------------------------------------

    def _shrink_tool_results(self, messages: list[Message]) -> list[Message]:
        shrunk = []
        for msg in messages:
            if (
                msg.role == "tool"
                and isinstance(msg.content, str)
                and len(msg.content) > self.tool_result_summary_chars
            ):
                preview = msg.content[: self.tool_result_summary_chars] + (
                    f"\n... [truncated, was {len(msg.content)} chars]"
                )
                shrunk.append(msg.model_copy(update={"content": preview}))
            else:
                shrunk.append(msg)
        return shrunk

    def _emergency_truncate(self, messages: list[Message]) -> list[Message]:
        truncated = []
        for msg in messages:
            if (
                isinstance(msg.content, str)
                and len(msg.content) > self.emergency_msg_max_chars
            ):
                keep_start = int(self.emergency_msg_max_chars * 0.6)
                keep_end = self.emergency_msg_max_chars - keep_start - 100
                orig_len = len(msg.content)
                notice = (
                    f"\n\n[... truncated {orig_len} -> "
                    f"{self.emergency_msg_max_chars} chars ...]\n\n"
                )
                new_content = (
                    msg.content[:keep_start] + notice + msg.content[-keep_end:]
                )
                truncated.append(msg.model_copy(update={"content": new_content}))
            else:
                truncated.append(msg)
        return truncated

    # -- Main entry point ----------------------------------------------------

    async def compress_if_needed(
        self,
        history: list[Message],
        provider: LLMModelAbstract,
    ) -> list[Message]:
        """Compress history when estimated tokens exceed the threshold.

        Returns the original list if no compression is needed.
        """
        estimated_tokens = self.estimate_token_count(history)
        if estimated_tokens <= self.token_threshold:
            return history

        system_msgs = [m for m in history if m.role == "system"]
        conversation_msgs = [m for m in history if m.role != "system"]

        msgs_to_keep = self._find_recent_complete_turns(
            conversation_msgs, self.keep_recent_turns
        )
        msgs_to_summarize = conversation_msgs[
            : len(conversation_msgs) - len(msgs_to_keep)
        ]

        if not msgs_to_summarize:
            return self._compress_single_turn(
                system_msgs, msgs_to_keep, estimated_tokens
            )

        return await self._compress_multi_turn(
            system_msgs, msgs_to_keep, msgs_to_summarize, provider, estimated_tokens
        )

    # -- Strategy implementations --------------------------------------------

    def _compress_single_turn(
        self,
        system_msgs: list[Message],
        msgs_to_keep: list[Message],
        estimated_tokens: int,
    ) -> list[Message]:
        compressed_turn = self._compress_tool_calls_in_turn(msgs_to_keep)
        new_history = system_msgs + compressed_turn
        new_tokens = self.estimate_token_count(new_history)

        if new_tokens <= self.token_threshold:
            logger.debug(
                "Compressed tool calls: {} -> {} msgs (tokens: {} -> {})",
                len(system_msgs) + len(msgs_to_keep),
                len(new_history),
                estimated_tokens,
                new_tokens,
            )
            return new_history

        logger.debug(
            "Tool-call compression insufficient ({} > {}), applying emergency truncation",
            new_tokens,
            self.token_threshold,
        )
        return system_msgs + self._emergency_truncate(compressed_turn)

    async def _compress_multi_turn(
        self,
        system_msgs: list[Message],
        msgs_to_keep: list[Message],
        msgs_to_summarize: list[Message],
        provider: LLMModelAbstract,
        estimated_tokens: int,
    ) -> list[Message]:
        shrunk = self._shrink_tool_results(msgs_to_summarize)
        try:
            request = GenerateRequest(
                messages=shrunk + [Message(role="user", content=self.summary_prompt)]
            )
            response = await provider.generate(request)
            summary = response.content or ""
            summary_text = f"[CONVERSATION HISTORY SUMMARY]\n{summary}\n[END SUMMARY]"

            compressed_keep = self._compress_tool_calls_in_turn(msgs_to_keep)
            new_history = (
                system_msgs
                + [Message(role="assistant", content=summary_text)]
                + compressed_keep
            )
            logger.debug(
                "Compressed history (LLM summary): tokens {} -> {}",
                estimated_tokens,
                self.estimate_token_count(new_history),
            )
            return new_history
        except Exception as e:
            logger.warning(
                f"LLM-based compression failed: {e}. Using mechanical fallback."
            )

        return self._mechanical_fallback(system_msgs, msgs_to_keep, estimated_tokens)

    def _mechanical_fallback(
        self,
        system_msgs: list[Message],
        msgs_to_keep: list[Message],
        estimated_tokens: int,
    ) -> list[Message]:
        compressed_keep = self._compress_tool_calls_in_turn(msgs_to_keep)
        new_history = (
            system_msgs
            + [Message(role="assistant", content=self.fallback_notice)]
            + compressed_keep
        )

        fallback_tokens = self.estimate_token_count(new_history)
        if fallback_tokens > self.token_threshold:
            logger.warning(
                "Fallback still exceeds threshold ({} > {}). Emergency-truncating.",
                fallback_tokens,
                self.token_threshold,
            )
            new_history = (
                system_msgs
                + [Message(role="assistant", content=self.fallback_notice)]
                + self._emergency_truncate(compressed_keep)
            )

        logger.debug(
            "Mechanical fallback: tokens {} -> {}",
            estimated_tokens,
            self.estimate_token_count(new_history),
        )
        return new_history


# -- Module-level convenience ------------------------------------------------

_default_compressor = HistoryCompressor()

# Backward-compatible aliases for constants
HISTORY_TOKEN_THRESHOLD = _default_compressor.token_threshold
HISTORY_KEEP_RECENT_TURNS = _default_compressor.keep_recent_turns
KEEP_RECENT_TOOL_PAIRS = _default_compressor.keep_recent_tool_pairs
TOOL_RESULT_SUMMARY_CHARS = _default_compressor.tool_result_summary_chars
EMERGENCY_MSG_MAX_CHARS = _default_compressor.emergency_msg_max_chars
HISTORY_SUMMARY_PROMPT = _default_compressor.summary_prompt
FALLBACK_TRUNCATION_NOTICE = _default_compressor.fallback_notice

# Backward-compatible free functions (delegate to default instance)
_estimate_token_count = _default_compressor.estimate_token_count
_find_recent_complete_turns = HistoryCompressor._find_recent_complete_turns
_compress_tool_calls_in_turn = _default_compressor._compress_tool_calls_in_turn
_shrink_tool_results = _default_compressor._shrink_tool_results
_emergency_truncate_messages = _default_compressor._emergency_truncate


async def compress_history_if_needed(
    history: list[Message],
    provider: LLMModelAbstract,
) -> list[Message]:
    """Compress history using the default compressor. Convenience wrapper."""
    return await _default_compressor.compress_if_needed(history, provider)
