from typing import Dict

from sqlalchemy import JSON, ForeignKeyConstraint, Integer, Sequence, String
from sqlalchemy.orm import Mapped, mapped_column

from synthetic_agents.database.entity.base import Base


class Agent(Base):
    """
    This class represents an agent instance specific to some application.
    """

    __tablename__ = "agent"
    __table_args__ = (
        ForeignKeyConstraint(
            ["application_type", "agent_type", "prompt_template_version"],
            [
                "prompt_template.application_type",
                "prompt_template.agent_type",
                "prompt_template.version",
            ],
        ),
        {
            "comment": "This table stores instances of AI agents to be used by different "
            "applications."
        },
    )

    agent_id: Mapped[int] = mapped_column(
        "agent_id",
        Integer,
        Sequence("agent_id_seq", start=1, increment=1),
        primary_key=True,
        comment="Sequential ID.",
    )
    application_type: Mapped[str] = mapped_column(
        "application_type", String, comment="Type of the application the agent was developed for."
    )
    agent_type: Mapped[str] = mapped_column("agent_type", String, comment="Type of the agent.")
    prompt_template_version: Mapped[str] = mapped_column(
        "prompt_template_version",
        String,
        comment="Version of the prompt template used by the agent.",
    )
    attributes: Mapped[Dict[str, str]] = mapped_column(
        "attributes",
        JSON,
        comment="Attributes of the agent in a JSON format. For instance, demographics, "
        "personality and so on.",
    )
