from typing import Optional

from vita.data_model.message import SystemMessage, UserMessage
from vita.utils.llm_utils import generate


TEMPLATE = """You are presented with a problem, a section of a dialogue that may contain the answer to the problem, and a previous memory. Please read the provided section carefully and update the memory with the new information that helps to answer the problem. Be sure to retain all relevant details from the previous memory while adding any new, useful information.

<problem> 
用户提出了哪些需求，助理哪些完成了，哪些没完成
</problem>

<memory>
{memory}
</memory>

<section>
{chunk}
</section>

Updated memory:
"""


class MemoryManager:
    def __init__(self, llm: Optional[str] = None, llm_args: Optional[dict] = None):
        self.memory = []
        self.llm = llm
        self.llm_args = llm_args or {}

    def extract_memory(self, chunk: str, previous_memory: str = "") -> str:
        """
        Extract and update memory using the LLM with the provided template.
        
        Args:
            chunk: The section of text to analyze
            previous_memory: Previous memory content (empty string if none)
            
        Returns:
            Updated memory content
        """
        if self.llm is None:
            raise ValueError("LLM is not set. Please provide an LLM when initializing MemoryManager.")
        
        # Format the template with the provided parameters
        formatted_prompt = TEMPLATE.format(
            memory=previous_memory,
            chunk=chunk
        )
        
        # Create messages for the LLM
        messages = [
            UserMessage(role="user", content=formatted_prompt)
        ]
        
        # Generate response using the LLM
        response = generate(
            model=self.llm,
            messages=messages,
            **self.llm_args
        )
        
        # Extract the updated memory from the response
        if response.content:
            return response.content.strip()
        else:
            return previous_memory  # Return previous memory if no new content generated
