# ========= Copyright 2023-2026 @ CAMEL-AI.org. All Rights Reserved. =========
# 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.
# ========= Copyright 2023-2026 @ CAMEL-AI.org. All Rights Reserved. =========
from typing import List, Optional

from camel.memories.base import MemoryBlock
from camel.memories.records import ContextRecord, MemoryRecord
from camel.storages.key_value_storages.base import BaseKeyValueStorage
from camel.storages.key_value_storages.in_memory import InMemoryKeyValueStorage
from camel.types import OpenAIBackendRole


class ChatHistoryBlock(MemoryBlock):
    r"""An implementation of the :obj:`MemoryBlock` abstract base class for
    maintaining a record of chat histories.

    This memory block helps manage conversation histories with a key-value
    storage backend, either provided by the user or using a default
    in-memory storage. It offers a windowed approach to retrieving chat
    histories, allowing users to specify how many recent messages they'd
    like to fetch.

    Args:
        storage (BaseKeyValueStorage, optional): A storage mechanism for
            storing chat history. If `None`, an :obj:`InMemoryKeyValueStorage`
            will be used. (default: :obj:`None`)
        keep_rate (float, optional): In historical messages, the score of the
            last message is 1.0, and with each step taken backward, the score
            of the message is multiplied by the `keep_rate`. Higher `keep_rate`
            leads to high possibility to keep history messages during context
            creation. (default: :obj:`0.9`)
    """

    def __init__(
        self,
        storage: Optional[BaseKeyValueStorage] = None,
        keep_rate: float = 0.9,
    ) -> None:
        if keep_rate > 1 or keep_rate < 0:
            raise ValueError("`keep_rate` should be in [0,1]")
        self.storage = storage or InMemoryKeyValueStorage()
        self.keep_rate = keep_rate

    def retrieve(
        self,
        window_size: Optional[int] = None,
    ) -> List[ContextRecord]:
        r"""Retrieves records with a proper size for the agent from the memory
        based on the window size or fetches the entire chat history if no
        window size is specified.

        Args:
            window_size (int, optional): Specifies the number of recent chat
                messages to retrieve. If not provided, the entire chat history
                will be retrieved. (default: :obj:`None`)

        Returns:
            List[ContextRecord]: A list of retrieved records.
        """
        record_dicts = self.storage.load()
        if len(record_dicts) == 0:
            # Empty memory is a valid state (e.g., during initialization).
            # Users can check if memory is empty by checking the returned list.
            return list()

        if window_size is not None and window_size >= 0:
            # Initial preserved index: Keep first message
            # if it's SYSTEM/DEVELOPER (index 0)
            start_index = (
                1
                if (
                    record_dicts
                    and record_dicts[0]['role_at_backend']
                    in {
                        OpenAIBackendRole.SYSTEM.value,
                        OpenAIBackendRole.DEVELOPER.value,
                    }
                )
                else 0
            )

            """
            Message Processing Logic:
            1. Preserve first system/developer message (if needed)
            2. Keep latest window_size messages from the rest

            Examples:
            - Case 1: First message is SYSTEM, total 5 messages, window_size=2
            Input: [system_msg, user_msg1, user_msg2, user_msg3, user_msg4]
            Result: [system_msg] + [user_msg3, user_msg4]

            - Case 2: First message is USER, total 5 messages, window_size=3
            Input: [user_msg1, user_msg2, user_msg3, user_msg4, , user_msg5]
            Result: [user_msg3, user_msg4, , user_msg5]
            """
            preserved_messages = record_dicts[
                :start_index
            ]  # Preserve system message (if exists)
            sliding_messages = record_dicts[
                start_index:
            ]  # Messages to be truncated

            # Take last window_size messages (if exceeds limit)
            truncated_messages = sliding_messages[-window_size:]

            # Combine preserved messages with truncated window messages
            final_records = preserved_messages + truncated_messages
        else:
            # Return full records when no window restriction
            final_records = record_dicts

        chat_records: List[MemoryRecord] = [
            MemoryRecord.from_dict(record) for record in final_records
        ]

        # We assume that, in the chat history memory, the closer the record is
        # to the current message, the more score it will be.
        output_records = []
        score = 1.0
        for record in reversed(chat_records):
            if record.role_at_backend == OpenAIBackendRole.SYSTEM:
                # System messages are always kept.
                output_records.append(
                    ContextRecord(
                        memory_record=record,
                        score=1.0,
                        timestamp=record.timestamp,
                    )
                )
            else:
                # Other messages' score drops down gradually
                score *= self.keep_rate
                output_records.append(
                    ContextRecord(
                        memory_record=record,
                        score=score,
                        timestamp=record.timestamp,
                    )
                )

        output_records.reverse()
        return output_records

    def write_records(self, records: List[MemoryRecord]) -> None:
        r"""Writes memory records to the memory. Additionally, performs
        validation checks on the messages.

        Args:
            records (List[MemoryRecord]): Memory records to be added to the
                memory.
        """
        stored_records = []
        for record in records:
            stored_records.append(record.to_dict())
        self.storage.save(stored_records)

    def clear(self) -> None:
        r"""Clears all chat messages from the memory."""
        self.storage.clear()

    def pop_records(self, count: int) -> List[MemoryRecord]:
        r"""Removes the most recent records from the memory.

        Args:
            count (int): Number of records to remove from the end of the
                conversation history. A value of 0 results in no changes.

        Returns:
            List[MemoryRecord]: The removed records in chronological order.
        """
        if not isinstance(count, int):
            raise TypeError("`count` must be an integer.")
        if count < 0:
            raise ValueError("`count` must be non-negative.")
        if count == 0:
            return []

        record_dicts = self.storage.load()
        if not record_dicts:
            return []

        # Preserve initial system/developer instruction if present.
        protected_prefix = (
            1
            if (
                record_dicts
                and record_dicts[0]['role_at_backend']
                in {
                    OpenAIBackendRole.SYSTEM.value,
                    OpenAIBackendRole.DEVELOPER.value,
                }
            )
            else 0
        )

        removable_count = max(len(record_dicts) - protected_prefix, 0)
        if removable_count == 0:
            return []

        pop_count = min(count, removable_count)
        split_index = len(record_dicts) - pop_count

        popped_dicts = record_dicts[split_index:]
        remaining_dicts = record_dicts[:split_index]

        self.storage.clear()
        if remaining_dicts:
            self.storage.save(remaining_dicts)

        return [MemoryRecord.from_dict(record) for record in popped_dicts]

    def remove_records_by_indices(
        self, indices: List[int]
    ) -> List[MemoryRecord]:
        r"""Removes records at specified indices from the memory.

        Args:
            indices (List[int]): List of indices to remove. Indices are
                positions in the current record list (0-based).
                System/developer messages at index 0 are protected and will
                not be removed.

        Returns:
            List[MemoryRecord]: The removed records in their original order.
        """
        if not indices:
            return []

        record_dicts = self.storage.load()
        if not record_dicts:
            return []

        # Preserve initial system/developer instruction if present.
        protected_prefix = (
            1
            if (
                record_dicts
                and record_dicts[0]['role_at_backend']
                in {
                    OpenAIBackendRole.SYSTEM.value,
                    OpenAIBackendRole.DEVELOPER.value,
                }
            )
            else 0
        )

        # Filter out protected indices and invalid ones
        valid_indices = sorted(
            {
                idx
                for idx in indices
                if idx >= protected_prefix and idx < len(record_dicts)
            }
        )

        if not valid_indices:
            return []

        # Extract records to remove (in original order)
        removed_records = [record_dicts[idx] for idx in valid_indices]

        # Build remaining records by excluding removed indices
        remaining_dicts = [
            record
            for idx, record in enumerate(record_dicts)
            if idx not in valid_indices
        ]

        # Save back to storage
        self.storage.clear()
        if remaining_dicts:
            self.storage.save(remaining_dicts)

        return [MemoryRecord.from_dict(record) for record in removed_records]
