"""Agent 会话上下文压缩中间件。"""

from collections.abc import Awaitable, Callable
from importlib import import_module
from typing import Any

from langchain.agents.middleware.summarization import (
    DEFAULT_SUMMARY_PROMPT,
    ContextSize,
    SummarizationMiddleware,
    TokenCounter,
    TriggerClause,
)
from langchain.agents.middleware.types import (
    AgentMiddleware,
    ExtendedModelResponse,
    ModelRequest,
    ModelResponse,
)
from langchain.chat_models import BaseChatModel
from langchain_core.messages import AnyMessage, HumanMessage, RemoveMessage, ToolMessage
from langchain_core.messages.utils import count_tokens_approximately, get_buffer_string
from langgraph.graph.message import REMOVE_ALL_MESSAGES
from langgraph.types import Command

from app.agent.middleware.usage import UsageMiddleware
from app.runtime.log import logger

try:
    _internal_call_metadata = import_module(
        "langchain.agents.middleware.internal_call_transformer"
    ).internal_call_metadata
except ImportError:

    def _internal_call_metadata() -> dict[str, Any]:
        """旧版 LangChain 没有内部模型调用的流式过滤标记。"""
        return {}


class ContextSummarizationError(RuntimeError):
    """摘要不可用且原有会话上下文未被替换。"""


class ContextPreservingSummarizationMiddleware(SummarizationMiddleware):
    """摘要失败时中止状态更新，避免永久丢失既有会话上下文。"""

    _ERROR_MESSAGE = "会话上下文压缩失败，原有上下文已保留，请稍后重试"
    _UNSUMMARIZABLE_MESSAGE = (
        "会话历史中存在无法压缩的超长内容，原有上下文已保留，"
        "请新建或清空会话后继续"
    )
    # LangChain 默认按 4000 token 裁剪待摘要消息，单条超长工具结果或用户输入
    # （超长无换行文本、非文本多模态块）会被整体丢弃，直接触发"无法压缩"报错。
    # 调大该上限可容纳更大单条消息，减少误报；16k 相对常见模型窗口仍然安全，
    # 摘要模型与主模型同窗口，过大会抬高摘要成本并挤占主模型预算。
    _DEFAULT_TRIM_TOKENS_TO_SUMMARIZE = 16000

    def __init__(
        self,
        model: str | BaseChatModel,
        *,
        trigger: (
            ContextSize | TriggerClause | list[ContextSize | TriggerClause] | None
        ) = None,
        keep: ContextSize = ("messages", 20),
        token_counter: TokenCounter = count_tokens_approximately,
        summary_prompt: str = DEFAULT_SUMMARY_PROMPT,
        trim_tokens_to_summarize: int | None = _DEFAULT_TRIM_TOKENS_TO_SUMMARIZE,
        **deprecated_kwargs: Any,
    ) -> None:
        """按 MoviePilot 的默认压缩策略构建摘要中间件。"""
        super().__init__(
            model=model,
            trigger=trigger,
            keep=keep,
            token_counter=token_counter,
            summary_prompt=summary_prompt,
            trim_tokens_to_summarize=trim_tokens_to_summarize,
            **deprecated_kwargs,
        )

    @classmethod
    def _require_valid_summary(cls, summary: str) -> str:
        """拒绝无法继续承载会话上下文的空摘要。"""
        if not summary:
            raise ContextSummarizationError(cls._ERROR_MESSAGE)
        return summary

    def _prepare_summary_input(
        self, messages_to_summarize: list[AnyMessage]
    ) -> str:
        """复用 LangChain 裁剪策略生成摘要模型输入。"""
        trimmed_messages = self._trim_messages_for_summary(messages_to_summarize)
        if not trimmed_messages:
            raise ContextSummarizationError(self._UNSUMMARIZABLE_MESSAGE)
        return get_buffer_string(trimmed_messages, format="xml")

    def _create_summary(self, messages_to_summarize: list[AnyMessage]) -> str:
        """同步摘要失败时保持原图状态。"""
        formatted_messages = self._prepare_summary_input(messages_to_summarize)
        summary_model = getattr(self, "_summary_model", self.model)
        try:
            response = summary_model.invoke(
                self.summary_prompt.format(messages=formatted_messages).rstrip(),
                config={
                    "metadata": {
                        "lc_source": "summarization",
                        **_internal_call_metadata(),
                    }
                },
            )
        except Exception as err:
            raise ContextSummarizationError(self._ERROR_MESSAGE) from err
        return self._require_valid_summary(response.text.strip())

    async def _acreate_summary(self, messages_to_summarize: list[AnyMessage]) -> str:
        """异步摘要失败时保持原图状态。"""
        formatted_messages = self._prepare_summary_input(messages_to_summarize)
        summary_model = getattr(self, "_summary_model", self.model)
        try:
            response = await summary_model.ainvoke(
                self.summary_prompt.format(messages=formatted_messages).rstrip(),
                config={
                    "metadata": {
                        "lc_source": "summarization",
                        **_internal_call_metadata(),
                    }
                },
            )
        except Exception as err:
            raise ContextSummarizationError(self._ERROR_MESSAGE) from err
        return self._require_valid_summary(response.text.strip())

    def partition_for_token_limit(
        self,
        messages: list[AnyMessage],
        token_limit: int,
        *,
        force: bool = False,
        minimum_cutoff: int = 1,
        strict_token_limit: bool = False,
    ) -> tuple[list[AnyMessage], list[AnyMessage]] | None:
        """按 token 上限拆分历史，并保持 LangChain 的工具调用事务边界。"""
        self._ensure_message_ids(messages)
        if self.token_counter(messages) <= token_limit:
            return (
                self._minimum_safe_partition(
                    messages,
                    minimum_cutoff=minimum_cutoff,
                    token_limit=token_limit if strict_token_limit else None,
                )
                if force
                else None
            )

        left, right = 0, len(messages)
        cutoff_candidate = len(messages)
        while left < right:
            midpoint = (left + right) // 2
            if self._partial_token_counter(messages[midpoint:]) <= token_limit:
                cutoff_candidate = midpoint
                right = midpoint
            else:
                left = midpoint + 1

        if cutoff_candidate >= len(messages):
            cutoff_candidate = len(messages)
        cutoff_index = self._find_safe_cutoff_point(messages, cutoff_candidate)
        if (
            cutoff_index <= 0
            or cutoff_index >= len(messages)
            or cutoff_index < minimum_cutoff
            or not self._contains_unsummarized_message(messages[:cutoff_index])
            or (
                strict_token_limit
                and self._partial_token_counter(messages[cutoff_index:]) > token_limit
            )
        ):
            if cutoff_candidate >= len(messages):
                if strict_token_limit:
                    return None
                return self._latest_safe_partition(
                    messages,
                    minimum_cutoff=minimum_cutoff,
                )
            return self._minimum_safe_partition(
                messages,
                minimum_cutoff=max(minimum_cutoff, cutoff_candidate),
                token_limit=token_limit if strict_token_limit else None,
            )
        return self._partition_messages(messages, cutoff_index)

    def partition_for_retention(
        self, messages: list[AnyMessage]
    ) -> tuple[list[AnyMessage], list[AnyMessage]] | None:
        """按摘要器既有触发和保留策略拆分历史。"""
        self._ensure_message_ids(messages)
        total_tokens = self.token_counter(messages)
        if not self._should_summarize(messages, total_tokens):
            return None
        cutoff_index = self._determine_cutoff_index(messages)
        if cutoff_index <= 0:
            return None
        return self._partition_messages(messages, cutoff_index)

    def _minimum_safe_partition(
        self,
        messages: list[AnyMessage],
        *,
        minimum_cutoff: int = 1,
        token_limit: int | None = None,
    ) -> tuple[list[AnyMessage], list[AnyMessage]] | None:
        """至少摘要一段旧历史，同时保留最新完整消息事务。"""
        for candidate in range(max(1, minimum_cutoff), len(messages)):
            cutoff_index = self._find_safe_cutoff_point(messages, candidate)
            if (
                minimum_cutoff <= cutoff_index < len(messages)
                and self._contains_unsummarized_message(messages[:cutoff_index])
                and (
                    token_limit is None
                    or self._partial_token_counter(messages[cutoff_index:])
                    <= token_limit
                )
            ):
                return self._partition_messages(messages, cutoff_index)
        return None

    def _latest_safe_partition(
        self,
        messages: list[AnyMessage],
        *,
        minimum_cutoff: int,
    ) -> tuple[list[AnyMessage], list[AnyMessage]] | None:
        """保留无法满足软预算时的最新完整消息事务。"""
        for candidate in range(len(messages) - 1, minimum_cutoff - 1, -1):
            cutoff_index = self._find_safe_cutoff_point(messages, candidate)
            if (
                minimum_cutoff <= cutoff_index < len(messages)
                and self._contains_unsummarized_message(messages[:cutoff_index])
            ):
                return self._partition_messages(messages, cutoff_index)
        return None

    @staticmethod
    def _contains_unsummarized_message(messages: list[AnyMessage]) -> bool:
        """确认待摘要段包含可推进上下文的原始消息。"""
        return any(
            message.additional_kwargs.get("lc_source") != "summarization"
            for message in messages
        )

    def build_summary_messages(self, summary: str) -> list[AnyMessage]:
        """将摘要转换为 LangChain 约定的可识别历史消息。"""
        return self._build_new_messages(summary)

    def ensure_message_ids(self, messages: list[AnyMessage]) -> None:
        """为压缩后消息补齐 LangGraph reducer 所需的稳定 ID。"""
        self._ensure_message_ids(messages)

    def create_summary(self, messages_to_summarize: list[AnyMessage]) -> str:
        """通过 MoviePilot 的失败保护合同生成同步摘要。"""
        return self._create_summary(messages_to_summarize)

    async def acreate_summary(self, messages_to_summarize: list[AnyMessage]) -> str:
        """通过 MoviePilot 的失败保护合同生成异步摘要。"""
        return await self._acreate_summary(messages_to_summarize)


class FinalRequestCompactionMiddleware(AgentMiddleware):
    """按最终模型请求预算压缩历史，并在模型成功后原子提交新状态。"""

    _COMPACTION_ANCHOR_KEY = "moviepilot_compaction_anchor_id"
    _UNCOMPRESSIBLE_REQUEST = (
        "最终模型请求压缩后仍超出上下文窗口，原有上下文已保留，"
        "请减少启用工具或切换更大上下文模型"
    )

    def __init__(
        self,
        *,
        summarizer: ContextPreservingSummarizationMiddleware,
        trigger_fraction: float = 0.85,
        keep_fraction: float = 0.10,
    ) -> None:
        self.summarizer = summarizer
        self.trigger_fraction = trigger_fraction
        self.keep_fraction = keep_fraction

    def _should_compact(self, budget: dict[str, Any]) -> bool:
        """以最终请求实际模型窗口判断是否需要压缩。"""
        estimated_tokens = budget.get("estimated_input_tokens")
        context_window = budget.get("context_window_tokens")
        return (
            isinstance(estimated_tokens, int)
            and isinstance(context_window, int)
            and estimated_tokens >= context_window * self.trigger_fraction
        )

    def _compaction_partition(
        self, request: ModelRequest
    ) -> tuple[list[AnyMessage], list[AnyMessage]] | None:
        """最终输入达到阈值时，拆分需要摘要和需要原样保留的消息。"""
        messages = list(request.messages)
        try:
            budget = UsageMiddleware.estimate_request(request)
        except Exception as error:
            logger.debug(
                "最终模型请求预算评估失败，继续原请求: error_type=%s",
                type(error).__name__,
            )
            return None

        context_window = budget.get("context_window_tokens")
        if self._should_skip_after_current_turn_compaction(messages, budget):
            return None
        if not self._should_compact(budget) or not isinstance(context_window, int):
            return None
        try:
            partition = self.summarizer.partition_for_retention(messages)
            if partition is None:
                partition = self.summarizer.partition_for_token_limit(
                    messages,
                    max(1, int(context_window * self.keep_fraction)),
                    force=budget["estimated_input_tokens"] > context_window,
                )
        except Exception as error:
            logger.debug(
                "最终请求历史拆分失败，继续原请求: error_type=%s",
                type(error).__name__,
            )
            if budget["estimated_input_tokens"] > context_window:
                raise ContextSummarizationError(
                    self._UNCOMPRESSIBLE_REQUEST
                ) from error
            return None
        if partition is None:
            if budget["estimated_input_tokens"] > context_window:
                raise ContextSummarizationError(self._UNCOMPRESSIBLE_REQUEST)
            return None
        messages_to_summarize, preserved_messages = partition
        if all(
            message.additional_kwargs.get("lc_source") == "summarization"
            for message in messages_to_summarize
        ):
            if budget["estimated_input_tokens"] > context_window:
                raise ContextSummarizationError(self._UNCOMPRESSIBLE_REQUEST)
            return None
        return messages_to_summarize, preserved_messages

    @classmethod
    def _should_skip_after_current_turn_compaction(
        cls, messages: list[AnyMessage], budget: dict[str, Any]
    ) -> bool:
        """同轮只在新工具结果已使请求超窗时再次压缩。"""
        for message in reversed(messages):
            anchor_id = message.additional_kwargs.get(cls._COMPACTION_ANCHOR_KEY)
            if not isinstance(anchor_id, str):
                continue
            anchor_index = next(
                (
                    index
                    for index, candidate in enumerate(messages)
                    if candidate.id == anchor_id
                ),
                None,
            )
            if anchor_index is None:
                return False
            messages_after_anchor = messages[anchor_index + 1 :]
            if any(
                isinstance(candidate, HumanMessage)
                for candidate in messages_after_anchor
            ):
                return False
            if any(isinstance(candidate, ToolMessage) for candidate in messages_after_anchor):
                estimated_tokens = budget.get("estimated_input_tokens")
                context_window = budget.get("context_window_tokens")
                return not (
                    isinstance(estimated_tokens, int)
                    and isinstance(context_window, int)
                    and estimated_tokens > context_window
                )
            return True
        return False

    def _build_compacted_messages(
        self, summary: str, preserved_messages: list[AnyMessage]
    ) -> list[AnyMessage]:
        """构造摘要与近期历史，并记录本轮压缩输入边界。"""
        summary_messages = self.summarizer.build_summary_messages(summary)
        if summary_messages:
            self.summarizer.ensure_message_ids(summary_messages)
            anchor_id = (
                preserved_messages[-1].id
                if preserved_messages
                else summary_messages[0].id
            )
            first_summary = summary_messages[0]
            summary_messages[0] = first_summary.model_copy(
                update={
                    "additional_kwargs": {
                        **first_summary.additional_kwargs,
                        self._COMPACTION_ANCHOR_KEY: anchor_id,
                    }
                }
            )
        return [*summary_messages, *preserved_messages]

    def _validate_or_repartition(
        self,
        request: ModelRequest,
        messages_to_summarize: list[AnyMessage],
        preserved_messages: list[AnyMessage],
        summary: str,
    ) -> tuple[list[AnyMessage], tuple[list[AnyMessage], list[AnyMessage]] | None]:
        """复核压缩后的最终预算，并计算一次更小的近期历史分区。"""
        compacted_messages = self._build_compacted_messages(summary, preserved_messages)
        compacted_budget = UsageMiddleware.estimate_request(
            request.override(messages=compacted_messages)
        )
        context_window = compacted_budget.get("context_window_tokens")
        estimated_tokens = compacted_budget.get("estimated_input_tokens")
        if not isinstance(context_window, int) or not isinstance(estimated_tokens, int):
            return compacted_messages, None

        if estimated_tokens <= context_window:
            return compacted_messages, None

        summary_messages = self.summarizer.build_summary_messages(summary)
        fixed_summary_budget = UsageMiddleware.estimate_request(
            request.override(messages=summary_messages)
        )
        available_recent_tokens = (
            context_window - fixed_summary_budget["estimated_input_tokens"]
        )
        if available_recent_tokens <= 0:
            raise ContextSummarizationError(self._UNCOMPRESSIBLE_REQUEST)
        repartition = self.summarizer.partition_for_token_limit(
            list(request.messages),
            available_recent_tokens,
            force=True,
            minimum_cutoff=len(messages_to_summarize) + 1,
            strict_token_limit=True,
        )
        if (
            repartition is None
            or len(repartition[0]) <= len(messages_to_summarize)
        ):
            raise ContextSummarizationError(self._UNCOMPRESSIBLE_REQUEST)
        return compacted_messages, repartition

    def _require_within_window(
        self, request: ModelRequest, compacted_messages: list[AnyMessage]
    ) -> None:
        """禁止把已知仍超过主模型窗口的请求发送给 provider。"""
        budget = UsageMiddleware.estimate_request(
            request.override(messages=compacted_messages)
        )
        estimated_tokens = budget.get("estimated_input_tokens")
        context_window = budget.get("context_window_tokens")
        if (
            isinstance(estimated_tokens, int)
            and isinstance(context_window, int)
            and estimated_tokens > context_window
        ):
            raise ContextSummarizationError(self._UNCOMPRESSIBLE_REQUEST)

    def _prepare_messages(self, request: ModelRequest) -> list[AnyMessage] | None:
        """同步生成摘要与需要原样保留的近期消息。"""
        partition = self._compaction_partition(request)
        if partition is None:
            return None
        messages_to_summarize, preserved_messages = partition
        summary = self.summarizer.create_summary(messages_to_summarize)
        compacted_messages, repartition = self._validate_or_repartition(
            request,
            messages_to_summarize,
            preserved_messages,
            summary,
        )
        if repartition is not None:
            messages_to_summarize, preserved_messages = repartition
            summary = self.summarizer.create_summary(messages_to_summarize)
            compacted_messages = self._build_compacted_messages(
                summary, preserved_messages
            )
        self._require_within_window(request, compacted_messages)
        return compacted_messages

    async def _aprepare_messages(
        self, request: ModelRequest
    ) -> list[AnyMessage] | None:
        """异步生成摘要与需要原样保留的近期消息。"""
        partition = self._compaction_partition(request)
        if partition is None:
            return None
        messages_to_summarize, preserved_messages = partition
        summary = await self.summarizer.acreate_summary(messages_to_summarize)
        compacted_messages, repartition = self._validate_or_repartition(
            request,
            messages_to_summarize,
            preserved_messages,
            summary,
        )
        if repartition is not None:
            messages_to_summarize, preserved_messages = repartition
            summary = await self.summarizer.acreate_summary(messages_to_summarize)
            compacted_messages = self._build_compacted_messages(
                summary, preserved_messages
            )
        self._require_within_window(request, compacted_messages)
        return compacted_messages

    @staticmethod
    def _with_state_update(
        response: ModelResponse, compacted_messages: list[AnyMessage]
    ) -> ExtendedModelResponse:
        """主模型成功后一次性替换历史，同时保留本次模型结果。"""
        return ExtendedModelResponse(
            model_response=response,
            command=Command(
                update={
                    "messages": [
                        RemoveMessage(id=REMOVE_ALL_MESSAGES),
                        *compacted_messages,
                        *response.result,
                    ]
                }
            ),
        )

    def wrap_model_call(
        self,
        request: ModelRequest,
        handler: Callable[[ModelRequest], ModelResponse],
    ) -> ModelResponse | ExtendedModelResponse:
        """同步压缩最终请求；模型失败时不提交摘要状态。"""
        compacted_messages = self._prepare_messages(request)
        if compacted_messages is None:
            return handler(request)
        response = handler(request.override(messages=compacted_messages))
        return self._with_state_update(response, compacted_messages)

    async def awrap_model_call(
        self,
        request: ModelRequest,
        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
    ) -> ModelResponse | ExtendedModelResponse:
        """异步压缩最终请求；模型失败时不提交摘要状态。"""
        compacted_messages = await self._aprepare_messages(request)
        if compacted_messages is None:
            return await handler(request)
        response = await handler(request.override(messages=compacted_messages))
        return self._with_state_update(response, compacted_messages)
