# %%

import ast
import os
import pickle
import weakref
from datetime import datetime
from typing import Any, Dict

from sqlalchemy.dialects.postgresql import insert
from sqlalchemy.orm import Session, joinedload

from arcan.datamodel.chat_history import ChatHistory
from arcan.datamodel.conversation import Conversation
from arcan.datamodel.engine import SessionLocal


class ArcanSession:
    def __init__(self, database: callable = SessionLocal):
        self.database = database
        self.database_uri = os.environ.get("SQLALCHEMY_URL")
        self.agents: Dict[str, weakref.ref] = weakref.WeakValueDictionary()

    def _get_session(self) -> Session:
        if self.database is None:
            raise ValueError("Database factory is not initialized.")
        return self.database()

    def store_message(self, user_id: str, body: str, response: str):
        with self._get_session() as db_session:
            conversation = Conversation(sender=user_id, message=body, response=response)
            db_session.add(conversation)
            db_session.commit()
            print(f"Conversation #{conversation.id} stored in database")

    def store_chat_history(self, user_id, agent_history):
        history = pickle.dumps(agent_history)
        stmt = (
            insert(ChatHistory)
            .values(
                sender=user_id,
                history=str(history),
                updated_at=datetime.utcnow(),
            )
            .on_conflict_do_update(
                index_elements=["sender"],
                set_={
                    "history": str(history),
                    "updated_at": datetime.utcnow(),
                },
            )
        )
        with self._get_session() as db:
            db.execute(stmt)
            db.commit()
            print(f"Upsert chat history for user {user_id} with statement {stmt}")

    def get_chat_history(self, user_id: str) -> list:
        with self._get_session() as db_session:
            history = (
                db_session.query(ChatHistory)
                # .options(joinedload(ChatHistory.history))
                .filter(ChatHistory.sender == user_id)
                .order_by(ChatHistory.updated_at.asc())
                .all()
            ) or []
        if not history:
            return []
        chat_history = history[0].history
        loaded = pickle.loads(ast.literal_eval(chat_history))
        return loaded

    def rollback(self):
        with self._get_session() as db:
            db.rollback()
            print("Rollback transaction")


# class ArcanSession:
#     def __init__(self, database: Session = None):
#         """
#         Initializes a new instance of the ArcanSession class.

#         :param database: A callable that returns a new SQLAlchemy Session instance when called.
#         """
#         self.database = database
#         self.database_uri = os.environ.get("SQLALCHEMY_URL")
#         self.agents: Dict[str, weakref.ref] = weakref.WeakValueDictionary()

# def store_message(self, user_id: str, body: str, response: str):
#     """
#     Stores a message in the database.

#     :param user_id: The unique identifier for the user.
#     :param Body: The body of the message sent by the user.
#     :param response: The response generated by the system.
#     """
#     with self.database as db_session:
#         conversation = Conversation(sender=user_id, message=body, response=response)
#         db_session.add(conversation)
#         db_session.commit()
#         print(f"Conversation #{conversation.id} stored in database")

#     def store_chat_history(self, user_id, agent_history):
#         """
#         Stores or updates the chat history for a user in the database.

#         :param user_id: The unique identifier for the user.
#         :param agent_history: The chat history to be stored.
#         """
#         history = pickle.dumps(agent_history)
#         # Upsert statement
#         stmt = (
#             insert(ChatHistory)
#             .values(
#                 sender=user_id,
#                 history=str(history),
#                 updated_at=datetime.utcnow(),  # Explicitly set updated_at on insert
#             )
#             .on_conflict_do_update(
#                 index_elements=["sender"],  # Specify the conflict target
#                 set_={
#                     "history": str(history),  # Update the history field upon conflict
#                     "updated_at": datetime.utcnow(),  # Update the updated_at field upon conflict
#                 },
#             )
#         )
#         # Execute the upsert
#         with self.database as db:
#             db.execute(stmt)
#             db.commit()
#             print(f"Upsert chat history for user {user_id} with statement {stmt}")

#     def get_chat_history(self, user_id: str) -> list:
#         """
#         Retrieves the chat history for a user from the database.

#         :param db_session: The SQLAlchemy Session instance.
#         :param user_id: The unique identifier for the user.
#         :return: A list representing the chat history.
#         """
#         with self.database as db_session:
#             history = (
#                 db_session.query(ChatHistory)
#                 .filter(ChatHistory.sender == user_id)
#                 .order_by(ChatHistory.updated_at.asc())
#                 .all()
#             ) or []
#         if not history:
#             return []
#         chat_history = history[0].history
#         loaded = pickle.loads(ast.literal_eval(chat_history))
#         return loaded
# %%
