import uuid
from typing import List
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.messages import AIMessage, HumanMessage
from sqlalchemy import select, asc, desc
from models.docchat import ChatMessage, ChatSession
from datetime import datetime


class MySQLChatMessageHistory(BaseChatMessageHistory):
    def __init__(self, db, session_id: str, socket_id: str, created_by: str):
        self.db = db  # Instance of MySQLDB
        self.session_id = session_id
        self.socket_id = socket_id
        self.created_by = created_by

    @property
    async def messages(self) -> List:
        filters = [
            ChatMessage.session_id == self.session_id,
            # ChatMessage.socket_id == self.socket_id,
        ]

        order_by = [desc(ChatMessage.created_at)]

        results = await self.db.execute_query(
            model=ChatMessage,
            filters=filters,
            order_by=order_by,
            limit=10,
        )

        messages = []
        for chat_message in results:
            if chat_message.question:
                messages.append(HumanMessage(content=chat_message.question))
            if chat_message.response:
                messages.append(AIMessage(content=chat_message.response))

        return messages

    def add_message(self, message):
        # This will be handled via add_interaction
        pass

    def clear(self):
        # No-op implementation to satisfy the abstract method requirement
        pass

    async def aget_messages(self) -> List:
        return await self.messages

    async def add_interaction(
        self,
        question: str,
        response: str,
        model: str,
        tokens_used: int = 0,
    ):

        new_chat = ChatMessage(
            id=str(uuid.uuid4()),
            socket_id=self.socket_id,
            question=question,
            response=response,
            tokens_used=tokens_used,
            ai_model_used=model,
            created_by=self.created_by,
            modified_by=self.created_by,
            session_id=self.session_id,
        )
        await self.db.execute_insert(new_chat)

    async def update_chat_session_name(self, name: str, user_id: str):
        filters = [
            ChatSession.session_id == self.session_id,
        ]
        await self.db.update_record(
            ChatSession,
            filters=filters,
            values={
                "name": name,
                "modified_by": user_id,
                "modified_at": datetime.now(),
            },
        )

    async def get_chat_session_name(self):
        return await self.db.get_record(
            ChatSession,
            record_id=self.session_id,
            primary_key_column="session_id",
        )