from langchain_postgres import PostgresChatMessageHistory

from config.pgDb import pg_pool

database_name = "ai_chat_history"


def insert_message(session_id: str, isAI: bool, content: str):
    with pg_pool.getConn() as conn:
        history = PostgresChatMessageHistory(
            database_name, session_id, sync_connection=conn
        )
        if isAI:
            history.add_ai_message(content)
        else:
            history.add_user_message(content)


def get_history(session_id: str):
    simplified = []
    with pg_pool.getConn() as conn:
        history = PostgresChatMessageHistory(
            database_name, session_id, sync_connection=conn
        )
        for msg in history.messages:
            simplified.append({"type": msg.type, "content": msg.content})
        return simplified


def get_history_with_time(session_id: str, number: int):
    simplified = []
    with pg_pool.getConn() as conn:
        with conn.cursor() as cur:
            cur.execute(
                f"SELECT message, created_at FROM ai_chat_history WHERE session_id = '{session_id}' ORDER BY created_at DESC LIMIT {number}"
            )
            rows = cur.fetchall()
            simplified = []

            for row in rows:
                msg_dict = row[0]
                simplified.append(
                    {
                        "type": msg_dict.get("type"),
                        "created_at": row[1].isoformat(),
                        "content": msg_dict.get("data", {}).get("content"),
                    }
                )

        return simplified