import logging

from mem0x.configs.base import ShortTermMemoryConfig
from mem0x.memory.base import ShortTermMemoryBase
from mem0x.memory.shortterm_storage import SQLiteManager
from mem0x.memory.shortterm_storage_postgres import PostgreSQLManager

logger = logging.getLogger(__name__)


class ShortTermMemory(ShortTermMemoryBase):
    def __init__(self, config: ShortTermMemoryConfig = ShortTermMemoryConfig()):
        self.config = config

        if self.config.postgres_storage_config:
            self.db = PostgreSQLManager(
                self.config.postgres_storage_config.dbname,
                self.config.postgres_storage_config.user,
                self.config.postgres_storage_config.password,
                self.config.postgres_storage_config.host,
                self.config.postgres_storage_config.port,
            )
        else:
            self.db = SQLiteManager(self.config.sqlite_db_path)

    def add(self, content, user_id, agent_id, run_id, table_name, created_at=None):
        self.db.add(content, user_id, agent_id, run_id, table_name, created_at)

    def get(self, table_name, run_id, timestamp):
        """
        Retrieve a memory by ID.

        Args:
            memory_id (str): ID of the memory to retrieve.

        Returns:
            List: Retrieved memory.
        """
        result = self.db.get(table_name, run_id, timestamp)
        return result

    def get_runids(self, timestamp):
        return self.db.get_runids(timestamp)

    def delete(self, id):
        self.db.delete(id)

    def reset(self):
        self.db.reset()
