from typing import Dict

from sqlalchemy import JSON, ForeignKey, String
from sqlalchemy.orm import Mapped, mapped_column

from synthetic_agents.database.entity.base import Base


class WorkingMemory(Base):
    """
    This class represents a memory that was retrieved, ranked and loaded into an agent’s working
    memory to produce messages.
    """

    __tablename__ = "working_memory"
    __table_args__ = {
        "comment": "This table stores the different memories that were retrieved, ranked and "
        "loaded into the agents’ working memory to produce messages."
    }

    message_id: Mapped[str] = mapped_column(
        "message_id",
        String,
        ForeignKey("application_message.message_id"),
        primary_key=True,
        comment="ID representing the message in the application.",
    )
    memory_id: Mapped[str] = mapped_column(
        "memory_id",
        String,
        ForeignKey("memory.memory_id"),
        primary_key=True,
        comment="ID representing the memory in the application.",
    )
    attributes: Mapped[Dict[str, str]] = mapped_column(
        "attributes",
        JSON,
        comment="Attributes of the memory in a JSON format. For instance, emotional state, scores "
        "and associated weights.",
    )
