import asyncio
import uuid
from abc import abstractmethod
from enum import Enum
from sqlalchemy.ext.asyncio import AsyncEngine
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

from llama_index.core.async_utils import asyncio_run
from llama_index.core.base.llms.types import ChatMessage, ContentBlock, TextBlock, AudioBlock, ImageBlock
from llama_index.core.bridge.pydantic import BaseModel, Field, model_validator, ConfigDict
from llama_index.core.memory.types import BaseMemory
from llama_index.core.prompts import RichPromptTemplate
from llama_index.core.storage.chat_store.sql import SQLAlchemyChatStore, MessageStatus
from llama_index.core.utils import get_tokenizer

DEFAULT_TOKEN_LIMIT = 30000
DEFAULT_FLUSH_SIZE = int(DEFAULT_TOKEN_LIMIT * 0.7)
DEFAULT_MEMORY_BLOCKS_TEMPLATE = RichPromptTemplate(
"""
<memory>
{% for (block_name, block_content) in memory_blocks %}
<{{ block_name }}>
  {% for block in block_content %}
    {% if block.block_type == "text" %}
      {{ block.text }}
    {% elif block.block_type == "image" %}
      {% if block.url %}
        {{ block.url | image }}
      {% elif block.path %}
        {{ block.path | image }}
      {% endif %}
    {% elif block.block_type == "audio" %}
      {% if block.url %}
        {{ block.url | audio }}
      {% elif block.path %}
        {{ block.path | audio }}
      {% endif %}
    {% endif %}
  {% endfor %}
</{{ block_name }}>
{% endfor %}
</memory>
"""
)

class InsertMethod(Enum):
    SYSTEM = "system"
    USER = "user"


def generate_chat_store_key() -> str:
    """Generate a unique chat store key."""
    return str(uuid.uuid4())


class BaseMemoryBlock(BaseModel):
    """A base class for memory blocks.
    
    Subclasses must implement the `aget` and `aput` methods.
    Optionally, subclasses can implement the `atruncate` method, which is used to reduce the size of the memory block.
    """

    model_config = ConfigDict(arbitrary_types_allowed=True)

    name: str = Field(description="The name/identifier of the memory block.")
    description: Optional[str] = Field(default=None, description="A description of the memory block.")
    priority: int = Field(default=0, description="Priority of this memory block (0 = never truncate, 1 = highest priority, etc.).")

    @abstractmethod
    async def aget(self, input: Optional[List[ChatMessage]] = None, **kwargs: Any) -> Union[str, List[ContentBlock], List[ChatMessage]]:
        """Pull the memory block (async).
        
        Returns:
            Union[str, List[ContentBlock], List[ChatMessage]]: The memory block content.
            - str: A simple text string to be inserted into the template.
            - List[ContentBlock]: A list of content blocks to be inserted into the template.
            - List[ChatMessage]: A list of chat messages to be directly appended to the chat history.
        """

    @abstractmethod   
    async def aput(self, messages: List[ChatMessage]) -> None:
        """Push to the memory block (async)."""
        
    async def atruncate(self, content: Union[List[ContentBlock], List[ChatMessage]]) -> Union[List[ContentBlock], List[ChatMessage]]:
        """Truncate the memory block content to the given token limit.
        
        By default, truncation will remove the entire block content.
        
        Args:
            content: Either a list of ContentBlock or a list of ChatMessage, depending on what the memory block returns.
            
        Returns:
            The truncated content of the same type as the input.
        """
        return []


class Memory(BaseMemory):
    """A memory module that waterfalls into memory blocks.
    
    Works by orchestrating around
    - a FIFO queue of messages
    - a list of memory blocks
    - various parameters (pressure size, token limit, etc.)

    When the FIFO queue reaches the token limit, the oldest messages within the pressure size are ejected from the FIFO queue.
    The messages are then processed by each memory block.

    When pulling messages from this memory, the memory blocks are processed in order, and the messages are injected into the system message or the latest user message.
    """
    model_config = ConfigDict(arbitrary_types_allowed=True)

    token_limit: int = Field(
        default=DEFAULT_TOKEN_LIMIT,
        description="The overall token limit of the memory.",
    )
    token_flush_size: int = Field(
        default=DEFAULT_FLUSH_SIZE,
        description="The token size to use for flushing the FIFO queue.",
    )
    chat_history_min_token_ratio: float = Field(
        default=0.7,
        description="Minimum percentage ratio of total token limit reserved for chat history.",
    )
    memory_blocks: List[BaseMemoryBlock] = Field(
        default_factory=list,
        description="The list of memory blocks to use.",
    )
    memory_blocks_template: RichPromptTemplate = Field(
        default=DEFAULT_MEMORY_BLOCKS_TEMPLATE,
        description="The template to use for formatting the memory blocks.",
    )
    insert_method: InsertMethod = Field(
        default=InsertMethod.SYSTEM,
        description="Whether to inject memory blocks into a system message or into the latest user message.",
    )
    image_token_size_estimate: int = Field(
        default=256,
        description="The token size estimate for images.",
    )
    audio_token_size_estimate: int = Field(
        default=256,
        description="The token size estimate for audio.",
    )
    tokenizer_fn: Callable[[str], List] = Field(
        default_factory=get_tokenizer,
        exclude=True,
        description="The tokenizer function to use for token counting.",
    )
    sql_store: SQLAlchemyChatStore = Field(
        default_factory=SQLAlchemyChatStore,
        exclude=True,
        description="The chat store to use for storing messages.",
    )
    user_id: str = Field(
        default_factory=generate_chat_store_key,
        description="The key to use for storing messages in the chat store.",
    )

    @classmethod
    def class_name(cls) -> str:
        return "Memory"

    @model_validator(mode="before")
    @classmethod
    def validate_memory(cls, values: dict) -> dict:
        # Validate token limit
        token_limit = values.get("token_limit", -1)
        if token_limit < 1:
            raise ValueError("Token limit must be set and greater than 0.")

        tokenizer_fn = values.get("tokenizer_fn", None)
        if tokenizer_fn is None:
            values["tokenizer_fn"] = get_tokenizer()
        
        if values.get("token_flush_size", -1) < 1:
            values["token_flush_size"] = int(token_limit * 0.7)
        elif values.get("token_flush_size", -1) > token_limit:
            values["token_flush_size"] = int(token_limit * 0.7)

        return values

    @classmethod
    def from_defaults(
        cls,
        user_id: str,
        chat_history: Optional[List[ChatMessage]] = None,
        token_limit: int = DEFAULT_TOKEN_LIMIT,
        memory_blocks: Optional[List[BaseMemoryBlock]] = None,
        tokenizer_fn: Optional[Callable[[str], List]] = None,
        chat_history_min_token_ratio: float = 0.7,
        token_flush_size: int = DEFAULT_FLUSH_SIZE,
        memory_blocks_template: RichPromptTemplate = DEFAULT_MEMORY_BLOCKS_TEMPLATE,
        insert_method: InsertMethod = InsertMethod.SYSTEM,
        image_token_size_estimate: int = 256,
        audio_token_size_estimate: int = 256,
        # SQLAlchemyChatStore parameters
        table_name: str = "llama_index_memory",
        async_database_uri: Optional[str] = None,
        async_engine: Optional[AsyncEngine] = None,
    ) -> "Memory":
        """Initialize Memory."""            
        # If not using the SQLAlchemyChatStore, provide an error
        sql_store = SQLAlchemyChatStore(
            table_name=table_name,
            async_database_uri=async_database_uri,
            async_engine=async_engine,
        )

        if chat_history is not None:
            sql_store.set_messages(user_id, chat_history)
        
        if token_flush_size > token_limit:
            token_flush_size = int(token_limit * 0.7)

        return cls(
            token_limit=token_limit,
            tokenizer_fn=tokenizer_fn or get_tokenizer(),
            sql_store=sql_store,
            user_id=user_id,
            memory_blocks=memory_blocks or [],
            chat_history_min_token_ratio=chat_history_min_token_ratio,
            token_flush_size=token_flush_size,
            memory_blocks_template=memory_blocks_template,
            insert_method=insert_method,
            image_token_size_estimate=image_token_size_estimate,
            audio_token_size_estimate=audio_token_size_estimate,
        )

    def _estimate_token_count(self, message_or_blocks: Union[str, ChatMessage, List[ContentBlock]]) -> int:
        """Estimate token count for a message."""
        token_count = 0

        # Normalize the input to a list of ContentBlocks
        if isinstance(message_or_blocks, ChatMessage):
            blocks = message_or_blocks.blocks

            # Estimate the token count for the additional kwargs
            token_count += len(self.tokenizer_fn(str(message_or_blocks.additional_kwargs)))
        elif isinstance(message_or_blocks, List) and all(isinstance(block, (TextBlock, ImageBlock, AudioBlock)) for block in message_or_blocks):
            blocks = message_or_blocks
        elif isinstance(message_or_blocks, str):
            blocks = [TextBlock(text=message_or_blocks)]
        else:
            raise ValueError(f"Invalid message type: {type(message_or_blocks)}")
        
        # Estimate the token count for each block
        for block in blocks:
            if isinstance(block, TextBlock):
                token_count += len(self.tokenizer_fn(block.text))
            elif isinstance(block, ImageBlock):
                token_count += self.image_token_size_estimate
            elif isinstance(block, AudioBlock):
                token_count += self.audio_token_size_estimate

        return token_count

    async def aget(self, **block_kwargs: Any) -> List[ChatMessage]:
        """Get messages with memory blocks included (async)."""        
        # Get chat history efficiently
        chat_history = await self.sql_store.aget_messages(self.user_id, status=MessageStatus.ACTIVE)
        chat_history_tokens = sum(self._estimate_token_count(message) for message in chat_history)
        
        # Get memory blocks content
        content_per_memory_block: Dict[str, Union[List[ContentBlock], List[ChatMessage]]] = {}
        memory_blocks_tokens = 0
        
        # Process memory blocks in priority order
        for memory_block in sorted(self.memory_blocks, key=lambda x: -x.priority):
            content = await memory_block.aget(chat_history, **block_kwargs)
            
            # Handle different return types from memory blocks
            if content and isinstance(content, list):
                # Memory block returned content blocks
                content_per_memory_block[memory_block.name] = content
                
                # Estimate tokens for each content block
                memory_blocks_tokens += self._estimate_token_count(content)
            elif content and isinstance(content, str):
                # Memory block returned a string
                content_per_memory_block[memory_block.name] = [TextBlock(text=content)]
                memory_blocks_tokens += self._estimate_token_count(content)
            elif not content:
                continue
            else:
                raise ValueError(f"Invalid content type received from memory block {memory_block.name}: {type(content)}")
        
        # Handle truncation if needed
        if memory_blocks_tokens + chat_history_tokens > self.token_limit:
            tokens_to_truncate = memory_blocks_tokens + chat_history_tokens - self.token_limit
            
            # Truncate memory blocks based on priority
            for memory_block in sorted(self.memory_blocks, key=lambda x: x.priority):  # Lower priority first
                if tokens_to_truncate <= 0:
                    break
                
                # Truncate content and messure tokens saved
                content = content_per_memory_block.get(memory_block.name, [])
                    
                truncated_content = await memory_block.atruncate(content)
                
                # Calculate tokens saved
                original_tokens = self._estimate_token_count(content)
                new_tokens = self._estimate_token_count(truncated_content)
                
                tokens_saved = original_tokens - new_tokens
                tokens_to_truncate -= tokens_saved

                # Update the content blocks
                content_per_memory_block[memory_block.name] = truncated_content
        
        # Format template-based memory blocks
        memory_blocks_data: List[Tuple[str, List[ContentBlock]]] = []
        chat_message_data: List[ChatMessage] = []
        for block in self.memory_blocks:
            if block.name in content_per_memory_block:
                content = content_per_memory_block[block.name]
                if isinstance(content, list) and content and isinstance(content[0], ChatMessage):
                    chat_message_data.extend(content)
                else:
                    memory_blocks_data.append((block.name, content))
        
        # Create messages from template content
        memory_content = []
        if memory_blocks_data:
            memory_block_messages = self.memory_blocks_template.format_messages(memory_blocks=memory_blocks_data)
            memory_content = memory_block_messages[0].blocks if memory_block_messages else []
        
        # Process chat messages
        if chat_message_data:
            chat_history = [*chat_message_data, *chat_history]

        # Process template-based memory blocks
        if memory_content:
            if self.insert_method == InsertMethod.SYSTEM:
                # Find system message or create a new one
                system_idx = next((i for i, msg in enumerate(chat_history) if msg.role == "system"), None)
                
                if system_idx is not None:
                    # Update existing system message
                    chat_history[system_idx].blocks = [*memory_content, *chat_history[system_idx].blocks]
                else:
                    # Create new system message at the beginning
                    chat_history.insert(0, ChatMessage(role="system", blocks=memory_content))
            elif self.insert_method == InsertMethod.USER:
                # Find the latest user message
                user_idx = next((i for i, msg in enumerate(reversed(chat_history)) if msg.role == "user"), None)
                
                if user_idx is not None:
                    # Get actual index (since we enumerated in reverse)
                    actual_idx = len(chat_history) - 1 - user_idx
                    # Update existing user message
                    chat_history[actual_idx].blocks = [*chat_history[actual_idx].blocks, *memory_content]
                else:
                    raise ValueError("No user message found in chat history!")
        
        return chat_history
    
    async def _manage_queue(self) -> None:
        """Manage the FIFO queue."""
         # Calculate if we need to waterfall
        current_queue = await self.sql_store.aget_messages(self.user_id, status=MessageStatus.ACTIVE)
        tokens_in_current_queue = sum(self._estimate_token_count(message) for message in current_queue)
        
        # If we're over the token limit, initiate waterfall
        token_limit = self.token_limit * self.chat_history_min_token_ratio
        if tokens_in_current_queue > token_limit:
            # Calculate approximate number of messages to remove
            tokens_to_remove = tokens_in_current_queue - token_limit
            
            while tokens_to_remove > 0:
                # Flush the last X messages that are within the flush size
                # Then check if we are still over the token limit or not
                messages_to_flush = []
                flushed_tokens = 0
                while flushed_tokens < self.token_flush_size:
                    if len(current_queue) == 0:
                        break

                    message = current_queue.pop()
                    messages_to_flush.append(message)
                    flushed_tokens += self._estimate_token_count(message)
                
                # Ensure we keep at least one message
                if len(current_queue) == 0:
                    current_queue.append(messages_to_flush.pop())
                
                # The last message MUST be a user message, default to flushing slightly less than the flush size
                while current_queue[-1].role != "user":
                    current_queue.append(messages_to_flush.pop())

                # Archive the flushed messages
                await self.sql_store.aarchive_oldest_messages(self.user_id, n=len(messages_to_flush))

                # Waterfall the flushed messages
                await asyncio.gather(*[block.aput(messages_to_flush) for block in self.memory_blocks])

                # Update the token count
                tokens_in_current_queue = sum(self._estimate_token_count(message) for message in current_queue)
                tokens_to_remove = tokens_in_current_queue - token_limit

    async def aput(self, message: ChatMessage) -> None:
        """Add a message to the chat store and process waterfall logic if needed."""
        # Add the message to the chat store
        await self.sql_store.async_add_message(self.user_id, message, status=MessageStatus.ACTIVE)
        
        # Ensure the active queue is managed
        await self._manage_queue()
    
    async def aput_messages(self, messages: List[ChatMessage]) -> None:
        """Add a list of messages to the chat store and process waterfall logic if needed."""
        # Add the messages to the chat store
        await self.sql_store.async_add_messages(self.user_id, messages, status=MessageStatus.ACTIVE)

        # Ensure the active queue is managed
        await self._manage_queue()

    async def aset(self, messages: List[ChatMessage]) -> None:
        """Set the chat history."""
        await self.sql_store.aset_messages(self.user_id, messages, status=MessageStatus.ACTIVE)
    
    async def aget_all(self, status: Optional[MessageStatus] = None) -> List[ChatMessage]:
        """Get all messages."""
        return await self.sql_store.aget_messages(self.user_id, status=status)
    
    async def areset(self, status: Optional[MessageStatus] = None) -> None:
        """Reset the memory."""
        await self.sql_store.adelete_messages(self.user_id, status=status)

    # ---- Sync methods ----

    def get(self, **block_kwargs: Any) -> List[ChatMessage]:
        """Get messages with memory blocks included."""
        return asyncio_run(self.aget(**block_kwargs))

    def get_all(self, status: Optional[MessageStatus] = None) -> List[ChatMessage]:
        """Get all messages."""
        return asyncio_run(self.aget_all(status=status))
    
    def put(self, message: ChatMessage) -> None:
        """Add a message to the chat store and process waterfall logic if needed."""
        return asyncio_run(self.aput(message))
    
    def set(self, messages: List[ChatMessage]) -> None:
        """Set the chat history."""
        return asyncio_run(self.aset(messages))
    
    def reset(self) -> None:
        """Reset the memory."""
        return asyncio_run(self.areset())   
