
# Copyright © 2023-2026 Cognizant Technology Solutions Corp, www.cognizant.com.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# END COPYRIGHT

from collections.abc import Callable
from collections.abc import Iterable
from typing import Any
from typing import Dict
from typing import List
from typing import Optional
from typing import Tuple
from typing import Union
from typing_extensions import override

from langchain.agents.middleware.summarization import DEFAULT_SUMMARY_PROMPT
from langchain.agents.middleware.summarization import SummarizationMiddleware
from langchain.agents.middleware.types import AgentState
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.messages.base import BaseMessage
from langchain_core.messages.utils import count_tokens_approximately
from langchain_core.messages.utils import MessageLikeRepresentation

# NOTE:
# The following defaults mirror LangChain's internal summarization defaults
# as of langchain==1.2.19. They are defined locally
# to avoid importing private module-level constants (e.g.,
# `_DEFAULT_MESSAGES_TO_KEEP`, `_DEFAULT_TRIM_TOKEN_LIMIT`), which are not
# part of the public API and may change across LangChain versions.
_DEFAULT_MESSAGES_TO_KEEP: int = 4
_DEFAULT_TRIM_TOKEN_LIMIT: int = 3000


class NeuroSanSummarizationMiddleware(SummarizationMiddleware):
    """
    Adapts LangChain's SummarizationMiddleware for Neuro-san's chat history.

    When chat history exceeds configured ``trigger`` thresholds (message count,
    token count, or both), the ``abefore_model`` hook summarizes older messages
    while preserving recent ones according to the ``keep`` policy.

    Summarization flow (inside ``abefore_model``):
        1. The base class partitions messages into *keep* (recent) and
           *summarize* (older) groups based on the ``keep`` param.
        2. A summary is generated by invoking the configured model.
        3. The base class returns a list of
           ``[RemoveMessage, summary, *kept_messages]``.
        4. This subclass strips the ``RemoveMessage`` (a LangChain-internal
           directive) and replaces Neuro-san's chat history with the remaining
           messages when ``keep_summary_in_context`` is True.

    Important:
        - The summarized message is always recorded in the calling agent's
          journal via ``JournalCallbackHandler.on_llm_end()``.
        - The system message is not part of the chat history — it is always
          injected at the first position separately.

    Note:
        The base ``SummarizationMiddleware`` modifies LangChain-managed state.
        Since Neuro-san manages its own chat history, this subclass adapts
        the output accordingly depending on ``keep_summary_in_context``.

    Example — effect of ``keep_summary_in_context`` on chat history
    (assuming ``keep=("messages", 2)`` and summarization is triggered):

        incoming messages:
            [HumanMessage, AIMessage, HumanMessage, AIMessage, HumanMessage]

        keep_summary_in_context=True:
            [HumanMessage(summary), AIMessage(kept), HumanMessage(kept), AIMessage(new)]

        keep_summary_in_context=False:
            [HumanMessage, AIMessage, HumanMessage, AIMessage, HumanMessage, AIMessage(new)]

    Reference:
        https://reference.langchain.com/python/langchain/agents/middleware/summarization/SummarizationMiddleware
    """

    # pylint: disable=too-many-arguments
    def __init__(
            self,
            *,
            model: Union[str, BaseChatModel],
            chat_history: List[BaseMessage],
            keep_summary_in_context: bool = True,
            trigger: Union[Tuple[str, Union[int, float]], List[Tuple[str, Union[int, float]]]] = None,
            keep: Tuple[str, Union[int, float]] = ("messages", _DEFAULT_MESSAGES_TO_KEEP),
            token_counter: Callable[[Iterable[MessageLikeRepresentation]], int] = count_tokens_approximately,
            summary_prompt: str = DEFAULT_SUMMARY_PROMPT,
            trim_tokens_to_summarize: Optional[int] = _DEFAULT_TRIM_TOKEN_LIMIT,
            **deprecated_kwargs: Any,
    ) -> None:
        """
        Initialize summarization middleware.
        :param model: Language model used to generate summaries. Passed to `init_chat_model`.
                    See https://reference.langchain.com/python/langchain/chat_models/base/init_chat_model
        :param keep_summary_in_context: Whether to keep the summarization in context
        :param chat_history: The chat history of the RunContext that the middleware is being created for.
        :param trigger: One or more thresholds that trigger summarization. (default: None)
        :param keep: Context retention policy applied after summarization.
                    (default: ("messages", _DEFAULT_MESSAGES_TO_KEEP))
        :param token_counter: Function to count tokens in messages. (default: count_tokens_approximately)
        :param summary_prompt: Prompt template for generating summaries. (default: DEFAULT_SUMMARY_PROMPT)
        :param trim_tokens_to_summarize: Maximum tokens to keep when preparing messages for
                the summarization call. (default: _DEFAULT_TRIM_TOKEN_LIMIT)
        """
        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
        )
        self.chat_history: List[BaseMessage] = chat_history
        self.keep_summary_in_context: bool = keep_summary_in_context

    @override
    async def abefore_model(
        self, state: AgentState[Any], runtime: Any
    ) -> Optional[Dict[str, Any]]:
        """
        Run summarization (if triggered) before the model processes messages.

        Delegates to the base class, which checks whether any ``trigger``
        threshold is exceeded. If so, older messages are summarized and the
        result is returned as a dict with a ``"messages"`` key containing
        ``[RemoveMessage, summary, *kept_messages]``.

        When ``keep_summary_in_context`` is True, this method also syncs
        Neuro-san's chat history by replacing it with ``[summary, *kept_messages]``
        (the ``RemoveMessage`` is stripped because it is a LangChain-internal
        directive with no meaning in Neuro-san).

        :param state: The current agent state containing messages.
        :param runtime: The runtime environment (passed through to the base class).
        :return: A dict with the summarized message list, or None if no
                 summarization was triggered.
        """
        # Delegate to the base class to check triggers and summarize if needed.
        messages_dict: Dict[str, Any] = await super().abefore_model(state, runtime)

        # None means no trigger threshold was exceeded — nothing to do.
        if not messages_dict:
            return None

        if self.keep_summary_in_context:
            # Clear the existing chat history so it can be replaced with the
            # condensed version.
            self.chat_history.clear()
            # The base class returns [RemoveMessage, summary, *kept_messages].
            # Skip index 0 (RemoveMessage) — it is a LangChain-internal
            # directive that has no meaning in Neuro-san's history management.
            summarized_messages: List[BaseMessage] = messages_dict.get("messages", [])
            self.chat_history.extend(summarized_messages[1:])

        # Return the full dict (including RemoveMessage) so LangChain's
        # state management can process it downstream.
        return messages_dict
