from sqlalchemy import Engine

from synthetic_agents.common.constants import (
    AI_CHAT_AAI_INTERVIEWER_AGENT_TYPE,
    AI_CHAT_USER_AGENT_TYPE,
    HUMAN_AGENT_TYPE,
)
from synthetic_agents.database.config import db_engine, get_db

# Suppressed F401 because the tables need to be imported even if they are not explicitly used in
# the code below first in order to Base.metadata.create_all recognize all the tables that need to
# be created.
from synthetic_agents.database.entity.agent import Agent  # noqa: F401
from synthetic_agents.database.entity.agent_type import AgentType  # noqa: F401
from synthetic_agents.database.entity.application import Application  # noqa: F401
from synthetic_agents.database.entity.application_agent import (  # noqa: F401
    ApplicationAgent,
)
from synthetic_agents.database.entity.application_log import (  # noqa: F401
    ApplicationLog,
)
from synthetic_agents.database.entity.application_message import (  # noqa: F401
    ApplicationMessage,
)
from synthetic_agents.database.entity.base import Base  # noqa: F401
from synthetic_agents.database.entity.chat import Chat  # noqa: F401
from synthetic_agents.database.entity.chat_memory import ChatMemory  # noqa: F401
from synthetic_agents.database.entity.life_memory import LifeMemory  # noqa: F401
from synthetic_agents.database.entity.memory import Memory  # noqa: F401
from synthetic_agents.database.entity.prompt_template import (  # noqa: F401
    PromptTemplate,
)
from synthetic_agents.database.entity.working_memory import WorkingMemory  # noqa: F401


def create_all_tables(engine: Engine = db_engine) -> None:
    """
    Creates all tables in the database if they do not yet exist.
    """
    Base.metadata.create_all(engine, checkfirst=True)


def create_aai_agent_types():
    """
    Creates user, AAI interviewer and human agent types in the database.
    """
    db = next(get_db())
    db.add(AgentType(type_id=AI_CHAT_USER_AGENT_TYPE, description="User"))
    db.add(AgentType(type_id=AI_CHAT_AAI_INTERVIEWER_AGENT_TYPE, description="AAI Interviewer"))
    db.add(AgentType(type_id=HUMAN_AGENT_TYPE, description="Human"))
    db.commit()
    db.close()
