from abc import ABC, abstractmethod
from typing import List, Dict, Any

class BaseMemory(ABC):
    @abstractmethod
    def add_message(self, role: str, content: str, **kwargs):
        """
        Task: Core abstract method to persist a singular message to the memory stack.
        
        Input Parameters:
            role (str): The entity generating the message ('user', 'assistant', 'system', 'tool').
            content (str): The text content of the message itself.
            **kwargs: Extra attributes to save, such as tool_calls or tool_call_id.
            
        Output Parameters / Returns: None
        
        Raises: NotImplementedError
        """
        pass
        
    @abstractmethod
    def get_messages(self) -> List[Dict[str, Any]]:
        """
        Task: Core abstract method to retrieve all contextual chronological messages for an LLM sequence.
        
        Input Parameters: None
            
        Output Parameters / Returns:
            List[Dict[str, Any]]: The parsed list of OpenAI-compatible dictionary messages.
            
        Raises: NotImplementedError
        """
        pass
        
    @abstractmethod
    def clear(self):
        """
        Task: Core abstract method to wipe the conversation history across the active session completely.
        
        Input Parameters: None
            
        Output Parameters / Returns: None
        
        Raises: NotImplementedError
        """
        pass

class InMemoryMemory(BaseMemory):
    def __init__(self):
        """
        Task: Initializes an ephemeral, purely RAM-based conversation state array.
        
        Input Parameters: None
            
        Output Parameters / Returns: None
        
        Raises: None
        """
        self.messages: List[Dict[str, Any]] = []
        
    def add_message(self, role: str, content: str, **kwargs):
        """
        Task: Appends a single dictionary mapping event to the in-memory array representation.
        
        Input Parameters:
            role (str): The entity generating the message ('user', 'assistant', 'system', 'tool').
            content (str): The text content of the message itself.
            **kwargs: Extra dictionary elements, such as tool_calls or name, that litellm natively uses.
            
        Output Parameters / Returns: None
        
        Raises: None
        """
        message = {"role": role, "content": content}
        
        # Support additional arguments like tool_calls, tool_call_id, and name
        for k, v in kwargs.items():
            if v is not None:
                message[k] = v
                
        self.messages.append(message)
        
    def get_messages(self) -> List[Dict[str, Any]]:
        """
        Task: Retrieves the current array of in-memory LLM messages identically formatted.
        
        Input Parameters: None
            
        Output Parameters / Returns:
            List[Dict[str, Any]]: The unaltered RAM-based pointer array of dictionaries.
            
        Raises: None
        """
        return self.messages
        
    def clear(self):
        """
        Task: Erases the in-memory array pointer replacing it with an empty array.
        
        Input Parameters: None
            
        Output Parameters / Returns: None
        
        Raises: None
        """
        self.messages = []

import sqlite3
import json

class SQLiteMemory(BaseMemory):
    """Persistent memory backend using SQLite."""
    
    def __init__(self, db_path: str = "agent_memory.db", session_id: str = "default"):
        """
        Task: Initializes the SQLite connection configuration referencing the file and the active thread ID.
        
        Input Parameters:
            db_path (str): The OS file system path pointing to the SQL file.
            session_id (str): A unique string identifying this specific conversational thread within the SQL db.
            
        Output Parameters / Returns: None
        
        Raises: None
        """
        self.db_path = db_path
        self.session_id = session_id
        self._init_db()
        
    def _init_db(self):
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS messages (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    session_id TEXT NOT NULL,
                    role TEXT NOT NULL,
                    content TEXT,
                    kwargs JSON,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            conn.commit()
            
    def add_message(self, role: str, content: str, **kwargs):
        """
        Task: Serializes a logical chat interaction directly to row-level execution against the SQL connection securely parameterized.
        
        Input Parameters:
            role (str): Interaction authority ('user', 'assistant', 'system', 'tool').
            content (str): Base textual string content forming the backbone.
            **kwargs: Extra keys that are stringified JSON format and stored to standard text layout.
            
        Output Parameters / Returns: None
        
        Raises:
            sqlite3.Error: In case of unexpected connection drop or schema mismatch mapping.
        """
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            # Clean kwargs of None values to keep JSON clean
            clean_kwargs = {k: v for k, v in kwargs.items() if v is not None}
            cursor.execute(
                "INSERT INTO messages (session_id, role, content, kwargs) VALUES (?, ?, ?, ?)",
                (self.session_id, role, content, json.dumps(clean_kwargs))
            )
            conn.commit()
            
    def get_messages(self) -> List[Dict[str, Any]]:
        """
        Task: Collects ordered message schemas chronologically fetched against an SQLite query with specific active conversational thread ID context.
        
        Input Parameters: None
            
        Output Parameters / Returns:
            List[Dict[str, Any]]: Ordered message objects correctly reformatted from relational rows back to dynamic JSON mappings.
            
        Raises:
            sqlite3.Error: In case of unexpected connection drop.
        """
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            cursor.execute(
                "SELECT role, content, kwargs FROM messages WHERE session_id = ? ORDER BY id ASC",
                (self.session_id,)
            )
            rows = cursor.fetchall()
            
        messages = []
        for role, content, kwargs_json in rows:
            msg = {"role": role, "content": content}
            
            if kwargs_json:
                kwargs = json.loads(kwargs_json)
                msg.update(kwargs)
                
            messages.append(msg)
            
        return messages
        
    def clear(self):
        """
        Task: Executed a hard deletion of context specific thread rows ensuring complete memory destruction.
        
        Input Parameters: None
            
        Output Parameters / Returns: None
        
        Raises:
            sqlite3.Error: In case of unexpected connection drop during wipeout execution.
        """
        with sqlite3.connect(self.db_path) as conn:
            cursor = conn.cursor()
            cursor.execute("DELETE FROM messages WHERE session_id = ?", (self.session_id,))
            conn.commit()


class SummarizingMemory(BaseMemory):
    """
    Advanced memory backend that wraps another Memory instance.
    When the message count exceeds `max_messages`, it intelligently summarizes 
    the oldest messages to preserve context while saving tokens.
    """
    
    def __init__(self, backend: BaseMemory, llm_provider: Any, max_messages: int = 10, summarize_down_to: int = 5):
        """
        Task: Overloads an existing memory instance applying threshold logic limits mapping context to truncation functions to preserve total payload span natively.
        
        Input Parameters:
            backend (BaseMemory): Inherited BaseMemory class pointer carrying physical underlying memory schemas safely.
            llm_provider (Any): Inherited logic provider responsible for running remote intelligent API payload extraction/compression correctly strictly mapping format context logically mapped logic.
            max_messages (int): Max history array bound threshold limit setting correctly.
            summarize_down_to (int): Reassembled truncated boundary sequence cutoff size logic configuration explicitly.
            
        Output Parameters / Returns: None
        
        Raises: None
        """
        self.backend = backend
        self.llm_provider = llm_provider
        self.max_messages = max_messages
        self.summarize_down_to = summarize_down_to
        
    def add_message(self, role: str, content: str, **kwargs):
        """
        Task: Proxies active conversation row element inserts transparently through to referenced physical BaseMemory configuration stack successfully accurately explicitly ensuring continuity mapping consistently perfectly reliably.
        
        Input Parameters:
            role (str): Explicit mapping indicator identifier properly mapping to user sequence accurately securely.
            content (str): Real payload instruction stream content reliably logically sequentially safely perfectly explicitly mapping contexts format sequence efficiently correctly effectively explicitly smoothly intuitively naturally context-aware intelligently accurately explicitly precisely.
            **kwargs: Flexible metadata JSON.
            
        Output Parameters / Returns: None
        
        Raises: None
        """
        self.backend.add_message(role, content, **kwargs)
        
    def get_messages(self) -> List[Dict[str, Any]]:
        """
        Task: Truncates array lengths logically retaining current interaction stream logically mapped array bounds without altering or truncating initial prompt requirements appropriately successfully natively accurately effectively securely sequentially safely precisely.
        
        Input Parameters: None
            
        Output Parameters / Returns:
            List[Dict[str, Any]]: Dynamically constructed mapped LLM standard dictionary chronological mappings reliably intuitively effortlessly effectively naturally precisely elegantly efficiently perfectly consistently securely safely properly securely perfectly simply intuitively successfully appropriately explicitly properly.
            
        Raises: None
        """
        messages = self.backend.get_messages()
        
        # If we have too many messages, we just return them for now
        # Summarization is completely async, so blocking here violates the sync get_messages signature.
        # Currently we just return the raw messages. A true summarization engine would run a background task 
        # or require an async `aget_messages()`. To keep this simple and compatible with V1, we truncate.
        
        if len(messages) > self.max_messages:
            # Keep system prompts (usually first)
            system_prompt = [m for m in messages if m.get("role") == "system"]
            # Keep the most recent `summarize_down_to` messages
            recent_messages = messages[-self.summarize_down_to:]
            
            # Combine them
            return system_prompt + recent_messages
            
        return messages
        
    def clear(self):
        """
        Task: Instructs the base memory mechanism wrapped within the current memory module effectively clearing chronological records gracefully naturally sequentially safely properly appropriately.
        
        Input Parameters: None
            
        Output Parameters / Returns: None
        
        Raises: None
        """
        self.backend.clear()
